94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
from openai import OpenAI
|
|
from Config.Config import *
|
|
import sys
|
|
|
|
from openai import AsyncOpenAI # 新增异步客户端导入
|
|
|
|
# 保留原同步函数，添加异步版本
|
|
async def get_llm_response_async(query_text: str, stream: bool = True):
|
|
"""
|
|
异步获取大模型的响应
|
|
@param query_text: 查询文本
|
|
@param stream: 是否使用流式输出
|
|
@return: 流式响应生成器或完整响应文本
|
|
"""
|
|
client = AsyncOpenAI(
|
|
api_key=LLM_API_KEY,
|
|
base_url=LLM_BASE_URL,
|
|
)
|
|
|
|
try:
|
|
# 创建请求
|
|
completion = await client.chat.completions.create(
|
|
model=LLM_MODEL_NAME,
|
|
messages=[
|
|
{'role': 'system', 'content': 'You are a helpful assistant.'},
|
|
{'role': 'user', 'content': query_text}
|
|
],
|
|
stream=stream
|
|
)
|
|
|
|
if stream:
|
|
# 流式输出模式，返回生成器
|
|
async for chunk in completion:
|
|
# 确保 chunk.choices 存在且不为空
|
|
if chunk and chunk.choices and len(chunk.choices) > 0:
|
|
# 确保 delta 存在
|
|
delta = chunk.choices[0].delta
|
|
if delta:
|
|
# 确保 content 存在且不为 None 或空字符串
|
|
content = delta.content
|
|
if content is not None and content.strip():
|
|
print(content, end='', flush=True)
|
|
yield content
|
|
else:
|
|
# 非流式处理
|
|
if completion and completion.choices and len(completion.choices) > 0:
|
|
message = completion.choices[0].message
|
|
if message:
|
|
content = message.content
|
|
if content is not None and content.strip():
|
|
yield content
|
|
except Exception as e:
|
|
print(f"大模型请求异常: {str(e)}", file=sys.stderr)
|
|
yield f"处理请求时发生异常: {str(e)}"
|
|
|
|
# 保留原同步函数
|
|
def get_llm_response(query_text: str, stream: bool = True):
|
|
"""
|
|
获取大模型的响应
|
|
@param query_text: 查询文本
|
|
@param stream: 是否使用流式输出
|
|
@return: 完整响应文本
|
|
"""
|
|
client = OpenAI(
|
|
api_key=LLM_API_KEY,
|
|
base_url=LLM_BASE_URL,
|
|
)
|
|
|
|
# 创建请求
|
|
completion = client.chat.completions.create(
|
|
model=LLM_MODEL_NAME,
|
|
messages=[
|
|
{'role': 'system', 'content': 'You are a helpful assistant.'},
|
|
{'role': 'user', 'content': query_text}
|
|
],
|
|
stream=stream
|
|
)
|
|
|
|
full_response = []
|
|
|
|
if stream:
|
|
for chunk in completion:
|
|
# 提取当前块的内容
|
|
if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
|
|
content = chunk.choices[0].delta.content
|
|
full_response.append(content)
|
|
# 实时输出内容，不换行
|
|
print(content, end='', flush=True)
|
|
else:
|
|
# 非流式处理
|
|
full_response.append(completion.choices[0].message.content)
|
|
|
|
return ''.join(full_response)
|