wiseflow/client/backend/llms/openai_wrapper.py

71 lines
2.5 KiB
Python
Raw Normal View History

2024-05-09 14:56:58 +08:00
"""
2024-05-08 09:57:48 +08:00
除了openai外很多大模型提供商也都使用openai的SDK对于这一类可以统一使用本wrapper
这里演示使用deepseek提供的DeepSeek-V2
2024-05-09 14:56:58 +08:00
"""
2024-05-08 09:57:48 +08:00
2024-04-07 09:37:47 +08:00
import os
2024-05-08 09:57:48 +08:00
from openai import OpenAI
2024-04-07 09:37:47 +08:00
2024-05-08 09:57:48 +08:00
token = os.environ.get('LLM_API_KEY', "")
if not token:
raise ValueError('请设置环境变量LLM_API_KEY')
2024-04-07 09:37:47 +08:00
2024-05-08 09:57:48 +08:00
base_url = os.environ.get('LLM_API_BASE', "")
2024-04-07 09:37:47 +08:00
2024-05-08 09:57:48 +08:00
client = OpenAI(api_key=token, base_url=base_url)
2024-04-07 09:37:47 +08:00
2024-05-08 09:57:48 +08:00
def openai_llm(messages: list, model: str, logger=None, **kwargs) -> str:
2024-04-07 09:37:47 +08:00
if logger:
logger.debug(f'messages:\n {messages}')
2024-05-08 09:57:48 +08:00
logger.debug(f'model: {model}')
logger.debug(f'kwargs:\n {kwargs}')
2024-05-09 14:56:58 +08:00
try:
response = client.chat.completions.create(messages=messages, model=model, **kwargs)
2024-05-08 09:57:48 +08:00
2024-05-09 14:56:58 +08:00
except Exception as e:
2024-05-08 09:57:48 +08:00
if logger:
2024-05-09 14:56:58 +08:00
logger.error(f'openai_llm error: {e}')
2024-05-08 09:57:48 +08:00
return ''
2024-04-07 09:37:47 +08:00
if logger:
2024-05-08 09:57:48 +08:00
logger.debug(f'result:\n {response.choices[0]}')
2024-04-07 09:37:47 +08:00
logger.debug(f'usage:\n {response.usage}')
return response.choices[0].message.content
if __name__ == '__main__':
import time
from pprint import pprint
# logging.basicConfig(level=logging.DEBUG)
system_content = ''
user_content = '''你是一名优秀的翻译,请帮我把如下新闻标题逐条(一行为一条)翻译为中文,你的输出也必须为一条一行的格式。
The New York Times reported on 2021-01-01 that the COVID-19 cases in China are increasing.
Cyber ops linked to Israel-Hamas conflict largely improvised, researchers say
Russian hackers disrupted Ukrainian electrical grid last year
Reform bill would overhaul controversial surveillance law
GitHub disables pro-Russian hacktivist DDoS pages
Notorious Russian hacking group appears to resurface with fresh cyberattacks on Ukraine
Russian hackers attempted to breach petroleum refining company in NATO country, researchers say
As agencies move towards multi-cloud networks, proactive security is key
Keeping a competitive edge in the cybersecurity game
Mud, sweat and data: The hard work of democratizing data at scale
SEC sues SolarWinds and CISO for fraud
Cyber workforce demand is outpacing supply, survey finds
Four dozen countries declare they won
White House executive order on AI seeks to address security risks
malware resembling NSA code
CISA budget cuts would be
Hackers that breached Las Vegas casinos rely on violent threats, research shows'''
data = [{'role': 'user', 'content': user_content}]
start_time = time.time()
2024-05-08 09:57:48 +08:00
pprint(openai_llm(data, "deepseek-chat"))
2024-04-07 09:37:47 +08:00
print(f'time cost: {time.time() - start_time}')