aiData/Util/LlmUtil.py

import sys
import asyncio

from openai import AsyncOpenAI

from Config import *
from Config.Config import ALY_LLM_API_KEY, ALY_LLM_BASE_URL, ALY_LLM_MODEL_NAME


# 保留原同步函数，添加异步版本
async def get_llm_response(query_text: str, stream: bool = True, system_prompt: str = 'You are a helpful assistant.', chat_history: list = None, temperature: float = None):
    """
    异步获取大模型的响应
    @param query_text: 查询文本
    @param stream: 是否使用流式输出
    @param system_prompt: 系统提示文本，默认为'You are a helpful assistant.'
    @param chat_history: 聊天历史记录，格式为[{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}, ...]
    @param temperature: 温度参数，控制输出的随机性，0-2之间，默认为None使用模型默认值
    @return: 流式响应生成器或完整响应文本
    """
    client = AsyncOpenAI(
        api_key=ALY_LLM_API_KEY,
        base_url=ALY_LLM_BASE_URL,
    )
    
    try:
        # 构建messages列表，包含系统提示和聊天历史
        messages = [{'role': 'system', 'content': system_prompt}]
        
        # 如果有聊天历史，添加到messages中
        if chat_history and isinstance(chat_history, list):
            messages.extend(chat_history)
        
        # 添加当前用户查询
        messages.append({'role': 'user', 'content': query_text})
        
        # 构建请求参数
        request_params = {
            'model': ALY_LLM_MODEL_NAME,
            'messages': messages,
            'stream': stream
        }
        
        # 如果指定了temperature参数，添加到请求中
        if temperature is not None:
            request_params['temperature'] = temperature
        
        # 创建请求
        completion = await asyncio.wait_for(client.chat.completions.create(**request_params), timeout=60.0)
        
        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:
                        # 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:
                    # if content is not None and content.strip():
                        yield content
    except Exception as e:
        error_msg = f"大模型请求异常: {str(e)}"
        print(error_msg, file=sys.stderr)
        # 抛出异常而不是返回错误信息，让调用方处理
        raise Exception(error_msg)