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.

chat_model.py 77KB

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