Files
my-v0-server/app/processor/nodes/qa_node.py
T
2026-09-29 19:21:29 +08:00

93 lines
3.3 KiB
Python

"""基于当前大屏和澄清历史回答问题"""
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.output_parsers import StrOutputParser
from app.clients.llm import get_llm
from app.processor.base import BaseNode
from app.processor.state import QueryState
from app.prompt.qa import QA_SYSTEM_PROMPT
class QANode(BaseNode):
name = "qa_node"
async def process(self, state: QueryState) -> QueryState:
# 浅拷贝状态
updated_state: QueryState = state.copy()
# 获取普通对话模型,问答只需文本回答,不绑定结构化输出协议
llm_client = get_llm()
# 异步调用模型,结合用户问题、当前大屏和澄清记录生成回答
llm_res = await llm_client.ainvoke(
[
# 系统消息提供稳定的问答规则和回答边界
SystemMessage(content=QA_SYSTEM_PROMPT),
# 将输入上下文序列化为 JSON 放入用户消息,只保留显式设置的字段
HumanMessage(
content=state["input_context"].model_dump_json(exclude_unset=True)
),
]
)
# 从模型响应中提取文本,兼容文本内容块,并去除回答首尾的空白
answer = StrOutputParser().invoke(llm_res).strip()
# 如果没有有效文本,抛出异常,避免将空回答标记为处理完成
if not answer:
raise ValueError("问答模型未返回文本回答")
# 保存最终回答,供后续输出给用户
updated_state["answer"] = answer
# 回答生成完成,问答分支在此结束
updated_state["status"] = "completed"
# 返回包含回答和完成状态的更新结果
return updated_state
if __name__ == "__main__":
import asyncio
import json
from pydantic import BaseModel
from app.processor.nodes.input_node import InputNode
async def main():
# 直接测试问答,仅需 LLM 配置
state: QueryState = {
"request": {
"prompt": "销售额统计的展示可以怎么优化?先给建议,不要修改",
"screen": {
"canvas": {"width": 1920, "height": 1080},
"nodes": [
{
"id": "sales-chart",
"type": "bar-chart",
"name": "销售额统计",
"layout": {"x": 26, "y": 34.8, "width": 400, "height": 260},
"props": {"option": {"title": {"text": "销售额统计"}}},
}
],
},
"context": {"selected_node_ids": ["sales-chart"]},
}
}
state = InputNode().process(state)
node = QANode()
final_state = await node.process(state)
print(
json.dumps(
{
key: value.model_dump(mode="json")
if isinstance(value, BaseModel)
else value
for key, value in final_state.items()
},
ensure_ascii=False,
indent=2,
)
)
asyncio.run(main())