77 lines
2.5 KiB
Python
77 lines
2.5 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:
|
|
llm_client = get_llm()
|
|
llm_res = await llm_client.ainvoke(
|
|
[
|
|
SystemMessage(content=QA_SYSTEM_PROMPT),
|
|
HumanMessage(
|
|
content=state["input_context"].model_dump_json(exclude_unset=True)
|
|
),
|
|
]
|
|
)
|
|
# 普通文本输出，同时兼容模型返回的文本内容块
|
|
answer = StrOutputParser().invoke(llm_res).strip()
|
|
if not answer:
|
|
raise ValueError("问答模型未返回文本回答")
|
|
return {**state, "answer": answer}
|
|
|
|
|
|
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())
|