konjacpotato
756f871b4e
All checks were successful
Gitea Actions Demo / deploy (push) Successful in 29s
144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
from typing import Optional, Dict
|
||
from dataclasses import dataclass
|
||
import os
|
||
from config import settings
|
||
from openai import OpenAI
|
||
from utils import logger
|
||
|
||
|
||
@dataclass
|
||
class LLMConfig:
|
||
"""LLM配置类"""
|
||
api_key: Optional[str] = None
|
||
base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||
model: str = settings.LLM_MODEL
|
||
enable_thinking: bool = True
|
||
temperature: float = 0.7
|
||
max_tokens: int = settings.LLM_MAX_TOKENS
|
||
|
||
|
||
class LLMThinkingEngine:
|
||
"""LLM驱动的思考引擎实现"""
|
||
|
||
def __init__(self, system_prompt_file: str = "system_prompt.txt", config: Optional[LLMConfig] = None):
|
||
"""
|
||
初始化LLMThinkingEngine
|
||
|
||
Args:
|
||
config: LLM配置对象，如果为None则使用默认配置
|
||
"""
|
||
self.system_prompt_file = system_prompt_file
|
||
self.config = config or LLMConfig()
|
||
self._init_client()
|
||
|
||
def _init_client(self):
|
||
"""初始化OpenAI客户端"""
|
||
api_key = self.config.api_key or settings.LLM_API_KEY
|
||
self.client = OpenAI(
|
||
api_key=api_key,
|
||
base_url=self.config.base_url,
|
||
)
|
||
|
||
def think(self, user_input: str) -> str:
|
||
"""
|
||
基于LLM进行思考，返回下一步的行动
|
||
|
||
Args:
|
||
user_input: 用户输入内容
|
||
|
||
Returns:
|
||
Thought: 思考结果，包含行动类型和内容
|
||
"""
|
||
# 构建适用于LLM的消息
|
||
messages = self._build_messages(user_input)
|
||
logger.info(f"LLM构建的消息: {messages}")
|
||
# 调用LLM进行思考
|
||
thinking_content, response_content = self._call_llm(messages)
|
||
|
||
# logger.info(f"LLM思考结果: thinking_content={thinking_content}, response_content={response_content}")
|
||
return response_content
|
||
|
||
def _build_messages(self, user_input: str) -> list[Dict[str, str]]:
|
||
"""
|
||
构建发送给LLM的消息
|
||
|
||
Args:
|
||
user_input: 用户输入内容
|
||
|
||
Returns:
|
||
消息列表，包含系统提示、历史和当前输入
|
||
"""
|
||
messages = []
|
||
|
||
# 系统提示
|
||
system_prompt = self._get_system_prompt()
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
|
||
# 用户输入
|
||
messages.append({
|
||
"role": "user",
|
||
"content": user_input
|
||
})
|
||
|
||
return messages
|
||
|
||
def _get_system_prompt(self) -> str:
|
||
"""
|
||
获取系统提示词
|
||
|
||
Returns:
|
||
系统提示词
|
||
"""
|
||
prompt_path = os.path.join(
|
||
os.path.dirname(__file__),
|
||
"prompts",
|
||
self.system_prompt_file
|
||
)
|
||
with open(prompt_path, "r", encoding="utf-8") as f:
|
||
return f.read()
|
||
|
||
def _call_llm(self, messages: list[Dict[str, str]]) -> tuple[str, str]:
|
||
"""
|
||
调用LLM API
|
||
|
||
Args:
|
||
messages: 消息列表
|
||
|
||
Returns:
|
||
(thinking_content, response_content): 思考内容和响应内容
|
||
"""
|
||
thinking_content = ""
|
||
response_content = ""
|
||
|
||
try:
|
||
completion = self.client.chat.completions.create(
|
||
model=self.config.model,
|
||
messages=messages,
|
||
temperature=self.config.temperature,
|
||
max_tokens=self.config.max_tokens,
|
||
extra_body={"enable_thinking": self.config.enable_thinking},
|
||
stream=True
|
||
)
|
||
|
||
# 流式处理响应
|
||
for chunk in completion:
|
||
delta = chunk.choices[0].delta
|
||
|
||
# 收集思考内容
|
||
if hasattr(delta, "reasoning_content") and delta.reasoning_content:
|
||
thinking_content += delta.reasoning_content
|
||
|
||
# 收集响应内容
|
||
if hasattr(delta, "content") and delta.content:
|
||
response_content += delta.content
|
||
|
||
except Exception as e:
|
||
# 错误处理
|
||
response_content = f"调用LLM时出错：{str(e)}"
|
||
|
||
return thinking_content, response_content
|
||
|
||
def set_config(self, config: LLMConfig):
|
||
"""更新LLM配置"""
|
||
self.config = config
|
||
self._init_client() |