186 lines
7.2 KiB
Python
186 lines
7.2 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
AI 代理模块
|
|||
|
|
负责与大模型进行交互，生成文本内容
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
import logging
|
|||
|
|
import traceback
|
|||
|
|
from openai import AsyncOpenAI, APITimeoutError, APIConnectionError, RateLimitError, APIStatusError
|
|||
|
|
import tiktoken
|
|||
|
|
from typing import Optional, Tuple
|
|||
|
|
|
|||
|
|
from core.config import AIModelConfig
|
|||
|
|
from core.exception import AIModelError, RetryableError, NonRetryableError
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AIAgent:
|
|||
|
|
"""
|
|||
|
|
AI代理类，负责与AI模型交互生成文本内容
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, config: AIModelConfig):
|
|||
|
|
"""
|
|||
|
|
初始化 AI 代理
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
config: AI模型配置
|
|||
|
|
"""
|
|||
|
|
self.config = config
|
|||
|
|
self.client = AsyncOpenAI(
|
|||
|
|
api_key=self.config.api_key,
|
|||
|
|
base_url=self.config.api_url,
|
|||
|
|
timeout=self.config.timeout
|
|||
|
|
)
|
|||
|
|
# try:
|
|||
|
|
# self.tokenizer = tiktoken.encoding_for_model(self.config.model)
|
|||
|
|
# except KeyError:
|
|||
|
|
# logger.warning(f"模型 '{self.config.model}' 没有找到对应的tokenizer，将使用 'cl100k_base'")
|
|||
|
|
# self.tokenizer = tiktoken.get_encoding("cl100k_base")
|
|||
|
|
|
|||
|
|
async def generate_text(
|
|||
|
|
self, system_prompt: str, user_prompt: str, use_stream: bool = False,
|
|||
|
|
temperature: Optional[float] = None, top_p: Optional[float] = None,
|
|||
|
|
presence_penalty: Optional[float] = None, stage: str = ""
|
|||
|
|
) -> Tuple[str, int, int, float]:
|
|||
|
|
"""
|
|||
|
|
生成文本 (支持流式和非流式)
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
system_prompt: 系统提示
|
|||
|
|
user_prompt: 用户提示
|
|||
|
|
use_stream: 是否流式返回
|
|||
|
|
temperature: 温度参数，控制随机性
|
|||
|
|
top_p: Top-p采样参数
|
|||
|
|
presence_penalty: 存在惩罚参数
|
|||
|
|
stage: 当前所处阶段，用于日志记录
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
一个元组 (generated_text, input_tokens, output_tokens, time_cost)
|
|||
|
|
"""
|
|||
|
|
messages = [
|
|||
|
|
{"role": "system", "content": system_prompt},
|
|||
|
|
{"role": "user", "content": user_prompt}
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# 使用传入的参数或默认配置
|
|||
|
|
temp = temperature if temperature is not None else self.config.temperature
|
|||
|
|
tp = top_p if top_p is not None else self.config.top_p
|
|||
|
|
pp = presence_penalty if presence_penalty is not None else self.config.presence_penalty
|
|||
|
|
|
|||
|
|
# 记录使用的模型参数
|
|||
|
|
stage_info = f"[{stage}]" if stage else ""
|
|||
|
|
logger.info(f"{stage_info} 使用模型参数: temperature={temp:.2f}, top_p={tp:.2f}, presence_penalty={pp:.2f}")
|
|||
|
|
|
|||
|
|
input_tokens = self.count_tokens(system_prompt + user_prompt)
|
|||
|
|
logger.info(f"{stage_info} 开始生成任务... 输入token数: {input_tokens}")
|
|||
|
|
|
|||
|
|
last_exception = None
|
|||
|
|
backoff_time = 1.0 # Start with 1 second
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
for attempt in range(self.config.max_retries):
|
|||
|
|
try:
|
|||
|
|
response = await self.client.chat.completions.create(
|
|||
|
|
model=self.config.model,
|
|||
|
|
messages=messages,
|
|||
|
|
temperature=temp,
|
|||
|
|
top_p=tp,
|
|||
|
|
presence_penalty=pp,
|
|||
|
|
stream=use_stream
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
time_cost = time.time() - start_time
|
|||
|
|
|
|||
|
|
if use_stream:
|
|||
|
|
# 流式处理需要异步迭代
|
|||
|
|
full_text = await self._process_stream(response)
|
|||
|
|
output_tokens = self.count_tokens(full_text)
|
|||
|
|
logger.info(f"{stage_info} 任务完成，耗时 {time_cost:.2f} 秒. 输出token数: {output_tokens}")
|
|||
|
|
return full_text, input_tokens, output_tokens, time_cost
|
|||
|
|
else:
|
|||
|
|
output_text = response.choices[0].message.content.strip()
|
|||
|
|
output_tokens = self.count_tokens(output_text)
|
|||
|
|
logger.info(f"{stage_info} 任务完成，耗时 {time_cost:.2f} 秒. 输出token数: {output_tokens}")
|
|||
|
|
return output_text, input_tokens, output_tokens, time_cost
|
|||
|
|
|
|||
|
|
except (APITimeoutError, APIConnectionError) as e:
|
|||
|
|
last_exception = RetryableError(f"AI模型连接或超时错误: {e}")
|
|||
|
|
logger.warning(f"{stage_info} 尝试 {attempt + 1}/{self.config.max_retries} 失败: {last_exception}. "
|
|||
|
|
f"将在 {backoff_time:.1f} 秒后重试...")
|
|||
|
|
time.sleep(backoff_time)
|
|||
|
|
backoff_time *= 2 # Exponential backoff
|
|||
|
|
except (RateLimitError, APIStatusError) as e:
|
|||
|
|
last_exception = NonRetryableError(f"AI模型API错误 (不可重试): {e}")
|
|||
|
|
logger.error(f"{stage_info} 发生不可重试的API错误: {last_exception}")
|
|||
|
|
break # Do not retry on these errors
|
|||
|
|
except Exception as e:
|
|||
|
|
last_exception = AIModelError(f"调用AI模型时发生未知错误: {e}")
|
|||
|
|
logger.error(f"{stage_info} 发生未知错误: {last_exception}\n{traceback.format_exc()}")
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
raise AIModelError(f"AI模型调用在 {self.config.max_retries} 次重试后失败") from last_exception
|
|||
|
|
|
|||
|
|
async def _process_stream(self, response):
|
|||
|
|
"""异步处理流式响应"""
|
|||
|
|
full_response = []
|
|||
|
|
async for chunk in response:
|
|||
|
|
content = chunk.choices[0].delta.content
|
|||
|
|
if content:
|
|||
|
|
full_response.append(content)
|
|||
|
|
# 如果需要在这里实现真正的流式处理，可以使用回调函数或其他方式
|
|||
|
|
|
|||
|
|
full_text = "".join(full_response)
|
|||
|
|
logger.info(f"流式响应接收完成，总长度: {len(full_text)}")
|
|||
|
|
return full_text
|
|||
|
|
|
|||
|
|
def count_tokens(self, text: str) -> int:
|
|||
|
|
"""
|
|||
|
|
计算文本的token数量
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
text: 输入文本
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
token数量
|
|||
|
|
"""
|
|||
|
|
return len(text) // 1.5
|
|||
|
|
# return len(self.tokenizer.encode(text))
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def read_folder_content(folder_path: str) -> str:
|
|||
|
|
"""
|
|||
|
|
读取指定文件夹下的所有文件内容并合并
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
folder_path: 文件夹路径
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
合并后的文件内容
|
|||
|
|
"""
|
|||
|
|
if not os.path.exists(folder_path):
|
|||
|
|
logger.warning(f"引用的文件夹不存在: {folder_path}")
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
context = ""
|
|||
|
|
try:
|
|||
|
|
for file_name in sorted(os.listdir(folder_path)):
|
|||
|
|
file_path = os.path.join(folder_path, file_name)
|
|||
|
|
if os.path.isfile(file_path):
|
|||
|
|
try:
|
|||
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|||
|
|
context += f"--- 文件名: {file_name} ---\n"
|
|||
|
|
context += f.read().strip()
|
|||
|
|
context += "\n\n"
|
|||
|
|
except Exception as read_err:
|
|||
|
|
logger.error(f"读取文件失败 {file_path}: {read_err}")
|
|||
|
|
except Exception as list_err:
|
|||
|
|
logger.error(f"列出目录失败 {folder_path}: {list_err}")
|
|||
|
|
return context.strip()
|