250 lines
7.5 KiB
Python
250 lines
7.5 KiB
Python
"""
|
||
主脑Agent - 简化版（无意图分类）
|
||
|
||
职责：
|
||
1. 输入处理：图片、音频、文本提取
|
||
2. 直接路由：根据外部传入的 route_flag 决定调用哪个工作流
|
||
|
||
工作流列表（来自 workflow_registry）：
|
||
- 故障排查与修理
|
||
- 舰船百科
|
||
- 操作使用
|
||
"""
|
||
|
||
from typing import TypedDict, Literal, Optional, Dict, Any, List
|
||
from langgraph.graph import StateGraph, START, END
|
||
from tools.function_tool import (
|
||
vlm_image_to_text,
|
||
asr_audio_to_text,
|
||
detect_input_type
|
||
)
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
from utils.function_tracker import track_function_calls
|
||
from workflow_registry import VALID_ROUTE_FLAGS
|
||
import json
|
||
import re
|
||
|
||
|
||
class MainAgentState(TypedDict):
|
||
"""主脑Agent状态"""
|
||
raw_input: Dict[str, Any]
|
||
has_image: bool
|
||
has_audio: bool
|
||
has_query: bool
|
||
has_file_text: bool
|
||
input_type: str
|
||
|
||
history_message: str
|
||
extracted_text: str
|
||
image_description: str
|
||
asr_text: str
|
||
file_text: str
|
||
combined_query: str
|
||
|
||
workflow_result: Optional[Dict[str, Any]]
|
||
final_response: str
|
||
error_message: str
|
||
|
||
route_flag: str
|
||
route_params: Dict[str, Any]
|
||
|
||
|
||
def detect_input_type_node(state: MainAgentState) -> Dict[str, Any]:
|
||
"""检测输入类型"""
|
||
raw = state["raw_input"]
|
||
print(raw)
|
||
print(22222222222222222222222222)
|
||
result = detect_input_type.invoke({"raw_input": raw})
|
||
|
||
has_image = result.get("has_image", False)
|
||
has_audio = result.get("has_audio", False)
|
||
has_query = result.get("has_query", False)
|
||
|
||
has_file_text = False
|
||
file_text_v = raw.get("file_text")
|
||
if isinstance(file_text_v, str) and file_text_v.strip():
|
||
has_file_text = True
|
||
|
||
input_type = "mixed"
|
||
count = sum([has_image, has_audio, has_query, has_file_text])
|
||
if count == 1:
|
||
if has_image:
|
||
input_type = "image"
|
||
elif has_audio:
|
||
input_type = "audio"
|
||
elif has_query:
|
||
input_type = "text"
|
||
elif has_file_text:
|
||
input_type = "file"
|
||
|
||
return {
|
||
"has_image": has_image,
|
||
"has_audio": has_audio,
|
||
"has_query": has_query,
|
||
"has_file_text": has_file_text,
|
||
"input_type": input_type,
|
||
}
|
||
|
||
|
||
async def process_multimodal_input(state: MainAgentState) -> Dict[str, Any]:
|
||
"""处理多模态输入"""
|
||
has_image = state.get("has_image", False)
|
||
has_audio = state.get("has_audio", False)
|
||
has_query = state.get("has_query", False)
|
||
has_file_text = state.get("has_file_text", False)
|
||
raw = state["raw_input"]
|
||
|
||
image_description = ""
|
||
asr_text = ""
|
||
extracted_text = ""
|
||
|
||
try:
|
||
if has_image:
|
||
image_path = raw.get("image", "")
|
||
if image_path:
|
||
try:
|
||
image_description = await vlm_image_to_text.ainvoke({
|
||
"image_path": image_path
|
||
})
|
||
if not image_description or image_description.startswith("VLM 调用失败"):
|
||
image_description = ""
|
||
except Exception:
|
||
image_description = ""
|
||
|
||
if has_audio:
|
||
audio_path = raw.get("audio", "")
|
||
if audio_path:
|
||
try:
|
||
asr_text = await asr_audio_to_text.ainvoke({
|
||
"audio_path": audio_path
|
||
})
|
||
if not asr_text or asr_text.startswith("ASR 调用失败"):
|
||
asr_text = ""
|
||
except Exception:
|
||
asr_text = ""
|
||
|
||
text_parts = []
|
||
if raw.get("query"):
|
||
text_parts.append(raw["query"])
|
||
if image_description:
|
||
text_parts.append(f"[图片描述] {image_description}")
|
||
if asr_text:
|
||
text_parts.append(f"[语音转文字] {asr_text}")
|
||
if raw.get("file_text"):
|
||
text_parts.append(f"[文件内容] {raw['file_text']}")
|
||
|
||
extracted_text = "\n".join(text_parts) if text_parts else ""
|
||
combined_query = raw.get("query", "") or extracted_text
|
||
|
||
return {
|
||
"extracted_text": extracted_text,
|
||
"image_description": image_description,
|
||
"asr_text": asr_text,
|
||
"combined_query": combined_query,
|
||
"file_text": raw.get("file_text", ""),
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
"error_message": f"输入处理失败: {str(e)}",
|
||
"extracted_text": "",
|
||
}
|
||
|
||
# @track_function_calls
|
||
def route_to_workflow(state: MainAgentState) -> Dict[str, Any]:
|
||
"""
|
||
🤖分析意图，调用相关智能体
|
||
"""
|
||
route_flag = state.get("route_flag", "问题排查")
|
||
|
||
if route_flag not in VALID_ROUTE_FLAGS:
|
||
route_flag = "问题排查"
|
||
|
||
query_text = state.get("extracted_text", "")
|
||
combined_query = state.get("combined_query", "") or query_text
|
||
|
||
raw_input = state.get("raw_input", {})
|
||
history_message = raw_input.get("history_message", "")
|
||
|
||
route_params = {
|
||
"extracted_text": query_text,
|
||
"combined_query": combined_query,
|
||
"history_message": history_message,
|
||
}
|
||
|
||
if raw_input.get("file_text"):
|
||
route_params["file_text"] = raw_input["file_text"]
|
||
|
||
print(f"[主脑Agent] 路由到: {route_flag}")
|
||
print(f"[主脑Agent] 参数: combined_query={combined_query[:100]}...")
|
||
|
||
if route_flag == "问题排查":
|
||
detail = "正在调用问题排查智能体"
|
||
elif route_flag == "百科问答":
|
||
detail = "正在调用百科问答智能体"
|
||
elif route_flag == "操作使用":
|
||
detail = "正在调用操作使用智能体"
|
||
elif route_flag == "统计":
|
||
detail = "正在调用统计智能体"
|
||
else:
|
||
detail = "正在调用通用智能体"
|
||
|
||
return {
|
||
"route_flag": route_flag,
|
||
"route_params": route_params,
|
||
"__detail__": detail,
|
||
}
|
||
|
||
|
||
def create_main_agent():
|
||
"""创建主脑Agent"""
|
||
workflow = StateGraph(MainAgentState)
|
||
|
||
workflow.add_node("detect_input_type", detect_input_type_node)
|
||
workflow.add_node("process_input", process_multimodal_input)
|
||
workflow.add_node("route", route_to_workflow)
|
||
|
||
workflow.add_edge(START, "detect_input_type")
|
||
workflow.add_edge("detect_input_type", "process_input")
|
||
workflow.add_edge("process_input", "route")
|
||
workflow.add_edge("route", END)
|
||
|
||
return workflow.compile()
|
||
|
||
|
||
from workflows.history_manager import filter_image_urls
|
||
|
||
|
||
async def extract_conversation_title(content: str) -> str:
|
||
"""提取对话标题"""
|
||
try:
|
||
prompt = f"请从以下对话内容中提取一个简短的标题（不超过20字）：\n\n{content[:500]}"
|
||
response = await OpenaiAPI.open_api_chat_without_thinking(
|
||
query=prompt,
|
||
model=None,
|
||
system_prompt="你是一个标题提取助手，请从对话内容中提取一个简短的标题。",
|
||
messages=[]
|
||
)
|
||
title = response.strip() if response else "新对话"
|
||
if len(title) > 30:
|
||
title = title[:30] + "..."
|
||
return title
|
||
except Exception:
|
||
return "新对话"
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import asyncio
|
||
|
||
async def test():
|
||
print("测试主脑Agent")
|
||
agent = create_main_agent()
|
||
result = await agent.ainvoke({
|
||
"raw_input": {"query": "101舰主机发生异响"},
|
||
"route_flag": "故障排查与修理"
|
||
})
|
||
print(f"结果: {result}")
|
||
|
||
asyncio.run(test())
|
||
|