You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669
  1. #
  2. # Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. #
  16. import re
  17. import random
  18. from openai.lib.azure import AzureOpenAI
  19. from zhipuai import ZhipuAI
  20. from dashscope import Generation
  21. from abc import ABC
  22. from openai import OpenAI
  23. import openai
  24. from ollama import Client
  25. from rag.nlp import is_chinese, is_english
  26. from rag.utils import num_tokens_from_string
  27. import os
  28. import json
  29. import requests
  30. import asyncio
  31. import logging
  32. import time
  33. # Error message constants
  34. ERROR_PREFIX = "**ERROR**"
  35. ERROR_RATE_LIMIT = "RATE_LIMIT_EXCEEDED"
  36. ERROR_AUTHENTICATION = "AUTH_ERROR"
  37. ERROR_INVALID_REQUEST = "INVALID_REQUEST"
  38. ERROR_SERVER = "SERVER_ERROR"
  39. ERROR_TIMEOUT = "TIMEOUT"
  40. ERROR_CONNECTION = "CONNECTION_ERROR"
  41. ERROR_MODEL = "MODEL_ERROR"
  42. ERROR_CONTENT_FILTER = "CONTENT_FILTERED"
  43. ERROR_QUOTA = "QUOTA_EXCEEDED"
  44. ERROR_MAX_RETRIES = "MAX_RETRIES_EXCEEDED"
  45. ERROR_GENERIC = "GENERIC_ERROR"
  46. LENGTH_NOTIFICATION_CN = "······\n由于大模型的上下文窗口大小限制,回答已经被大模型截断。"
  47. LENGTH_NOTIFICATION_EN = "...\nThe answer is truncated by your chosen LLM due to its limitation on context length."
  48. class Base(ABC):
  49. def __init__(self, key, model_name, base_url):
  50. timeout = int(os.environ.get('LM_TIMEOUT_SECONDS', 600))
  51. self.client = OpenAI(api_key=key, base_url=base_url, timeout=timeout)
  52. self.model_name = model_name
  53. # Configure retry parameters
  54. self.max_retries = int(os.environ.get('LLM_MAX_RETRIES', 5))
  55. self.base_delay = float(os.environ.get('LLM_BASE_DELAY', 2.0))
  56. def _get_delay(self, attempt):
  57. """Calculate retry delay time"""
  58. return self.base_delay * (2 ** attempt) + random.uniform(0, 0.5)
  59. def _classify_error(self, error):
  60. """Classify error based on error message content"""
  61. error_str = str(error).lower()
  62. if "rate limit" in error_str or "429" in error_str or "tpm limit" in error_str or "too many requests" in error_str or "requests per minute" in error_str:
  63. return ERROR_RATE_LIMIT
  64. elif "auth" in error_str or "key" in error_str or "apikey" in error_str or "401" in error_str or "forbidden" in error_str or "permission" in error_str:
  65. return ERROR_AUTHENTICATION
  66. elif "invalid" in error_str or "bad request" in error_str or "400" in error_str or "format" in error_str or "malformed" in error_str or "parameter" in error_str:
  67. return ERROR_INVALID_REQUEST
  68. elif "server" in error_str or "502" in error_str or "503" in error_str or "504" in error_str or "500" in error_str or "unavailable" in error_str:
  69. return ERROR_SERVER
  70. elif "timeout" in error_str or "timed out" in error_str:
  71. return ERROR_TIMEOUT
  72. elif "connect" in error_str or "network" in error_str or "unreachable" in error_str or "dns" in error_str:
  73. return ERROR_CONNECTION
  74. elif "quota" in error_str or "capacity" in error_str or "credit" in error_str or "billing" in error_str or "limit" in error_str and "rate" not in error_str:
  75. return ERROR_QUOTA
  76. elif "filter" in error_str or "content" in error_str or "policy" in error_str or "blocked" in error_str or "safety" in error_str:
  77. return ERROR_CONTENT_FILTER
  78. elif "model" in error_str or "not found" in error_str or "does not exist" in error_str or "not available" in error_str:
  79. return ERROR_MODEL
  80. else:
  81. return ERROR_GENERIC
  82. def chat(self, system, history, gen_conf):
  83. if system:
  84. history.insert(0, {"role": "system", "content": system})
  85. if "max_tokens" in gen_conf:
  86. del gen_conf["max_tokens"]
  87. # Implement exponential backoff retry strategy
  88. for attempt in range(self.max_retries):
  89. try:
  90. response = self.client.chat.completions.create(
  91. model=self.model_name,
  92. messages=history,
  93. **gen_conf)
  94. if any([not response.choices, not response.choices[0].message, not response.choices[0].message.content]):
  95. return "", 0
  96. ans = response.choices[0].message.content.strip()
  97. if response.choices[0].finish_reason == "length":
  98. if is_chinese(ans):
  99. ans += LENGTH_NOTIFICATION_CN
  100. else:
  101. ans += LENGTH_NOTIFICATION_EN
  102. return ans, self.total_token_count(response)
  103. except Exception as e:
  104. # Classify the error
  105. error_code = self._classify_error(e)
  106. # Check if it's a rate limit error or server error and not the last attempt
  107. should_retry = (error_code == ERROR_RATE_LIMIT or error_code == ERROR_SERVER) and attempt < self.max_retries - 1
  108. if should_retry:
  109. delay = self._get_delay(attempt)
  110. logging.warning(f"Error: {error_code}. Retrying in {delay:.2f} seconds... (Attempt {attempt+1}/{self.max_retries})")
  111. time.sleep(delay)
  112. else:
  113. # For non-rate limit errors or the last attempt, return an error message
  114. if attempt == self.max_retries - 1:
  115. error_code = ERROR_MAX_RETRIES
  116. return f"{ERROR_PREFIX}: {error_code} - {str(e)}", 0
  117. def chat_streamly(self, system, history, gen_conf):
  118. if system:
  119. history.insert(0, {"role": "system", "content": system})
  120. if "max_tokens" in gen_conf:
  121. del gen_conf["max_tokens"]
  122. ans = ""
  123. total_tokens = 0
  124. try:
  125. response = self.client.chat.completions.create(
  126. model=self.model_name,
  127. messages=history,
  128. stream=True,
  129. **gen_conf)
  130. for resp in response:
  131. if not resp.choices:
  132. continue
  133. if not resp.choices[0].delta.content:
  134. resp.choices[0].delta.content = ""
  135. if hasattr(resp.choices[0].delta, "reasoning_content") and resp.choices[0].delta.reasoning_content:
  136. if ans.find("<think>") < 0:
  137. ans += "<think>"
  138. ans = ans.replace("</think>", "")
  139. ans += resp.choices[0].delta.reasoning_content + "</think>"
  140. else:
  141. ans += resp.choices[0].delta.content
  142. tol = self.total_token_count(resp)
  143. if not tol:
  144. total_tokens += num_tokens_from_string(resp.choices[0].delta.content)
  145. else:
  146. total_tokens = tol
  147. if resp.choices[0].finish_reason == "length":
  148. if is_chinese(ans):
  149. ans += LENGTH_NOTIFICATION_CN
  150. else:
  151. ans += LENGTH_NOTIFICATION_EN
  152. yield ans
  153. except openai.APIError as e:
  154. yield ans + "\n**ERROR**: " + str(e)
  155. yield total_tokens
  156. def total_token_count(self, resp):
  157. try:
  158. return resp.usage.total_tokens
  159. except Exception:
  160. pass
  161. try:
  162. return resp["usage"]["total_tokens"]
  163. except Exception:
  164. pass
  165. return 0
  166. class GptTurbo(Base):
  167. def __init__(self, key, model_name="gpt-3.5-turbo", base_url="https://api.openai.com/v1"):
  168. if not base_url:
  169. base_url = "https://api.openai.com/v1"
  170. super().__init__(key, model_name, base_url)
  171. class MoonshotChat(Base):
  172. def __init__(self, key, model_name="moonshot-v1-8k", base_url="https://api.moonshot.cn/v1"):
  173. if not base_url:
  174. base_url = "https://api.moonshot.cn/v1"
  175. super().__init__(key, model_name, base_url)
  176. class XinferenceChat(Base):
  177. def __init__(self, key=None, model_name="", base_url=""):
  178. if not base_url:
  179. raise ValueError("Local llm url cannot be None")
  180. if base_url.split("/")[-1] != "v1":
  181. base_url = os.path.join(base_url, "v1")
  182. super().__init__(key, model_name, base_url)
  183. class HuggingFaceChat(Base):
  184. def __init__(self, key=None, model_name="", base_url=""):
  185. if not base_url:
  186. raise ValueError("Local llm url cannot be None")
  187. if base_url.split("/")[-1] != "v1":
  188. base_url = os.path.join(base_url, "v1")
  189. super().__init__(key, model_name.split("___")[0], base_url)
  190. class ModelScopeChat(Base):
  191. def __init__(self, key=None, model_name="", base_url=""):
  192. if not base_url:
  193. raise ValueError("Local llm url cannot be None")
  194. base_url = base_url.rstrip('/')
  195. if base_url.split("/")[-1] != "v1":
  196. base_url = os.path.join(base_url, "v1")
  197. super().__init__(key, model_name.split("___")[0], base_url)
  198. class DeepSeekChat(Base):
  199. def __init__(self, key, model_name="deepseek-chat", base_url="https://api.deepseek.com/v1"):
  200. if not base_url:
  201. base_url = "https://api.deepseek.com/v1"
  202. super().__init__(key, model_name, base_url)
  203. class AzureChat(Base):
  204. def __init__(self, key, model_name, **kwargs):
  205. api_key = json.loads(key).get('api_key', '')
  206. api_version = json.loads(key).get('api_version', '2024-02-01')
  207. self.client = AzureOpenAI(api_key=api_key, azure_endpoint=kwargs["base_url"], api_version=api_version)
  208. self.model_name = model_name
  209. class BaiChuanChat(Base):
  210. def __init__(self, key, model_name="Baichuan3-Turbo", base_url="https://api.baichuan-ai.com/v1"):
  211. if not base_url:
  212. base_url = "https://api.baichuan-ai.com/v1"
  213. super().__init__(key, model_name, base_url)
  214. @staticmethod
  215. def _format_params(params):
  216. return {
  217. "temperature": params.get("temperature", 0.3),
  218. "top_p": params.get("top_p", 0.85),
  219. }
  220. def chat(self, system, history, gen_conf):
  221. if system:
  222. history.insert(0, {"role": "system", "content": system})
  223. if "max_tokens" in gen_conf:
  224. del gen_conf["max_tokens"]
  225. try:
  226. response = self.client.chat.completions.create(
  227. model=self.model_name,
  228. messages=history,
  229. extra_body={
  230. "tools": [{
  231. "type": "web_search",
  232. "web_search": {
  233. "enable": True,
  234. "search_mode": "performance_first"
  235. }
  236. }]
  237. },
  238. **self._format_params(gen_conf))
  239. ans = response.choices[0].message.content.strip()
  240. if response.choices[0].finish_reason == "length":
  241. if is_chinese([ans]):
  242. ans += LENGTH_NOTIFICATION_CN
  243. else:
  244. ans += LENGTH_NOTIFICATION_EN
  245. return ans, self.total_token_count(response)
  246. except openai.APIError as e:
  247. return "**ERROR**: " + str(e), 0
  248. def chat_streamly(self, system, history, gen_conf):
  249. if system:
  250. history.insert(0, {"role": "system", "content": system})
  251. if "max_tokens" in gen_conf:
  252. del gen_conf["max_tokens"]
  253. ans = ""
  254. total_tokens = 0
  255. try:
  256. response = self.client.chat.completions.create(
  257. model=self.model_name,
  258. messages=history,
  259. extra_body={
  260. "tools": [{
  261. "type": "web_search",
  262. "web_search": {
  263. "enable": True,
  264. "search_mode": "performance_first"
  265. }
  266. }]
  267. },
  268. stream=True,
  269. **self._format_params(gen_conf))
  270. for resp in response:
  271. if not resp.choices:
  272. continue
  273. if not resp.choices[0].delta.content:
  274. resp.choices[0].delta.content = ""
  275. ans += resp.choices[0].delta.content
  276. tol = self.total_token_count(resp)
  277. if not tol:
  278. total_tokens += num_tokens_from_string(resp.choices[0].delta.content)
  279. else:
  280. total_tokens = tol
  281. if resp.choices[0].finish_reason == "length":
  282. if is_chinese([ans]):
  283. ans += LENGTH_NOTIFICATION_CN
  284. else:
  285. ans += LENGTH_NOTIFICATION_EN
  286. yield ans
  287. except Exception as e:
  288. yield ans + "\n**ERROR**: " + str(e)
  289. yield total_tokens
  290. class QWenChat(Base):
  291. def __init__(self, key, model_name=Generation.Models.qwen_turbo, **kwargs):
  292. import dashscope
  293. dashscope.api_key = key
  294. self.model_name = model_name
  295. if self.is_reasoning_model(self.model_name):
  296. super().__init__(key, model_name, "https://dashscope.aliyuncs.com/compatible-mode/v1")
  297. def chat(self, system, history, gen_conf):
  298. if "max_tokens" in gen_conf:
  299. del gen_conf["max_tokens"]
  300. if self.is_reasoning_model(self.model_name):
  301. return super().chat(system, history, gen_conf)
  302. stream_flag = str(os.environ.get('QWEN_CHAT_BY_STREAM', 'true')).lower() == 'true'
  303. if not stream_flag:
  304. from http import HTTPStatus
  305. if system:
  306. history.insert(0, {"role": "system", "content": system})
  307. response = Generation.call(
  308. self.model_name,
  309. messages=history,
  310. result_format='message',
  311. **gen_conf
  312. )
  313. ans = ""
  314. tk_count = 0
  315. if response.status_code == HTTPStatus.OK:
  316. ans += response.output.choices[0]['message']['content']
  317. tk_count += self.total_token_count(response)
  318. if response.output.choices[0].get("finish_reason", "") == "length":
  319. if is_chinese([ans]):
  320. ans += LENGTH_NOTIFICATION_CN
  321. else:
  322. ans += LENGTH_NOTIFICATION_EN
  323. return ans, tk_count
  324. return "**ERROR**: " + response.message, tk_count
  325. else:
  326. g = self._chat_streamly(system, history, gen_conf, incremental_output=True)
  327. result_list = list(g)
  328. error_msg_list = [item for item in result_list if str(item).find("**ERROR**") >= 0]
  329. if len(error_msg_list) > 0:
  330. return "**ERROR**: " + "".join(error_msg_list), 0
  331. else:
  332. return "".join(result_list[:-1]), result_list[-1]
  333. def _chat_streamly(self, system, history, gen_conf, incremental_output=False):
  334. from http import HTTPStatus
  335. if system:
  336. history.insert(0, {"role": "system", "content": system})
  337. if "max_tokens" in gen_conf:
  338. del gen_conf["max_tokens"]
  339. ans = ""
  340. tk_count = 0
  341. try:
  342. response = Generation.call(
  343. self.model_name,
  344. messages=history,
  345. result_format='message',
  346. stream=True,
  347. incremental_output=incremental_output,
  348. **gen_conf
  349. )
  350. for resp in response:
  351. if resp.status_code == HTTPStatus.OK:
  352. ans = resp.output.choices[0]['message']['content']
  353. tk_count = self.total_token_count(resp)
  354. if resp.output.choices[0].get("finish_reason", "") == "length":
  355. if is_chinese(ans):
  356. ans += LENGTH_NOTIFICATION_CN
  357. else:
  358. ans += LENGTH_NOTIFICATION_EN
  359. yield ans
  360. else:
  361. yield ans + "\n**ERROR**: " + resp.message if not re.search(r" (key|quota)",
  362. str(resp.message).lower()) else "Out of credit. Please set the API key in **settings > Model providers.**"
  363. except Exception as e:
  364. yield ans + "\n**ERROR**: " + str(e)
  365. yield tk_count
  366. def chat_streamly(self, system, history, gen_conf):
  367. if "max_tokens" in gen_conf:
  368. del gen_conf["max_tokens"]
  369. if self.is_reasoning_model(self.model_name):
  370. return super().chat_streamly(system, history, gen_conf)
  371. return self._chat_streamly(system, history, gen_conf)
  372. @staticmethod
  373. def is_reasoning_model(model_name: str) -> bool:
  374. return any([
  375. model_name.lower().find("deepseek") >= 0,
  376. model_name.lower().find("qwq") >= 0 and model_name.lower() != 'qwq-32b-preview',
  377. ])
  378. class ZhipuChat(Base):
  379. def __init__(self, key, model_name="glm-3-turbo", **kwargs):
  380. self.client = ZhipuAI(api_key=key)
  381. self.model_name = model_name
  382. def chat(self, system, history, gen_conf):
  383. if system:
  384. history.insert(0, {"role": "system", "content": system})
  385. if "max_tokens" in gen_conf:
  386. del gen_conf["max_tokens"]
  387. try:
  388. if "presence_penalty" in gen_conf:
  389. del gen_conf["presence_penalty"]
  390. if "frequency_penalty" in gen_conf:
  391. del gen_conf["frequency_penalty"]
  392. response = self.client.chat.completions.create(
  393. model=self.model_name,
  394. messages=history,
  395. **gen_conf
  396. )
  397. ans = response.choices[0].message.content.strip()
  398. if response.choices[0].finish_reason == "length":
  399. if is_chinese(ans):
  400. ans += LENGTH_NOTIFICATION_CN
  401. else:
  402. ans += LENGTH_NOTIFICATION_EN
  403. return ans, self.total_token_count(response)
  404. except Exception as e:
  405. return "**ERROR**: " + str(e), 0
  406. def chat_streamly(self, system, history, gen_conf):
  407. if system:
  408. history.insert(0, {"role": "system", "content": system})
  409. if "max_tokens" in gen_conf:
  410. del gen_conf["max_tokens"]
  411. if "presence_penalty" in gen_conf:
  412. del gen_conf["presence_penalty"]
  413. if "frequency_penalty" in gen_conf:
  414. del gen_conf["frequency_penalty"]
  415. ans = ""
  416. tk_count = 0
  417. try:
  418. response = self.client.chat.completions.create(
  419. model=self.model_name,
  420. messages=history,
  421. stream=True,
  422. **gen_conf
  423. )
  424. for resp in response:
  425. if not resp.choices[0].delta.content:
  426. continue
  427. delta = resp.choices[0].delta.content
  428. ans += delta
  429. if resp.choices[0].finish_reason == "length":
  430. if is_chinese(ans):
  431. ans += LENGTH_NOTIFICATION_CN
  432. else:
  433. ans += LENGTH_NOTIFICATION_EN
  434. tk_count = self.total_token_count(resp)
  435. if resp.choices[0].finish_reason == "stop":
  436. tk_count = self.total_token_count(resp)
  437. yield ans
  438. except Exception as e:
  439. yield ans + "\n**ERROR**: " + str(e)
  440. yield tk_count
  441. class OllamaChat(Base):
  442. def __init__(self, key, model_name, **kwargs):
  443. self.client = Client(host=kwargs["base_url"]) if not key or key == "x" else \
  444. Client(host=kwargs["base_url"], headers={"Authorization": f"Bear {key}"})
  445. self.model_name = model_name
  446. def chat(self, system, history, gen_conf):
  447. if system:
  448. history.insert(0, {"role": "system", "content": system})
  449. if "max_tokens" in gen_conf:
  450. del gen_conf["max_tokens"]
  451. try:
  452. options = {
  453. "num_ctx": 32768
  454. }
  455. if "temperature" in gen_conf:
  456. options["temperature"] = gen_conf["temperature"]
  457. if "max_tokens" in gen_conf:
  458. options["num_predict"] = gen_conf["max_tokens"]
  459. if "top_p" in gen_conf:
  460. options["top_p"] = gen_conf["top_p"]
  461. if "presence_penalty" in gen_conf:
  462. options["presence_penalty"] = gen_conf["presence_penalty"]
  463. if "frequency_penalty" in gen_conf:
  464. options["frequency_penalty"] = gen_conf["frequency_penalty"]
  465. response = self.client.chat(
  466. model=self.model_name,
  467. messages=history,
  468. options=options,
  469. keep_alive=-1
  470. )
  471. ans = response["message"]["content"].strip()
  472. return ans, response.get("eval_count", 0) + response.get("prompt_eval_count", 0)
  473. except Exception as e:
  474. return "**ERROR**: " + str(e), 0
  475. def chat_streamly(self, system, history, gen_conf):
  476. if system:
  477. history.insert(0, {"role": "system", "content": system})
  478. if "max_tokens" in gen_conf:
  479. del gen_conf["max_tokens"]
  480. options = {}
  481. if "temperature" in gen_conf:
  482. options["temperature"] = gen_conf["temperature"]
  483. if "max_tokens" in gen_conf:
  484. options["num_predict"] = gen_conf["max_tokens"]
  485. if "top_p" in gen_conf:
  486. options["top_p"] = gen_conf["top_p"]
  487. if "presence_penalty" in gen_conf:
  488. options["presence_penalty"] = gen_conf["presence_penalty"]
  489. if "frequency_penalty" in gen_conf:
  490. options["frequency_penalty"] = gen_conf["frequency_penalty"]
  491. ans = ""
  492. try:
  493. response = self.client.chat(
  494. model=self.model_name,
  495. messages=history,
  496. stream=True,
  497. options=options,
  498. keep_alive=-1
  499. )
  500. for resp in response:
  501. if resp["done"]:
  502. yield resp.get("prompt_eval_count", 0) + resp.get("eval_count", 0)
  503. ans += resp["message"]["content"]
  504. yield ans
  505. except Exception as e:
  506. yield ans + "\n**ERROR**: " + str(e)
  507. yield 0
  508. class LocalAIChat(Base):
  509. def __init__(self, key, model_name, base_url):
  510. if not base_url:
  511. raise ValueError("Local llm url cannot be None")
  512. if base_url.split("/")[-1] != "v1":
  513. base_url = os.path.join(base_url, "v1")
  514. self.client = OpenAI(api_key="empty", base_url=base_url)
  515. self.model_name = model_name.split("___")[0]
  516. class LocalLLM(Base):
  517. class RPCProxy:
  518. def __init__(self, host, port):
  519. self.host = host
  520. self.port = int(port)
  521. self.__conn()
  522. def __conn(self):
  523. from multiprocessing.connection import Client
  524. self._connection = Client(
  525. (self.host, self.port), authkey=b"infiniflow-token4kevinhu"
  526. )
  527. def __getattr__(self, name):
  528. import pickle
  529. def do_rpc(*args, **kwargs):
  530. for _ in range(3):
  531. try:
  532. self._connection.send(pickle.dumps((name, args, kwargs)))
  533. return pickle.loads(self._connection.recv())
  534. except Exception:
  535. self.__conn()
  536. raise Exception("RPC connection lost!")
  537. return do_rpc
  538. def __init__(self, key, model_name):
  539. from jina import Client
  540. self.client = Client(port=12345, protocol="grpc", asyncio=True)
  541. def _prepare_prompt(self, system, history, gen_conf):
  542. from rag.svr.jina_server import Prompt
  543. if system:
  544. history.insert(0, {"role": "system", "content": system})
  545. return Prompt(message=history, gen_conf=gen_conf)
  546. def _stream_response(self, endpoint, prompt):
  547. from rag.svr.jina_server import Generation
  548. answer = ""
  549. try:
  550. res = self.client.stream_doc(
  551. on=endpoint, inputs=prompt, return_type=Generation
  552. )
  553. loop = asyncio.get_event_loop()
  554. try:
  555. while True:
  556. answer = loop.run_until_complete(res.__anext__()).text
  557. yield answer
  558. except StopAsyncIteration:
  559. pass
  560. except Exception as e:
  561. yield answer + "\n**ERROR**: " + str(e)
  562. yield num_tokens_from_string(answer)
  563. def chat(self, system, history, gen_conf):
  564. if "max_tokens" in gen_conf:
  565. del gen_conf["max_tokens"]
  566. prompt = self._prepare_prompt(system, history, gen_conf)
  567. chat_gen = self._stream_response("/chat", prompt)
  568. ans = next(chat_gen)
  569. total_tokens = next(chat_gen)
  570. return ans, total_tokens
  571. def chat_streamly(self, system, history, gen_conf):
  572. if "max_tokens" in gen_conf:
  573. del gen_conf["max_tokens"]
  574. prompt = self._prepare_prompt(system, history, gen_conf)
  575. return self._stream_response("/stream", prompt)
  576. class VolcEngineChat(Base):
  577. def __init__(self, key, model_name, base_url='https://ark.cn-beijing.volces.com/api/v3'):
  578. """
  579. Since do not want to modify the original database fields, and the VolcEngine authentication method is quite special,
  580. Assemble ark_api_key, ep_id into api_key, store it as a dictionary type, and parse it for use
  581. model_name is for display only
  582. """
  583. base_url = base_url if base_url else 'https://ark.cn-beijing.volces.com/api/v3'
  584. ark_api_key = json.loads(key).get('ark_api_key', '')
  585. model_name = json.loads(key).get('ep_id', '') + json.loads(key).get('endpoint_id', '')
  586. super().__init__(ark_api_key, model_name, base_url)
  587. class MiniMaxChat(Base):
  588. def __init__(
  589. self,
  590. key,
  591. model_name,
  592. base_url="https://api.minimax.chat/v1/text/chatcompletion_v2",
  593. ):
  594. if not base_url:
  595. base_url = "https://api.minimax.chat/v1/text/chatcompletion_v2"
  596. self.base_url = base_url
  597. self.model_name = model_name
  598. self.api_key = key
  599. def chat(self, system, history, gen_conf):
  600. if system:
  601. history.insert(0, {"role": "system", "content": system})
  602. for k in list(gen_conf.keys()):
  603. if k not in ["temperature", "top_p", "max_tokens"]:
  604. del gen_conf[k]
  605. headers = {
  606. "Authorization": f"Bearer {self.api_key}",
  607. "Content-Type": "application/json",
  608. }
  609. payload = json.dumps(
  610. {"model": self.model_name, "messages": history, **gen_conf}
  611. )
  612. try:
  613. response = requests.request(
  614. "POST", url=self.base_url, headers=headers, data=payload
  615. )
  616. response = response.json()
  617. ans = response["choices"][0]["message"]["content"].strip()
  618. if response["choices"][0]["finish_reason"] == "length":
  619. if is_chinese(ans):
  620. ans += LENGTH_NOTIFICATION_CN
  621. else:
  622. ans += LENGTH_NOTIFICATION_EN
  623. return ans, self.total_token_count(response)
  624. except Exception as e:
  625. return "**ERROR**: " + str(e), 0
  626. def chat_streamly(self, system, history, gen_conf):
  627. if system:
  628. history.insert(0, {"role": "system", "content": system})
  629. for k in list(gen_conf.keys()):
  630. if k not in ["temperature", "top_p", "max_tokens"]:
  631. del gen_conf[k]
  632. ans = ""
  633. total_tokens = 0
  634. try:
  635. headers = {
  636. "Authorization": f"Bearer {self.api_key}",
  637. "Content-Type": "application/json",
  638. }
  639. payload = json.dumps(
  640. {
  641. "model": self.model_name,
  642. "messages": history,
  643. "stream": True,
  644. **gen_conf,
  645. }
  646. )
  647. response = requests.request(
  648. "POST",
  649. url=self.base_url,
  650. headers=headers,
  651. data=payload,
  652. )
  653. for resp in response.text.split("\n\n")[:-1]:
  654. resp = json.loads(resp[6:])
  655. text = ""
  656. if "choices" in resp and "delta" in resp["choices"][0]:
  657. text = resp["choices"][0]["delta"]["content"]
  658. ans += text
  659. tol = self.total_token_count(resp)
  660. if not tol:
  661. total_tokens += num_tokens_from_string(text)
  662. else:
  663. total_tokens = tol
  664. yield ans
  665. except Exception as e:
  666. yield ans + "\n**ERROR**: " + str(e)
  667. yield total_tokens
  668. class MistralChat(Base):
  669. def __init__(self, key, model_name, base_url=None):
  670. from mistralai.client import MistralClient
  671. self.client = MistralClient(api_key=key)
  672. self.model_name = model_name
  673. def chat(self, system, history, gen_conf):
  674. if system:
  675. history.insert(0, {"role": "system", "content": system})
  676. for k in list(gen_conf.keys()):
  677. if k not in ["temperature", "top_p", "max_tokens"]:
  678. del gen_conf[k]
  679. try:
  680. response = self.client.chat(
  681. model=self.model_name,
  682. messages=history,
  683. **gen_conf)
  684. ans = response.choices[0].message.content
  685. if response.choices[0].finish_reason == "length":
  686. if is_chinese(ans):
  687. ans += LENGTH_NOTIFICATION_CN
  688. else:
  689. ans += LENGTH_NOTIFICATION_EN
  690. return ans, self.total_token_count(response)
  691. except openai.APIError as e:
  692. return "**ERROR**: " + str(e), 0
  693. def chat_streamly(self, system, history, gen_conf):
  694. if system:
  695. history.insert(0, {"role": "system", "content": system})
  696. for k in list(gen_conf.keys()):
  697. if k not in ["temperature", "top_p", "max_tokens"]:
  698. del gen_conf[k]
  699. ans = ""
  700. total_tokens = 0
  701. try:
  702. response = self.client.chat_stream(
  703. model=self.model_name,
  704. messages=history,
  705. **gen_conf)
  706. for resp in response:
  707. if not resp.choices or not resp.choices[0].delta.content:
  708. continue
  709. ans += resp.choices[0].delta.content
  710. total_tokens += 1
  711. if resp.choices[0].finish_reason == "length":
  712. if is_chinese(ans):
  713. ans += LENGTH_NOTIFICATION_CN
  714. else:
  715. ans += LENGTH_NOTIFICATION_EN
  716. yield ans
  717. except openai.APIError as e:
  718. yield ans + "\n**ERROR**: " + str(e)
  719. yield total_tokens
  720. class BedrockChat(Base):
  721. def __init__(self, key, model_name, **kwargs):
  722. import boto3
  723. self.bedrock_ak = json.loads(key).get('bedrock_ak', '')
  724. self.bedrock_sk = json.loads(key).get('bedrock_sk', '')
  725. self.bedrock_region = json.loads(key).get('bedrock_region', '')
  726. self.model_name = model_name
  727. if self.bedrock_ak == '' or self.bedrock_sk == '' or self.bedrock_region == '':
  728. # Try to create a client using the default credentials (AWS_PROFILE, AWS_DEFAULT_REGION, etc.)
  729. self.client = boto3.client('bedrock-runtime')
  730. else:
  731. self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
  732. aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
  733. def chat(self, system, history, gen_conf):
  734. from botocore.exceptions import ClientError
  735. for k in list(gen_conf.keys()):
  736. if k not in ["top_p", "max_tokens"]:
  737. del gen_conf[k]
  738. for item in history:
  739. if not isinstance(item["content"], list) and not isinstance(item["content"], tuple):
  740. item["content"] = [{"text": item["content"]}]
  741. try:
  742. # Send the message to the model, using a basic inference configuration.
  743. response = self.client.converse(
  744. modelId=self.model_name,
  745. messages=history,
  746. inferenceConfig=gen_conf,
  747. system=[{"text": (system if system else "Answer the user's message.")}],
  748. )
  749. # Extract and print the response text.
  750. ans = response["output"]["message"]["content"][0]["text"]
  751. return ans, num_tokens_from_string(ans)
  752. except (ClientError, Exception) as e:
  753. return f"ERROR: Can't invoke '{self.model_name}'. Reason: {e}", 0
  754. def chat_streamly(self, system, history, gen_conf):
  755. from botocore.exceptions import ClientError
  756. for k in list(gen_conf.keys()):
  757. if k not in ["top_p", "max_tokens"]:
  758. del gen_conf[k]
  759. for item in history:
  760. if not isinstance(item["content"], list) and not isinstance(item["content"], tuple):
  761. item["content"] = [{"text": item["content"]}]
  762. if self.model_name.split('.')[0] == 'ai21':
  763. try:
  764. response = self.client.converse(
  765. modelId=self.model_name,
  766. messages=history,
  767. inferenceConfig=gen_conf,
  768. system=[{"text": (system if system else "Answer the user's message.")}]
  769. )
  770. ans = response["output"]["message"]["content"][0]["text"]
  771. return ans, num_tokens_from_string(ans)
  772. except (ClientError, Exception) as e:
  773. return f"ERROR: Can't invoke '{self.model_name}'. Reason: {e}", 0
  774. ans = ""
  775. try:
  776. # Send the message to the model, using a basic inference configuration.
  777. streaming_response = self.client.converse_stream(
  778. modelId=self.model_name,
  779. messages=history,
  780. inferenceConfig=gen_conf,
  781. system=[{"text": (system if system else "Answer the user's message.")}]
  782. )
  783. # Extract and print the streamed response text in real-time.
  784. for resp in streaming_response["stream"]:
  785. if "contentBlockDelta" in resp:
  786. ans += resp["contentBlockDelta"]["delta"]["text"]
  787. yield ans
  788. except (ClientError, Exception) as e:
  789. yield ans + f"ERROR: Can't invoke '{self.model_name}'. Reason: {e}"
  790. yield num_tokens_from_string(ans)
  791. class GeminiChat(Base):
  792. def __init__(self, key, model_name, base_url=None):
  793. from google.generativeai import client, GenerativeModel
  794. client.configure(api_key=key)
  795. _client = client.get_default_generative_client()
  796. self.model_name = 'models/' + model_name
  797. self.model = GenerativeModel(model_name=self.model_name)
  798. self.model._client = _client
  799. def chat(self, system, history, gen_conf):
  800. from google.generativeai.types import content_types
  801. if system:
  802. self.model._system_instruction = content_types.to_content(system)
  803. for k in list(gen_conf.keys()):
  804. if k not in ["temperature", "top_p", "max_tokens"]:
  805. del gen_conf[k]
  806. for item in history:
  807. if 'role' in item and item['role'] == 'assistant':
  808. item['role'] = 'model'
  809. if 'role' in item and item['role'] == 'system':
  810. item['role'] = 'user'
  811. if 'content' in item:
  812. item['parts'] = item.pop('content')
  813. try:
  814. response = self.model.generate_content(
  815. history,
  816. generation_config=gen_conf)
  817. ans = response.text
  818. return ans, response.usage_metadata.total_token_count
  819. except Exception as e:
  820. return "**ERROR**: " + str(e), 0
  821. def chat_streamly(self, system, history, gen_conf):
  822. from google.generativeai.types import content_types
  823. if system:
  824. self.model._system_instruction = content_types.to_content(system)
  825. for k in list(gen_conf.keys()):
  826. if k not in ["temperature", "top_p", "max_tokens"]:
  827. del gen_conf[k]
  828. for item in history:
  829. if 'role' in item and item['role'] == 'assistant':
  830. item['role'] = 'model'
  831. if 'content' in item:
  832. item['parts'] = item.pop('content')
  833. ans = ""
  834. try:
  835. response = self.model.generate_content(
  836. history,
  837. generation_config=gen_conf, stream=True)
  838. for resp in response:
  839. ans += resp.text
  840. yield ans
  841. yield response._chunks[-1].usage_metadata.total_token_count
  842. except Exception as e:
  843. yield ans + "\n**ERROR**: " + str(e)
  844. yield 0
  845. class GroqChat(Base):
  846. def __init__(self, key, model_name, base_url=''):
  847. from groq import Groq
  848. self.client = Groq(api_key=key)
  849. self.model_name = model_name
  850. def chat(self, system, history, gen_conf):
  851. if system:
  852. history.insert(0, {"role": "system", "content": system})
  853. for k in list(gen_conf.keys()):
  854. if k not in ["temperature", "top_p", "max_tokens"]:
  855. del gen_conf[k]
  856. ans = ""
  857. try:
  858. response = self.client.chat.completions.create(
  859. model=self.model_name,
  860. messages=history,
  861. **gen_conf
  862. )
  863. ans = response.choices[0].message.content
  864. if response.choices[0].finish_reason == "length":
  865. if is_chinese(ans):
  866. ans += LENGTH_NOTIFICATION_CN
  867. else:
  868. ans += LENGTH_NOTIFICATION_EN
  869. return ans, self.total_token_count(response)
  870. except Exception as e:
  871. return ans + "\n**ERROR**: " + str(e), 0
  872. def chat_streamly(self, system, history, gen_conf):
  873. if system:
  874. history.insert(0, {"role": "system", "content": system})
  875. for k in list(gen_conf.keys()):
  876. if k not in ["temperature", "top_p", "max_tokens"]:
  877. del gen_conf[k]
  878. ans = ""
  879. total_tokens = 0
  880. try:
  881. response = self.client.chat.completions.create(
  882. model=self.model_name,
  883. messages=history,
  884. stream=True,
  885. **gen_conf
  886. )
  887. for resp in response:
  888. if not resp.choices or not resp.choices[0].delta.content:
  889. continue
  890. ans += resp.choices[0].delta.content
  891. total_tokens += 1
  892. if resp.choices[0].finish_reason == "length":
  893. if is_chinese(ans):
  894. ans += LENGTH_NOTIFICATION_CN
  895. else:
  896. ans += LENGTH_NOTIFICATION_EN
  897. yield ans
  898. except Exception as e:
  899. yield ans + "\n**ERROR**: " + str(e)
  900. yield total_tokens
  901. ## openrouter
  902. class OpenRouterChat(Base):
  903. def __init__(self, key, model_name, base_url="https://openrouter.ai/api/v1"):
  904. if not base_url:
  905. base_url = "https://openrouter.ai/api/v1"
  906. super().__init__(key, model_name, base_url)
  907. class StepFunChat(Base):
  908. def __init__(self, key, model_name, base_url="https://api.stepfun.com/v1"):
  909. if not base_url:
  910. base_url = "https://api.stepfun.com/v1"
  911. super().__init__(key, model_name, base_url)
  912. class NvidiaChat(Base):
  913. def __init__(self, key, model_name, base_url="https://integrate.api.nvidia.com/v1"):
  914. if not base_url:
  915. base_url = "https://integrate.api.nvidia.com/v1"
  916. super().__init__(key, model_name, base_url)
  917. class LmStudioChat(Base):
  918. def __init__(self, key, model_name, base_url):
  919. if not base_url:
  920. raise ValueError("Local llm url cannot be None")
  921. if base_url.split("/")[-1] != "v1":
  922. base_url = os.path.join(base_url, "v1")
  923. self.client = OpenAI(api_key="lm-studio", base_url=base_url)
  924. self.model_name = model_name
  925. class OpenAI_APIChat(Base):
  926. def __init__(self, key, model_name, base_url):
  927. if not base_url:
  928. raise ValueError("url cannot be None")
  929. model_name = model_name.split("___")[0]
  930. super().__init__(key, model_name, base_url)
  931. class PPIOChat(Base):
  932. def __init__(self, key, model_name, base_url="https://api.ppinfra.com/v3/openai"):
  933. if not base_url:
  934. base_url = "https://api.ppinfra.com/v3/openai"
  935. super().__init__(key, model_name, base_url)
  936. class CoHereChat(Base):
  937. def __init__(self, key, model_name, base_url=""):
  938. from cohere import Client
  939. self.client = Client(api_key=key)
  940. self.model_name = model_name
  941. def chat(self, system, history, gen_conf):
  942. if system:
  943. history.insert(0, {"role": "system", "content": system})
  944. if "max_tokens" in gen_conf:
  945. del gen_conf["max_tokens"]
  946. if "top_p" in gen_conf:
  947. gen_conf["p"] = gen_conf.pop("top_p")
  948. if "frequency_penalty" in gen_conf and "presence_penalty" in gen_conf:
  949. gen_conf.pop("presence_penalty")
  950. for item in history:
  951. if "role" in item and item["role"] == "user":
  952. item["role"] = "USER"
  953. if "role" in item and item["role"] == "assistant":
  954. item["role"] = "CHATBOT"
  955. if "content" in item:
  956. item["message"] = item.pop("content")
  957. mes = history.pop()["message"]
  958. ans = ""
  959. try:
  960. response = self.client.chat(
  961. model=self.model_name, chat_history=history, message=mes, **gen_conf
  962. )
  963. ans = response.text
  964. if response.finish_reason == "MAX_TOKENS":
  965. ans += (
  966. "...\nFor the content length reason, it stopped, continue?"
  967. if is_english([ans])
  968. else "······\n由于长度的原因,回答被截断了,要继续吗?"
  969. )
  970. return (
  971. ans,
  972. response.meta.tokens.input_tokens + response.meta.tokens.output_tokens,
  973. )
  974. except Exception as e:
  975. return ans + "\n**ERROR**: " + str(e), 0
  976. def chat_streamly(self, system, history, gen_conf):
  977. if system:
  978. history.insert(0, {"role": "system", "content": system})
  979. if "max_tokens" in gen_conf:
  980. del gen_conf["max_tokens"]
  981. if "top_p" in gen_conf:
  982. gen_conf["p"] = gen_conf.pop("top_p")
  983. if "frequency_penalty" in gen_conf and "presence_penalty" in gen_conf:
  984. gen_conf.pop("presence_penalty")
  985. for item in history:
  986. if "role" in item and item["role"] == "user":
  987. item["role"] = "USER"
  988. if "role" in item and item["role"] == "assistant":
  989. item["role"] = "CHATBOT"
  990. if "content" in item:
  991. item["message"] = item.pop("content")
  992. mes = history.pop()["message"]
  993. ans = ""
  994. total_tokens = 0
  995. try:
  996. response = self.client.chat_stream(
  997. model=self.model_name, chat_history=history, message=mes, **gen_conf
  998. )
  999. for resp in response:
  1000. if resp.event_type == "text-generation":
  1001. ans += resp.text
  1002. total_tokens += num_tokens_from_string(resp.text)
  1003. elif resp.event_type == "stream-end":
  1004. if resp.finish_reason == "MAX_TOKENS":
  1005. ans += (
  1006. "...\nFor the content length reason, it stopped, continue?"
  1007. if is_english([ans])
  1008. else "······\n由于长度的原因,回答被截断了,要继续吗?"
  1009. )
  1010. yield ans
  1011. except Exception as e:
  1012. yield ans + "\n**ERROR**: " + str(e)
  1013. yield total_tokens
  1014. class LeptonAIChat(Base):
  1015. def __init__(self, key, model_name, base_url=None):
  1016. if not base_url:
  1017. base_url = os.path.join("https://" + model_name + ".lepton.run", "api", "v1")
  1018. super().__init__(key, model_name, base_url)
  1019. class TogetherAIChat(Base):
  1020. def __init__(self, key, model_name, base_url="https://api.together.xyz/v1"):
  1021. if not base_url:
  1022. base_url = "https://api.together.xyz/v1"
  1023. super().__init__(key, model_name, base_url)
  1024. class PerfXCloudChat(Base):
  1025. def __init__(self, key, model_name, base_url="https://cloud.perfxlab.cn/v1"):
  1026. if not base_url:
  1027. base_url = "https://cloud.perfxlab.cn/v1"
  1028. super().__init__(key, model_name, base_url)
  1029. class UpstageChat(Base):
  1030. def __init__(self, key, model_name, base_url="https://api.upstage.ai/v1/solar"):
  1031. if not base_url:
  1032. base_url = "https://api.upstage.ai/v1/solar"
  1033. super().__init__(key, model_name, base_url)
  1034. class NovitaAIChat(Base):
  1035. def __init__(self, key, model_name, base_url="https://api.novita.ai/v3/openai"):
  1036. if not base_url:
  1037. base_url = "https://api.novita.ai/v3/openai"
  1038. super().__init__(key, model_name, base_url)
  1039. class SILICONFLOWChat(Base):
  1040. def __init__(self, key, model_name, base_url="https://api.siliconflow.cn/v1"):
  1041. if not base_url:
  1042. base_url = "https://api.siliconflow.cn/v1"
  1043. super().__init__(key, model_name, base_url)
  1044. class YiChat(Base):
  1045. def __init__(self, key, model_name, base_url="https://api.lingyiwanwu.com/v1"):
  1046. if not base_url:
  1047. base_url = "https://api.lingyiwanwu.com/v1"
  1048. super().__init__(key, model_name, base_url)
  1049. class ReplicateChat(Base):
  1050. def __init__(self, key, model_name, base_url=None):
  1051. from replicate.client import Client
  1052. self.model_name = model_name
  1053. self.client = Client(api_token=key)
  1054. self.system = ""
  1055. def chat(self, system, history, gen_conf):
  1056. if "max_tokens" in gen_conf:
  1057. del gen_conf["max_tokens"]
  1058. if system:
  1059. self.system = system
  1060. prompt = "\n".join(
  1061. [item["role"] + ":" + item["content"] for item in history[-5:]]
  1062. )
  1063. ans = ""
  1064. try:
  1065. response = self.client.run(
  1066. self.model_name,
  1067. input={"system_prompt": self.system, "prompt": prompt, **gen_conf},
  1068. )
  1069. ans = "".join(response)
  1070. return ans, num_tokens_from_string(ans)
  1071. except Exception as e:
  1072. return ans + "\n**ERROR**: " + str(e), 0
  1073. def chat_streamly(self, system, history, gen_conf):
  1074. if "max_tokens" in gen_conf:
  1075. del gen_conf["max_tokens"]
  1076. if system:
  1077. self.system = system
  1078. prompt = "\n".join(
  1079. [item["role"] + ":" + item["content"] for item in history[-5:]]
  1080. )
  1081. ans = ""
  1082. try:
  1083. response = self.client.run(
  1084. self.model_name,
  1085. input={"system_prompt": self.system, "prompt": prompt, **gen_conf},
  1086. )
  1087. for resp in response:
  1088. ans += resp
  1089. yield ans
  1090. except Exception as e:
  1091. yield ans + "\n**ERROR**: " + str(e)
  1092. yield num_tokens_from_string(ans)
  1093. class HunyuanChat(Base):
  1094. def __init__(self, key, model_name, base_url=None):
  1095. from tencentcloud.common import credential
  1096. from tencentcloud.hunyuan.v20230901 import hunyuan_client
  1097. key = json.loads(key)
  1098. sid = key.get("hunyuan_sid", "")
  1099. sk = key.get("hunyuan_sk", "")
  1100. cred = credential.Credential(sid, sk)
  1101. self.model_name = model_name
  1102. self.client = hunyuan_client.HunyuanClient(cred, "")
  1103. def chat(self, system, history, gen_conf):
  1104. from tencentcloud.hunyuan.v20230901 import models
  1105. from tencentcloud.common.exception.tencent_cloud_sdk_exception import (
  1106. TencentCloudSDKException,
  1107. )
  1108. _gen_conf = {}
  1109. _history = [{k.capitalize(): v for k, v in item.items()} for item in history]
  1110. if system:
  1111. _history.insert(0, {"Role": "system", "Content": system})
  1112. if "max_tokens" in gen_conf:
  1113. del gen_conf["max_tokens"]
  1114. if "temperature" in gen_conf:
  1115. _gen_conf["Temperature"] = gen_conf["temperature"]
  1116. if "top_p" in gen_conf:
  1117. _gen_conf["TopP"] = gen_conf["top_p"]
  1118. req = models.ChatCompletionsRequest()
  1119. params = {"Model": self.model_name, "Messages": _history, **_gen_conf}
  1120. req.from_json_string(json.dumps(params))
  1121. ans = ""
  1122. try:
  1123. response = self.client.ChatCompletions(req)
  1124. ans = response.Choices[0].Message.Content
  1125. return ans, response.Usage.TotalTokens
  1126. except TencentCloudSDKException as e:
  1127. return ans + "\n**ERROR**: " + str(e), 0
  1128. def chat_streamly(self, system, history, gen_conf):
  1129. from tencentcloud.hunyuan.v20230901 import models
  1130. from tencentcloud.common.exception.tencent_cloud_sdk_exception import (
  1131. TencentCloudSDKException,
  1132. )
  1133. _gen_conf = {}
  1134. _history = [{k.capitalize(): v for k, v in item.items()} for item in history]
  1135. if system:
  1136. _history.insert(0, {"Role": "system", "Content": system})
  1137. if "max_tokens" in gen_conf:
  1138. del gen_conf["max_tokens"]
  1139. if "temperature" in gen_conf:
  1140. _gen_conf["Temperature"] = gen_conf["temperature"]
  1141. if "top_p" in gen_conf:
  1142. _gen_conf["TopP"] = gen_conf["top_p"]
  1143. req = models.ChatCompletionsRequest()
  1144. params = {
  1145. "Model": self.model_name,
  1146. "Messages": _history,
  1147. "Stream": True,
  1148. **_gen_conf,
  1149. }
  1150. req.from_json_string(json.dumps(params))
  1151. ans = ""
  1152. total_tokens = 0
  1153. try:
  1154. response = self.client.ChatCompletions(req)
  1155. for resp in response:
  1156. resp = json.loads(resp["data"])
  1157. if not resp["Choices"] or not resp["Choices"][0]["Delta"]["Content"]:
  1158. continue
  1159. ans += resp["Choices"][0]["Delta"]["Content"]
  1160. total_tokens += 1
  1161. yield ans
  1162. except TencentCloudSDKException as e:
  1163. yield ans + "\n**ERROR**: " + str(e)
  1164. yield total_tokens
  1165. class SparkChat(Base):
  1166. def __init__(
  1167. self, key, model_name, base_url="https://spark-api-open.xf-yun.com/v1"
  1168. ):
  1169. if not base_url:
  1170. base_url = "https://spark-api-open.xf-yun.com/v1"
  1171. model2version = {
  1172. "Spark-Max": "generalv3.5",
  1173. "Spark-Lite": "general",
  1174. "Spark-Pro": "generalv3",
  1175. "Spark-Pro-128K": "pro-128k",
  1176. "Spark-4.0-Ultra": "4.0Ultra",
  1177. }
  1178. version2model = {v: k for k, v in model2version.items()}
  1179. assert model_name in model2version or model_name in version2model, f"The given model name is not supported yet. Support: {list(model2version.keys())}"
  1180. if model_name in model2version:
  1181. model_version = model2version[model_name]
  1182. else:
  1183. model_version = model_name
  1184. super().__init__(key, model_version, base_url)
  1185. class BaiduYiyanChat(Base):
  1186. def __init__(self, key, model_name, base_url=None):
  1187. import qianfan
  1188. key = json.loads(key)
  1189. ak = key.get("yiyan_ak", "")
  1190. sk = key.get("yiyan_sk", "")
  1191. self.client = qianfan.ChatCompletion(ak=ak, sk=sk)
  1192. self.model_name = model_name.lower()
  1193. self.system = ""
  1194. def chat(self, system, history, gen_conf):
  1195. if system:
  1196. self.system = system
  1197. gen_conf["penalty_score"] = (
  1198. (gen_conf.get("presence_penalty", 0) + gen_conf.get("frequency_penalty",
  1199. 0)) / 2
  1200. ) + 1
  1201. if "max_tokens" in gen_conf:
  1202. del gen_conf["max_tokens"]
  1203. ans = ""
  1204. try:
  1205. response = self.client.do(
  1206. model=self.model_name,
  1207. messages=history,
  1208. system=self.system,
  1209. **gen_conf
  1210. ).body
  1211. ans = response['result']
  1212. return ans, self.total_token_count(response)
  1213. except Exception as e:
  1214. return ans + "\n**ERROR**: " + str(e), 0
  1215. def chat_streamly(self, system, history, gen_conf):
  1216. if system:
  1217. self.system = system
  1218. gen_conf["penalty_score"] = (
  1219. (gen_conf.get("presence_penalty", 0) + gen_conf.get("frequency_penalty",
  1220. 0)) / 2
  1221. ) + 1
  1222. if "max_tokens" in gen_conf:
  1223. del gen_conf["max_tokens"]
  1224. ans = ""
  1225. total_tokens = 0
  1226. try:
  1227. response = self.client.do(
  1228. model=self.model_name,
  1229. messages=history,
  1230. system=self.system,
  1231. stream=True,
  1232. **gen_conf
  1233. )
  1234. for resp in response:
  1235. resp = resp.body
  1236. ans += resp['result']
  1237. total_tokens = self.total_token_count(resp)
  1238. yield ans
  1239. except Exception as e:
  1240. return ans + "\n**ERROR**: " + str(e), 0
  1241. yield total_tokens
  1242. class AnthropicChat(Base):
  1243. def __init__(self, key, model_name, base_url=None):
  1244. import anthropic
  1245. self.client = anthropic.Anthropic(api_key=key)
  1246. self.model_name = model_name
  1247. self.system = ""
  1248. def chat(self, system, history, gen_conf):
  1249. if system:
  1250. self.system = system
  1251. if "presence_penalty" in gen_conf:
  1252. del gen_conf["presence_penalty"]
  1253. if "frequency_penalty" in gen_conf:
  1254. del gen_conf["frequency_penalty"]
  1255. gen_conf["max_tokens"] = 8196
  1256. if "haiku" in self.model_name or "opus" in self.model_name:
  1257. gen_conf["max_tokens"] = 4096
  1258. ans = ""
  1259. try:
  1260. response = self.client.messages.create(
  1261. model=self.model_name,
  1262. messages=history,
  1263. system=self.system,
  1264. stream=False,
  1265. **gen_conf,
  1266. ).to_dict()
  1267. ans = response["content"][0]["text"]
  1268. if response["stop_reason"] == "max_tokens":
  1269. ans += (
  1270. "...\nFor the content length reason, it stopped, continue?"
  1271. if is_english([ans])
  1272. else "······\n由于长度的原因,回答被截断了,要继续吗?"
  1273. )
  1274. return (
  1275. ans,
  1276. response["usage"]["input_tokens"] + response["usage"]["output_tokens"],
  1277. )
  1278. except Exception as e:
  1279. return ans + "\n**ERROR**: " + str(e), 0
  1280. def chat_streamly(self, system, history, gen_conf):
  1281. if system:
  1282. self.system = system
  1283. if "presence_penalty" in gen_conf:
  1284. del gen_conf["presence_penalty"]
  1285. if "frequency_penalty" in gen_conf:
  1286. del gen_conf["frequency_penalty"]
  1287. gen_conf["max_tokens"] = 8196
  1288. if "haiku" in self.model_name or "opus" in self.model_name:
  1289. gen_conf["max_tokens"] = 4096
  1290. ans = ""
  1291. total_tokens = 0
  1292. try:
  1293. response = self.client.messages.create(
  1294. model=self.model_name,
  1295. messages=history,
  1296. system=system,
  1297. stream=True,
  1298. **gen_conf,
  1299. )
  1300. for res in response:
  1301. if res.type == 'content_block_delta':
  1302. if res.delta.type == "thinking_delta" and res.delta.thinking:
  1303. if ans.find("<think>") < 0:
  1304. ans += "<think>"
  1305. ans = ans.replace("</think>", "")
  1306. ans += res.delta.thinking + "</think>"
  1307. else:
  1308. text = res.delta.text
  1309. ans += text
  1310. total_tokens += num_tokens_from_string(text)
  1311. yield ans
  1312. except Exception as e:
  1313. yield ans + "\n**ERROR**: " + str(e)
  1314. yield total_tokens
  1315. class GoogleChat(Base):
  1316. def __init__(self, key, model_name, base_url=None):
  1317. from google.oauth2 import service_account
  1318. import base64
  1319. key = json.loads(key)
  1320. access_token = json.loads(
  1321. base64.b64decode(key.get("google_service_account_key", ""))
  1322. )
  1323. project_id = key.get("google_project_id", "")
  1324. region = key.get("google_region", "")
  1325. scopes = ["https://www.googleapis.com/auth/cloud-platform"]
  1326. self.model_name = model_name
  1327. self.system = ""
  1328. if "claude" in self.model_name:
  1329. from anthropic import AnthropicVertex
  1330. from google.auth.transport.requests import Request
  1331. if access_token:
  1332. credits = service_account.Credentials.from_service_account_info(
  1333. access_token, scopes=scopes
  1334. )
  1335. request = Request()
  1336. credits.refresh(request)
  1337. token = credits.token
  1338. self.client = AnthropicVertex(
  1339. region=region, project_id=project_id, access_token=token
  1340. )
  1341. else:
  1342. self.client = AnthropicVertex(region=region, project_id=project_id)
  1343. else:
  1344. from google.cloud import aiplatform
  1345. import vertexai.generative_models as glm
  1346. if access_token:
  1347. credits = service_account.Credentials.from_service_account_info(
  1348. access_token
  1349. )
  1350. aiplatform.init(
  1351. credentials=credits, project=project_id, location=region
  1352. )
  1353. else:
  1354. aiplatform.init(project=project_id, location=region)
  1355. self.client = glm.GenerativeModel(model_name=self.model_name)
  1356. def chat(self, system, history, gen_conf):
  1357. if system:
  1358. self.system = system
  1359. if "claude" in self.model_name:
  1360. if "max_tokens" in gen_conf:
  1361. del gen_conf["max_tokens"]
  1362. try:
  1363. response = self.client.messages.create(
  1364. model=self.model_name,
  1365. messages=history,
  1366. system=self.system,
  1367. stream=False,
  1368. **gen_conf,
  1369. ).json()
  1370. ans = response["content"][0]["text"]
  1371. if response["stop_reason"] == "max_tokens":
  1372. ans += (
  1373. "...\nFor the content length reason, it stopped, continue?"
  1374. if is_english([ans])
  1375. else "······\n由于长度的原因,回答被截断了,要继续吗?"
  1376. )
  1377. return (
  1378. ans,
  1379. response["usage"]["input_tokens"]
  1380. + response["usage"]["output_tokens"],
  1381. )
  1382. except Exception as e:
  1383. return "\n**ERROR**: " + str(e), 0
  1384. else:
  1385. self.client._system_instruction = self.system
  1386. if "max_tokens" in gen_conf:
  1387. gen_conf["max_output_tokens"] = gen_conf["max_tokens"]
  1388. for k in list(gen_conf.keys()):
  1389. if k not in ["temperature", "top_p", "max_output_tokens"]:
  1390. del gen_conf[k]
  1391. for item in history:
  1392. if "role" in item and item["role"] == "assistant":
  1393. item["role"] = "model"
  1394. if "content" in item:
  1395. item["parts"] = item.pop("content")
  1396. try:
  1397. response = self.client.generate_content(
  1398. history, generation_config=gen_conf
  1399. )
  1400. ans = response.text
  1401. return ans, response.usage_metadata.total_token_count
  1402. except Exception as e:
  1403. return "**ERROR**: " + str(e), 0
  1404. def chat_streamly(self, system, history, gen_conf):
  1405. if system:
  1406. self.system = system
  1407. if "claude" in self.model_name:
  1408. if "max_tokens" in gen_conf:
  1409. del gen_conf["max_tokens"]
  1410. ans = ""
  1411. total_tokens = 0
  1412. try:
  1413. response = self.client.messages.create(
  1414. model=self.model_name,
  1415. messages=history,
  1416. system=self.system,
  1417. stream=True,
  1418. **gen_conf,
  1419. )
  1420. for res in response.iter_lines():
  1421. res = res.decode("utf-8")
  1422. if "content_block_delta" in res and "data" in res:
  1423. text = json.loads(res[6:])["delta"]["text"]
  1424. ans += text
  1425. total_tokens += num_tokens_from_string(text)
  1426. except Exception as e:
  1427. yield ans + "\n**ERROR**: " + str(e)
  1428. yield total_tokens
  1429. else:
  1430. self.client._system_instruction = self.system
  1431. if "max_tokens" in gen_conf:
  1432. gen_conf["max_output_tokens"] = gen_conf["max_tokens"]
  1433. for k in list(gen_conf.keys()):
  1434. if k not in ["temperature", "top_p", "max_output_tokens"]:
  1435. del gen_conf[k]
  1436. for item in history:
  1437. if "role" in item and item["role"] == "assistant":
  1438. item["role"] = "model"
  1439. if "content" in item:
  1440. item["parts"] = item.pop("content")
  1441. ans = ""
  1442. try:
  1443. response = self.model.generate_content(
  1444. history, generation_config=gen_conf, stream=True
  1445. )
  1446. for resp in response:
  1447. ans += resp.text
  1448. yield ans
  1449. except Exception as e:
  1450. yield ans + "\n**ERROR**: " + str(e)
  1451. yield response._chunks[-1].usage_metadata.total_token_count
  1452. class GPUStackChat(Base):
  1453. def __init__(self, key=None, model_name="", base_url=""):
  1454. if not base_url:
  1455. raise ValueError("Local llm url cannot be None")
  1456. if base_url.split("/")[-1] != "v1-openai":
  1457. base_url = os.path.join(base_url, "v1-openai")
  1458. super().__init__(key, model_name, base_url)