128 lines
5.3 KiB
Python
128 lines
5.3 KiB
Python
import os
|
||
from openai import OpenAI
|
||
from dotenv import load_dotenv
|
||
import time
|
||
|
||
# === Config LLM call ===
|
||
load_dotenv()
|
||
deepseek_url = 'https://api.deepseek.com/v1' # set to be compatible with the OpenAI API
|
||
deepseek_api = os.getenv("DEEPSEEK_API_KEY")
|
||
if not deepseek_api:
|
||
raise ValueError("DEEPSEEK_API_KEY not set!")
|
||
|
||
def call_deepseek(
|
||
messages=None,
|
||
system_prompt="你是一个擅长前端开发的AI，专注于生成Vue.js代码。",
|
||
prompt=None,
|
||
model='deepseek-chat',
|
||
temperature=0.6,
|
||
max_tokens=1000,
|
||
stream=False,
|
||
max_retries=3,
|
||
):
|
||
"""
|
||
调用 DeepSeek API,支持多轮对话和流式/非流式响应
|
||
|
||
Args:
|
||
messages (list): 对话消息列表，格式为[{"role": "user", "content": "..."},...]。
|
||
如果提供，则优先使用，否则根据 system_prompt和 prompt 构造.
|
||
system_prompt (str): 系统提示词，仅在 message 未提供时使用.
|
||
prompt (str): 用户提示词，仅在 message 未提供时使用.
|
||
model (str): 模型名称，默认为 'deepseek-chat'。
|
||
temperature (float): 温度参数，控制生成文本的随机性，默认为 0.6。
|
||
max_tokens (int): 最大生成 token 数量，默认为 1000。
|
||
stream (bool): 是否使用流式响应，默认为 False。
|
||
max_retries (int): 最大重试次数，默认为 3。
|
||
|
||
Returns:
|
||
tuple: (result, usage)
|
||
- result (str): 非流式返回字符串(回复内容),流式返回生成器(逐块内容).
|
||
- usage (dict): token使用统计(非流式返回 Usage 对象,流式返回 None).
|
||
|
||
Raises:
|
||
ValueError: 如果 message 和 prompt 都为空或参数无效
|
||
Exception: 如果 API 调用失败且达到最大重试次数
|
||
"""
|
||
# 初始化 OpenAI 客户端
|
||
client = OpenAI(
|
||
api_key=deepseek_api,
|
||
base_url=deepseek_url,
|
||
)
|
||
|
||
# 参数验证
|
||
if messages is None:
|
||
if not prompt:
|
||
raise ValueError("必须提供 message 或 prompt 参数")
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": prompt}
|
||
]
|
||
elif not isinstance(messages, list) or not messages:
|
||
raise ValueError("message 参数必须是非空列表")
|
||
|
||
# 模型验证
|
||
models = ["deepseek-chat", "deepseek-reasoner"]
|
||
if model not in models:
|
||
raise ValueError(f"无效的模型名称: {model}，可用模型: {models}")
|
||
|
||
# 调用 API
|
||
for attempt in range(max_retries):
|
||
try:
|
||
response = client.chat.completions.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
stream=stream
|
||
)
|
||
|
||
# 流式响应
|
||
if stream:
|
||
def stream_generator():
|
||
usage = None
|
||
for chunk in response:
|
||
if chunk.choices[0].delta.content is not None:
|
||
yield chunk.choices[0].delta.content
|
||
return usage
|
||
return stream_generator(), None
|
||
else:
|
||
# 非流式响应
|
||
content = response.choices[0].message.content
|
||
return content, response.usage
|
||
|
||
except Exception as e:
|
||
if hasattr-------------('status_code') and e.status_code == 429: # 限流
|
||
print(f"请求过于频繁，正在重试... (尝试 {attempt + 1}/{max_retries})")
|
||
time.sleep(2 ** attempt) # 指数退避
|
||
else:
|
||
raise
|
||
raise Exception("达到最大重试次数，API 调用失败")
|
||
|
||
def generate_vue_code(prompt=None):
|
||
prompt = (
|
||
"生成一个Vue组件代码，用于端午节活动海报，包含以下部分并指定排版位置："
|
||
"1. 背景图层：div，占据整个组件区域。"
|
||
"2. 主体图层：div，位于顶部1/4处，居中，包含标题和副标题。"
|
||
"3. 活动亮点：div，位于底部1/4处，居中，使用网格布局展示三项活动（每项包含图标、标题和描述）。"
|
||
"4. 页脚：div，位于底部，居中，包含主办单位信息和logo图片。"
|
||
"组件尺寸为1080x1920px，布局使用absolute定位，仅关注排版位置，不包含任何样式描述（如颜色、字体、阴影、动画等）。"
|
||
"仅生成Vue代码本身，不包含说明性文字、注释或Markdown格式。"
|
||
)
|
||
system_prompt = (
|
||
"你是一个擅长前端开发的AI，专注于生成Vue.js代码。"
|
||
"生成的代码仅关注排版位置，使用absolute定位，不包含任何样式描述（如颜色、字体、阴影、动画等）。"
|
||
"确保代码符合Vue最佳实践，仅生成代码本身。"
|
||
)
|
||
|
||
result, usage = call_deepseek(prompt=prompt, system_prompt=system_prompt, temperature=0.4)
|
||
return result
|
||
|
||
def save_code(code, file_path="../outputs/generated_code.vue"):
|
||
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
||
with open(file_path, "w", encoding="utf-8") as f:
|
||
f.write(code)
|
||
|
||
if __name__ == "__main__":
|
||
vue_code = generate_vue_code()
|
||
save_code(vue_code)
|
||
print("Vue组件代码已生成并保存到outputs/generated_code.vue") |