llm_client.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156
  1. import asyncio
  2. from openai import OpenAI, AsyncOpenAI
  3. import copy
  4. # from qwen_agent.llm.base import LLMBase
  5. import re
  6. class LLMClient:
  7. def __init__(self, model='qwen-plus', api_key='none', model_server=''):
  8. self.model_server = model_server
  9. self._client = OpenAI(
  10. api_key=api_key, base_url=model_server
  11. )
  12. self.model = model
  13. def chat(self, query=None, model=None, stream=False, messages=None, stop=None, functions=None):
  14. if not model:
  15. model = self.model
  16. if stream:
  17. return self._chat_stream(query, model, messages, stop=stop, functions=functions)
  18. else:
  19. return self._chat_no_stream(query, model, messages, stop=stop, functions=functions)
  20. def _chat_stream(self, query, model, messages=None, stop=None, functions=None, **kvargs):
  21. pattern = r'<think>.*?</think>'
  22. print(f'begin: stream to lianqi client, model name: {model}')
  23. if messages:
  24. response = self._client.chat.completions.create(
  25. model=model, messages=messages, stop=stop, stream=True, functions=functions, **kvargs)
  26. else:
  27. response = self._client.chat.completions.create(
  28. model=model,
  29. messages=[{'role': 'user', 'content': query}],
  30. stop=stop,
  31. stream=True,
  32. functions=functions,
  33. **kvargs
  34. )
  35. for chunk in response:
  36. if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
  37. print(chunk.choices[0].delta.content, end='', flush=True)
  38. result = re.sub(pattern, '', chunk.choices[0].delta.content, flags=re.DOTALL)
  39. yield result
  40. def _chat_no_stream(self, query, model=None, messages=None, stop=None, functions=None, **kvargs):
  41. print(f'begin: no stream to lianqi client, model name: {model}')
  42. if not model:
  43. model = self.model
  44. if messages:
  45. if query:
  46. local_message = copy.deepcopy(messages)
  47. local_message.append({'role': 'user','content': query})
  48. _messages = local_message
  49. else:
  50. _messages = messages
  51. response = self._client.chat.completions.create(model=model,
  52. messages=_messages,
  53. stop=stop,
  54. stream=False,
  55. functions=functions,
  56. **kvargs)
  57. else:
  58. response = self._client.chat.completions.create(model=model,
  59. messages=[{
  60. 'role': 'user',
  61. 'content': query
  62. }],
  63. stop=stop,
  64. stream=False,
  65. functions=functions,
  66. **kvargs)
  67. # print(response.choices[0].message.content)
  68. return response.choices[0].message.content
  69. def chat_func(self, model=None, messages=[], functions=None):
  70. print('begin: no stream in lianqi chat_func')
  71. if not model:
  72. model = self.model
  73. if not functions:
  74. functions = []
  75. if functions:
  76. response = self._client.chat.completions.create(
  77. model=model, messages=messages, functions=functions
  78. )
  79. else:
  80. response = self._client.chat.completions.create(model=self, messages=messages)
  81. return response.choices[0].message
  82. class LLMAsyncClient:
  83. def __init__(self, model='qwen-plus', api_key='none', model_server=''):
  84. self.model_server = model_server
  85. self._client = AsyncOpenAI(
  86. api_key=api_key, base_url=model_server
  87. )
  88. self.model = model
  89. async def chat(self, query=None, model=None, stream=False, messages=None, stop=[], functions=None):
  90. if not model:
  91. model = self.model
  92. if stream:
  93. response = self._chat_stream(query, model, messages, stop=stop, functions=functions)
  94. else:
  95. response = self._chat_no_stream(query, model, messages, stop=stop, functions=functions)
  96. # await asyncio.sleep(0.1)
  97. return response
  98. async def _chat_stream(self, query, model, messages=None, stop=[], functions=None, **kvargs):
  99. print(f'begin: stream to lianqi client, model name: {model}')
  100. if messages:
  101. response = await self._client.chat.completions.create(
  102. model=model, messages=messages, stop=stop, stream=True, **kvargs)
  103. else:
  104. response = await self._client.chat.completions.create(
  105. model=model,
  106. messages=[{'role': 'user', 'content': query}],
  107. stop=stop,
  108. stream=True,
  109. functions=functions,
  110. **kvargs
  111. )
  112. async for chunk in response:
  113. if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
  114. print(chunk.choices[0].delta.content, end='', flush=True)
  115. pattern = r'<think>.*?</think>\n*'
  116. result = re.sub(pattern, '', chunk.choices[0].delta.content, flags=re.DOTALL)
  117. yield result
  118. async def _chat_no_stream(self, query, model=None, messages=None, stop=None, functions=None, **kvargs):
  119. print(f'begin: no stream to lianqi client, model name: {model}')
  120. if not model:
  121. model = self.model
  122. if messages:
  123. if query:
  124. local_message = copy.deepcopy(messages)
  125. local_message.append({'role': 'user','content': query})
  126. _messages = local_message
  127. else:
  128. _messages = messages
  129. response = await self._client.chat.completions.create(model=model,
  130. messages=_messages,
  131. stop=stop,
  132. stream=False,
  133. functions=functions,
  134. **kvargs)
  135. else:
  136. response = await self._client.chat.completions.create(model=model,
  137. messages=[{
  138. 'role': 'user',
  139. 'content': query
  140. }],
  141. stop=stop,
  142. stream=False,
  143. functions=functions,
  144. **kvargs)
  145. # print(response.choices[0].message.content)
  146. return response.choices[0].message.content