163 lines
5.1 KiB
Python
163 lines
5.1 KiB
Python
"""
|
||
LLM 故事生成服务 — 基于火山引擎豆包大模型
|
||
"""
|
||
import json
|
||
import logging
|
||
from django.conf import settings
|
||
|
||
try:
|
||
from volcenginesdkarkruntime import Ark
|
||
ARK_AVAILABLE = True
|
||
except ImportError:
|
||
ARK_AVAILABLE = False
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
STORY_SYSTEM_PROMPT = """# 角色
|
||
|
||
你是「卡皮巴拉故事工坊」的首席故事大师。你为 3-8 岁的小朋友创作原创童话故事。
|
||
|
||
# 任务
|
||
|
||
根据用户提供的**角色、场景、道具**素材，创作一个完整的儿童故事。
|
||
|
||
# 输出格式
|
||
|
||
你 **必须** 只返回如下 JSON，不要返回任何其他内容（不要 markdown 代码块，不要解释）：
|
||
|
||
{"title": "故事标题（6字以内）", "content": "故事正文"}
|
||
|
||
# 故事创作规范
|
||
|
||
1. **字数**：正文 400-600 字，不要太短也不要太长
|
||
2. **段落**：用 `\\n\\n` 分段，每段 2-4 句话
|
||
3. **语言**：简单易懂，适合给小朋友朗读；可以包含拟声词（"哗啦啦"、"咕噜噜"）和语气词（"哇！"、"嘿嘿"）
|
||
4. **结构**：开头引入角色和场景 → 中间遇到挑战或趣事 → 结尾温馨圆满
|
||
5. **情感**：温暖、有趣、充满想象力，带一点小幽默
|
||
6. **教育**：自然融入一个小道理（勇气、友谊、分享等），不要说教
|
||
7. **创意**：即使收到相同的素材组合，每次也要创作全新的、不同的故事情节
|
||
8. **角色融合**：所有用户选择的角色、场景、道具都必须在故事中出现并发挥作用
|
||
9. **标题**：简短有趣，6 个字以内，能引起小朋友的好奇心"""
|
||
|
||
|
||
def build_user_prompt(characters, scenes, props):
|
||
"""构建用户提示词"""
|
||
parts = []
|
||
if characters:
|
||
parts.append(f"角色：{', '.join(characters)}")
|
||
if scenes:
|
||
parts.append(f"场景：{', '.join(scenes)}")
|
||
if props:
|
||
parts.append(f"道具：{', '.join(props)}")
|
||
return '请根据以下元素创作一个儿童故事：\n' + '\n'.join(parts)
|
||
|
||
|
||
def generate_story_stream(characters, scenes, props):
|
||
"""
|
||
流式生成故事，yield SSE 事件字符串。
|
||
使用火山引擎豆包大模型（Ark SDK）。
|
||
|
||
Yields:
|
||
str: SSE 格式的事件数据行
|
||
"""
|
||
config = settings.LLM_CONFIG
|
||
|
||
if not config.get('API_KEY'):
|
||
yield _sse_event('error', {'message': 'Volcengine API Key 未配置'})
|
||
return
|
||
|
||
if not ARK_AVAILABLE:
|
||
yield _sse_event('error', {'message': 'volcengine SDK 未安装，请运行 pip install "volcengine-python-sdk[ark]"'})
|
||
return
|
||
|
||
yield _sse_event('stage', {
|
||
'stage': 'connecting',
|
||
'progress': 0,
|
||
'message': '正在收集灵感碎片...',
|
||
})
|
||
|
||
client = Ark(api_key=config['API_KEY'])
|
||
|
||
user_prompt = build_user_prompt(characters, scenes, props)
|
||
|
||
try:
|
||
yield _sse_event('stage', {
|
||
'stage': 'generating',
|
||
'progress': 10,
|
||
'message': '故事正在诞生...',
|
||
})
|
||
|
||
stream = client.chat.completions.create(
|
||
model=config['MODEL_NAME'],
|
||
messages=[
|
||
{'role': 'system', 'content': STORY_SYSTEM_PROMPT},
|
||
{'role': 'user', 'content': user_prompt},
|
||
],
|
||
max_tokens=2048,
|
||
stream=True,
|
||
)
|
||
|
||
full_content = ''
|
||
chunk_count = 0
|
||
|
||
for chunk in stream:
|
||
delta = chunk.choices[0].delta if chunk.choices else None
|
||
if delta and delta.content:
|
||
full_content += delta.content
|
||
chunk_count += 1
|
||
|
||
if chunk_count % 5 == 0:
|
||
progress = min(10 + int(chunk_count * 0.5), 80)
|
||
yield _sse_event('stage', {
|
||
'stage': 'generating',
|
||
'progress': progress,
|
||
'message': '故事正在诞生...',
|
||
})
|
||
|
||
yield _sse_event('stage', {
|
||
'stage': 'parsing',
|
||
'progress': 85,
|
||
'message': '正在编制最后的魔法...',
|
||
})
|
||
|
||
result = _parse_story_json(full_content)
|
||
|
||
yield _sse_event('done', {
|
||
'stage': 'done',
|
||
'progress': 100,
|
||
'message': '大功告成！',
|
||
'title': result['title'],
|
||
'content': result['content'],
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f'LLM story generation failed: {e}')
|
||
yield _sse_event('error', {'message': f'故事生成失败: {str(e)}'})
|
||
|
||
|
||
def _parse_story_json(text):
|
||
"""从 LLM 输出中解析故事 JSON"""
|
||
text = text.strip()
|
||
if text.startswith('```'):
|
||
text = text.split('\n', 1)[1] if '\n' in text else text[3:]
|
||
if text.endswith('```'):
|
||
text = text[:-3]
|
||
text = text.strip()
|
||
|
||
try:
|
||
data = json.loads(text)
|
||
return {
|
||
'title': data.get('title', '未命名故事'),
|
||
'content': data.get('content', text),
|
||
}
|
||
except json.JSONDecodeError:
|
||
return {
|
||
'title': '新故事',
|
||
'content': text,
|
||
}
|
||
|
||
|
||
def _sse_event(event, data):
|
||
"""格式化 SSE 事件"""
|
||
return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
|