107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
|
import os
|
|||
|
|
from openai import OpenAI
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AI_Agent():
|
|||
|
|
"""AI代理类，负责与AI模型交互生成文本内容"""
|
|||
|
|
|
|||
|
|
def __init__(self, base_url, model_name, api):
|
|||
|
|
self.url_list = {
|
|||
|
|
"ali": "https://dashscope.aliyuncs.com/compatible-mode/v1",
|
|||
|
|
"kimi": "https://api.moonshot.cn/v1",
|
|||
|
|
"doubao": "https://ark.cn-beijing.volces.com/api/v3/",
|
|||
|
|
"deepseek": "https://api.deepseek.com",
|
|||
|
|
"vllm": "http://localhost:8000/v1",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
self.base_url = self.url_list[base_url] if base_url in self.url_list else base_url
|
|||
|
|
self.api = api
|
|||
|
|
self.model_name = model_name
|
|||
|
|
|
|||
|
|
self.client = OpenAI(
|
|||
|
|
api_key=self.api,
|
|||
|
|
base_url=self.base_url,
|
|||
|
|
# timeout=10
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def generate_text(self, system_prompt, user_prompt, temperature, top_p, presence_penalty):
|
|||
|
|
"""生成文本内容，并返回完整响应和token估计值"""
|
|||
|
|
print("系统提示词:")
|
|||
|
|
print(system_prompt)
|
|||
|
|
print("\n用户提示词:")
|
|||
|
|
print(user_prompt)
|
|||
|
|
print(f"\nAPI Key: {self.api}")
|
|||
|
|
print(f"Base URL: {self.base_url}")
|
|||
|
|
print(f"Model: {self.model_name}")
|
|||
|
|
|
|||
|
|
response = self.client.chat.completions.create(
|
|||
|
|
model=self.model_name,
|
|||
|
|
messages=[{"role": "system", "content": system_prompt},
|
|||
|
|
{"role": "user", "content": user_prompt}],
|
|||
|
|
temperature=temperature,
|
|||
|
|
top_p=top_p,
|
|||
|
|
# presence_penalty=presence_penalty,
|
|||
|
|
stream=True,
|
|||
|
|
max_tokens=8192,
|
|||
|
|
extra_body={
|
|||
|
|
"repetition_penalty": 1.05,
|
|||
|
|
},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 收集完整的输出内容
|
|||
|
|
full_response = ""
|
|||
|
|
for chunk in response:
|
|||
|
|
if chunk.choices and len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None:
|
|||
|
|
content = chunk.choices[0].delta.content
|
|||
|
|
full_response += content
|
|||
|
|
print(content, end="", flush=True)
|
|||
|
|
if chunk.choices and len(chunk.choices) > 0 and chunk.choices[0].finish_reason == "stop":
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
print("\n完成生成，正在处理结果...")
|
|||
|
|
|
|||
|
|
# 由于使用流式输出，无法获取真实的token计数，因此返回估计值
|
|||
|
|
estimated_tokens = len(full_response.split()) * 1.3 # 简单估算token数量
|
|||
|
|
|
|||
|
|
return full_response, estimated_tokens
|
|||
|
|
|
|||
|
|
def read_folder(self, file_folder):
|
|||
|
|
"""读取指定文件夹下的所有文件内容"""
|
|||
|
|
if not os.path.exists(file_folder):
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
context = ""
|
|||
|
|
for file in os.listdir(file_folder):
|
|||
|
|
file_path = os.path.join(file_folder, file)
|
|||
|
|
if os.path.isfile(file_path):
|
|||
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|||
|
|
context += f"文件名: {file}\n"
|
|||
|
|
for line in f.readlines():
|
|||
|
|
context += line
|
|||
|
|
context += "\n"
|
|||
|
|
return context
|
|||
|
|
|
|||
|
|
def work(self, system_prompt, user_prompt, file_folder, temperature, top_p, presence_penalty):
|
|||
|
|
"""完整的工作流程：生成文本并返回结果"""
|
|||
|
|
# 生成时间戳
|
|||
|
|
date_time = time.strftime("%Y-%m-%d_%H-%M-%S")
|
|||
|
|
result_file = f"/root/autodl-tmp/xhsTweetGene/result/{date_time}.md"
|
|||
|
|
|
|||
|
|
# 如果提供了参考文件夹，则读取其内容
|
|||
|
|
if file_folder:
|
|||
|
|
context = self.read_folder(file_folder)
|
|||
|
|
if context:
|
|||
|
|
user_prompt = f"{user_prompt}\n参考资料:\n{context}"
|
|||
|
|
|
|||
|
|
# 计时
|
|||
|
|
time_start = time.time()
|
|||
|
|
|
|||
|
|
# 生成文本
|
|||
|
|
result, tokens = self.generate_text(system_prompt, user_prompt, temperature, top_p, presence_penalty)
|
|||
|
|
|
|||
|
|
# 计算耗时
|
|||
|
|
time_end = time.time()
|
|||
|
|
time_cost = time_end - time_start
|
|||
|
|
|
|||
|
|
return result, system_prompt, user_prompt, file_folder, result_file, tokens, time_cost
|