160 lines
6.7 KiB
Python
160 lines
6.7 KiB
Python
"""基于当前大屏、上传的数据和会话历史流式回答问题
|
||
|
||
由问答 agent 执行：会话里有数据文件时先查询再回答，否则只调用一次模型；只把模型输出的文字推送给前端
|
||
"""
|
||
|
||
import json
|
||
|
||
from langchain.agents.middleware.model_call_limit import ModelCallLimitExceededError
|
||
from langchain_core.messages import (
|
||
AIMessage,
|
||
AIMessageChunk,
|
||
BaseMessage,
|
||
HumanMessage,
|
||
ToolMessage,
|
||
)
|
||
from langchain_core.output_parsers import StrOutputParser
|
||
from langgraph.config import get_stream_writer
|
||
|
||
from app.agent.qa_agent import QAContext, get_qa_agent
|
||
from app.processor.base import BaseNode
|
||
from app.processor.history import history_messages, history_records
|
||
from app.processor.state import QueryState, run_data
|
||
from app.processor.summary import data_file_summary
|
||
from app.schemas.base import JsonObject
|
||
from app.schemas.conversation import TurnQuery
|
||
|
||
# 模型在一次回答里分几次输出文字（例如查询前后各说一段）时，段落之间空一行
|
||
PARAGRAPH_BREAK = "\n\n"
|
||
# 带数据的问答回答完后给的后续请求：意图判断会把"展示一下 9 月销售额"这类说法分到问答，
|
||
# 用户其实想在大屏上看到；回答之后给一个按钮，点了按普通请求走大屏 agent，会话历史里有这次回答
|
||
PUT_ON_SCREEN = "把上面的结果放到大屏上"
|
||
# 一轮最多记下的查询；Langfuse 里每条 trace 的 query_data 调用 p50 3 次、p90 7 次（含重试）
|
||
MAX_QUERY_RECORDS = 10
|
||
|
||
|
||
class QANode(BaseNode):
|
||
name = "qa_node"
|
||
|
||
async def process(self, state: QueryState) -> QueryState:
|
||
updated_state: QueryState = state.copy()
|
||
|
||
try:
|
||
# 图执行上下文中通过 custom 流逐段推送回答，独立调用节点时没有写入器
|
||
writer = get_stream_writer()
|
||
except RuntimeError:
|
||
writer = None
|
||
|
||
context = state["input_context"]
|
||
payload = context.model_dump(
|
||
mode="json", exclude_unset=True, exclude={"history", "data_files"}
|
||
)
|
||
if context.data_files:
|
||
# 只给表名、行数和列名，文件路径和列概况不进提示词，需要时由 describe_data 查看
|
||
payload["data_files"] = [
|
||
data_file_summary(item) for item in context.data_files
|
||
]
|
||
agent = get_qa_agent(with_data=bool(context.data_files))
|
||
|
||
# 此前各轮的修改、查询和澄清记录跟着本次请求，不拼进历史消息；前端带来的同名字段不保留
|
||
payload.pop("history_records", None)
|
||
if records := history_records(context.history, context.screen):
|
||
payload["history_records"] = records
|
||
|
||
parser = StrOutputParser()
|
||
parts: list[str] = []
|
||
last_message_id = None
|
||
messages: list[BaseMessage] = []
|
||
try:
|
||
async for mode, data in agent.astream(
|
||
{
|
||
"messages": [
|
||
# 会话历史作为此前的对话消息，本次请求是最后一条用户消息
|
||
*history_messages(context.history),
|
||
HumanMessage(content=json.dumps(payload, ensure_ascii=False)),
|
||
]
|
||
},
|
||
context=QAContext(data=run_data()),
|
||
# messages 推送文字，values 留下完整的消息，用来记下执行过的查询
|
||
stream_mode=["messages", "values"],
|
||
):
|
||
if mode == "values":
|
||
messages = data["messages"]
|
||
continue
|
||
chunk, metadata = data
|
||
# 工具结果、工具调用和中间件插入的消息都不推送
|
||
if metadata.get("langgraph_node") != "model" or not isinstance(
|
||
chunk, AIMessageChunk
|
||
):
|
||
continue
|
||
if not (delta := parser.invoke(chunk)):
|
||
continue
|
||
if parts and chunk.id != last_message_id:
|
||
delta = PARAGRAPH_BREAK + delta
|
||
last_message_id = chunk.id
|
||
parts.append(delta)
|
||
if writer is not None:
|
||
writer(
|
||
{
|
||
"type": "answer_delta",
|
||
"message_id": state.get("message_id"),
|
||
"delta": delta,
|
||
}
|
||
)
|
||
except ModelCallLimitExceededError:
|
||
self.logger.warning("问答 agent 超出模型调用上限（MAX_QA_STEPS）")
|
||
updated_state["status"] = "failed"
|
||
updated_state["answer"] = (
|
||
"查询步骤超出上限，这次没有完成，可以缩小范围或换个问法再试"
|
||
)
|
||
return updated_state
|
||
|
||
# 推送了什么就保存什么，会话历史和前端看到的回答一致
|
||
answer = "".join(parts).strip()
|
||
# 没有有效文本时报错，避免将空回答标记为处理完成
|
||
if not answer:
|
||
raise ValueError("问答模型未返回文本回答")
|
||
|
||
updated_state["answer"] = answer
|
||
updated_state["status"] = "completed"
|
||
if context.data_files:
|
||
updated_state["suggestions"] = [PUT_ON_SCREEN]
|
||
# 只记成功的问答：步数用尽时回答里也没有数字可用
|
||
updated_state["queries"] = query_records(messages)
|
||
return updated_state
|
||
|
||
|
||
def query_records(messages: list[BaseMessage]) -> list[JsonObject]:
|
||
"""问答里成功执行的 query_data：调用参数和工具结果按 tool_call_id 配对
|
||
|
||
只存 SQL 和结果的形状；按执行顺序，相同的 SQL 只留最后一次，最多 MAX_QUERY_RECORDS 条（留最近的）
|
||
"""
|
||
sql_of = {
|
||
call["id"]: call["args"].get("sql")
|
||
for message in messages
|
||
if isinstance(message, AIMessage)
|
||
for call in message.tool_calls
|
||
if call["name"] == "query_data"
|
||
}
|
||
found: dict[str, JsonObject] = {}
|
||
for message in messages:
|
||
if not isinstance(message, ToolMessage) or message.status == "error":
|
||
continue
|
||
sql = sql_of.get(message.tool_call_id)
|
||
if not isinstance(sql, str):
|
||
continue
|
||
try:
|
||
result = json.loads(str(message.content))
|
||
record = TurnQuery(
|
||
sql=sql,
|
||
columns=result["columns"],
|
||
rows=result["row_count"],
|
||
truncated=result["truncated"],
|
||
).model_dump(mode="json")
|
||
except (ValueError, KeyError, TypeError):
|
||
continue
|
||
# 重新插入，相同 SQL 的位置按最后一次执行算；按原文比较，字面量里的空白也有意义
|
||
found.pop(sql, None)
|
||
found[sql] = record
|
||
return list(found.values())[-MAX_QUERY_RECORDS:]
|