import asyncio import json import uuid from collections.abc import AsyncIterator from datetime import UTC, datetime from typing import Any from langchain.messages import AIMessage, AIMessageChunk, HumanMessage from langgraph.types import Command from yuxi import config as conf from yuxi.agents.buildin import agent_manager from yuxi.agents.context import build_agent_input_context, normalize_agent_context_config from yuxi.agents.state import AgentStatePayload from yuxi.repositories.agent_repository import AgentRepository from yuxi.repositories.agent_run_repository import AgentRunRepository from yuxi.repositories.conversation_repository import ConversationRepository from yuxi.services.conversation_service import serialize_attachment from yuxi.services.langfuse_service import ( LangfuseRunContext, build_run_context, flush_langfuse, get_trace_info, ) from yuxi.storage.postgres.manager import pg_manager from yuxi.storage.postgres.models_business import Agent, User from yuxi.utils.guard import content_guard from yuxi.utils.logging_config import logger from yuxi.utils.question_utils import ( normalize_questions as _normalize_interrupt_questions, ) def _build_state_files(attachments: list[dict]) -> dict: """将附件列表转换为 StateBackend 格式的 files 字典 StateBackend 期望的格式: { "/attachments/file.md": { "content": ["line1", "line2", ...], "created_at": "...", "modified_at": "...", } } """ files = {} for attachment in attachments: if attachment.get("status") != "parsed": continue file_path = attachment.get("file_path") markdown = attachment.get("markdown") if not file_path or not markdown: continue now = datetime.now(UTC).isoformat() # 将 markdown 内容按行拆分 content_lines = markdown.split("\n") files[file_path] = { "content": content_lines, "created_at": attachment.get("uploaded_at", now), "modified_at": attachment.get("uploaded_at", now), } return files async def _get_langgraph_messages(agent_instance, config_dict): graph = await agent_instance.get_graph() state = await graph.aget_state(config_dict) if not state or not state.values: logger.warning("No state found in LangGraph") return None return state.values.get("messages", []) def _build_langfuse_run_context( *, current_user, thread_id: str, agent_id: str, request_id: str, operation: str, backend_id: str | None = None, message_type: str | None = None, ) -> LangfuseRunContext: return build_run_context( user_id=str(getattr(current_user, "uid", current_user.id)), thread_id=thread_id, agent_id=agent_id, request_id=request_id, operation=operation, backend_id=backend_id, message_type=message_type, username=getattr(current_user, "username", None), login_user_id=getattr(current_user, "uid", None), department_id=getattr(current_user, "department_id", None), ) def extract_agent_state(values: dict) -> AgentStatePayload: """从 LangGraph state 中提取 agent 状态""" if not isinstance(values, dict): return {"todos": [], "files": {}, "artifacts": [], "subagent_runs": []} # 直接获取，信任 state 的数据结构 todos = values.get("todos") artifacts = values.get("artifacts") subagent_runs = values.get("subagent_runs") result: AgentStatePayload = { "todos": list(todos)[:20] if todos else [], "files": values.get("files") or {}, "artifacts": list(artifacts) if artifacts else [], "subagent_runs": list(subagent_runs) if subagent_runs else [], } return result def _agent_state_signature(agent_state: AgentStatePayload | dict | None) -> str: if not agent_state: return "" try: return json.dumps(agent_state, ensure_ascii=False, sort_keys=True) except Exception: return str(agent_state) def _metadata_thread_id(metadata: dict | None, fallback: str | None = None) -> str | None: if not isinstance(metadata, dict): return fallback for source in ( metadata, metadata.get("configurable"), metadata.get("metadata"), metadata.get("stream_event"), ): if isinstance(source, dict): value = source.get("thread_id") if isinstance(value, str) and value.strip(): return value.strip() return fallback def _metadata_namespace(metadata: dict | None) -> list[str]: if not isinstance(metadata, dict): return [] namespace = metadata.get("namespace") if isinstance(namespace, list): return [str(item) for item in namespace] return [] def _json_safe(value: Any) -> Any: if value is None or isinstance(value, str | int | float | bool): return value if isinstance(value, dict): return {str(key): _json_safe(child) for key, child in value.items()} if isinstance(value, list | tuple): return [_json_safe(child) for child in value] if hasattr(value, "model_dump"): return _json_safe(value.model_dump()) return str(value) def _stream_message_key(metadata: dict | None, namespace: list[str], thread_id: str | None) -> tuple[str, str]: if not isinstance(metadata, dict): return thread_id or "", "/".join(namespace) return thread_id or "", str(metadata.get("run_id") or metadata.get("langgraph_node") or "/".join(namespace)) def _stream_message_id( message_ids: dict[tuple[str, str], str], key: tuple[str, str], preferred: str | None = None, ) -> str: if preferred: message_ids[key] = preferred return preferred return message_ids.setdefault(key, str(uuid.uuid4())) def _message_chunk_yuxi_events( msg_dict: dict[str, Any], *, message_id: str, thread_id: str | None, namespace: list[str], ) -> list[dict[str, Any]]: events: list[dict[str, Any]] = [] route = {"thread_id": thread_id, "namespace": namespace} content = msg_dict.get("content") additional_kwargs = msg_dict.get("additional_kwargs") if isinstance(msg_dict.get("additional_kwargs"), dict) else {} reasoning_content = msg_dict.get("reasoning_content") additional_reasoning_content = additional_kwargs.get("reasoning_content") message_event: dict[str, Any] = {"type": "message_delta", "message_id": message_id, **route} if isinstance(content, str) and content: message_event["content"] = content if isinstance(reasoning_content, str) and reasoning_content: message_event["reasoning_content"] = reasoning_content if isinstance(additional_reasoning_content, str) and additional_reasoning_content: message_event["additional_reasoning_content"] = additional_reasoning_content if len(message_event) > 4: events.append(message_event) tool_call_chunks = msg_dict.get("tool_call_chunks") if isinstance(tool_call_chunks, list): for tool_call_chunk in tool_call_chunks: if not isinstance(tool_call_chunk, dict): continue args_delta = tool_call_chunk.get("args") if args_delta is None: args_delta = "" elif not isinstance(args_delta, str): args_delta = json.dumps(args_delta, ensure_ascii=False) if not tool_call_chunk.get("id") and not tool_call_chunk.get("name") and not args_delta: continue events.append( { "type": "tool_call_delta", "message_id": message_id, "tool_call_id": tool_call_chunk.get("id"), "name": tool_call_chunk.get("name"), "args_delta": args_delta, "index": tool_call_chunk.get("index") if tool_call_chunk.get("index") is not None else 0, **route, } ) return events def _protocol_event_yuxi_event( event: dict[str, Any], *, message_id: str | None, thread_id: str | None, namespace: list[str], ) -> dict[str, Any] | None: event_name = event.get("event") if event_name in {"message-start", "content-block-start", "message-finish"} or not message_id: return None route = {"thread_id": thread_id, "namespace": namespace} if event_name == "content-block-delta": delta = event.get("delta") if isinstance(event.get("delta"), dict) else {} text = delta.get("text") if delta.get("type") == "text-delta" and isinstance(text, str) and text: return {"type": "message_delta", "message_id": message_id, "content": text, **route} return None if event_name == "content-block-finish": content = event.get("content") if isinstance(event.get("content"), dict) else {} if content.get("type") != "tool_call" or not content.get("id") and not content.get("name"): return None return { "type": "tool_call", "message_id": message_id, "tool_call_id": content.get("id"), "name": content.get("name"), "args": content.get("args") if content.get("args") is not None else {}, "index": event.get("index") if event.get("index") is not None else 0, **route, } return None def _stream_event_response(event: dict[str, Any]) -> str: if event.get("type") != "message_delta": return "" return str(event.get("content") or "") def _message_payload_yuxi_events( msg: Any, *, metadata: dict[str, Any], namespace: list[str], thread_id: str | None, protocol_message_ids: dict[tuple[str, str], str], ) -> list[dict[str, Any]]: message_key = _stream_message_key(metadata, namespace, thread_id) if isinstance(msg, dict) and isinstance(msg.get("event"), str): preferred_message_id = str(msg["id"]) if msg.get("event") == "message-start" and msg.get("id") else None message_id = _stream_message_id(protocol_message_ids, message_key, preferred_message_id) stream_event = _protocol_event_yuxi_event( msg, message_id=message_id, thread_id=thread_id, namespace=namespace, ) return [stream_event] if stream_event else [] if isinstance(msg, AIMessageChunk) or hasattr(msg, "model_dump"): msg_dict = msg.model_dump() elif isinstance(msg, dict): msg_dict = dict(msg) else: msg_dict = {"content": str(msg)} message_id = str(msg_dict.get("id") or _stream_message_id(protocol_message_ids, message_key)) return _message_chunk_yuxi_events( msg_dict, message_id=message_id, thread_id=thread_id, namespace=namespace, ) async def _stream_agent_events(agent, messages, *, input_context=None, **kwargs): if hasattr(agent, "stream_messages_with_state"): async for mode, payload in agent.stream_messages_with_state( messages, input_context=input_context, **kwargs, ): yield mode, payload return async for msg, metadata in agent.stream_messages(messages, input_context=input_context, **kwargs): yield "messages", (msg, metadata) async def _get_existing_message_ids(conv_repo: ConversationRepository, thread_id: str) -> set[str]: existing_messages = await conv_repo.get_messages_by_thread_id(thread_id) return { msg.extra_metadata["id"] for msg in existing_messages if msg.extra_metadata and "id" in msg.extra_metadata and isinstance(msg.extra_metadata["id"], str) } async def _save_ai_message( conv_repo: ConversationRepository, thread_id: str, msg_dict: dict, trace_info: dict[str, Any] | None = None, ) -> None: content = msg_dict.get("content", "") tool_calls_data = msg_dict.get("tool_calls") or [] if isinstance(content, list): if not tool_calls_data: tool_calls_data = [ {"id": item.get("id"), "name": item.get("name"), "args": item.get("args") or {}} for item in content if isinstance(item, dict) and item.get("type") == "tool_call" ] content = "\n".join( item.get("text", "") for item in content if isinstance(item, dict) and isinstance(item.get("text"), str) ) elif not isinstance(content, str): content = str(content) extra_metadata = dict(msg_dict) if trace_info: extra_metadata.update(trace_info) ai_msg = await conv_repo.add_message_by_thread_id( thread_id=thread_id, role="assistant", content=content, message_type="text", extra_metadata=extra_metadata, ) if ai_msg and tool_calls_data: for tc in tool_calls_data: await conv_repo.add_tool_call( message_id=ai_msg.id, tool_name=tc.get("name", "unknown"), tool_input=tc.get("args", {}), status="pending", langgraph_tool_call_id=tc.get("id"), ) async def _save_tool_message(conv_repo: ConversationRepository, msg_dict: dict) -> None: tool_call_id = msg_dict.get("tool_call_id") content = msg_dict.get("content", "") if not tool_call_id: return if isinstance(content, list): tool_output = json.dumps(content) if content else "" else: tool_output = str(content) await conv_repo.update_tool_call_output( langgraph_tool_call_id=tool_call_id, tool_output=tool_output, status="success", ) async def save_partial_message( conv_repo: ConversationRepository, thread_id: str, full_msg=None, error_message: str | None = None, error_type: str = "interrupted", trace_info: dict[str, Any] | None = None, ): try: extra_metadata = { "error_type": error_type, "is_error": True, "error_message": error_message or f"发生错误: {error_type}", } if full_msg: msg_dict = full_msg.model_dump() if hasattr(full_msg, "model_dump") else {} content = full_msg.content if hasattr(full_msg, "content") else str(full_msg) extra_metadata = msg_dict | extra_metadata else: content = "" if trace_info: extra_metadata.update(trace_info) return await conv_repo.add_message_by_thread_id( thread_id=thread_id, role="assistant", content=content, message_type="text", extra_metadata=extra_metadata, ) except Exception as e: logger.exception(f"Error saving message: {e}") return None async def save_messages_from_langgraph_state( agent_instance, thread_id: str, conv_repo: ConversationRepository, config_dict: dict, trace_info: dict[str, Any] | None = None, ) -> None: messages = await _get_langgraph_messages(agent_instance, config_dict) if messages is None: return existing_ids = await _get_existing_message_ids(conv_repo, thread_id) for msg in messages: if hasattr(msg, "model_dump"): msg_dict = msg.model_dump() elif isinstance(msg, dict): msg_dict = dict(msg) else: continue msg_type = msg_dict.get("type", "unknown") if msg_type == "unknown": role = msg_dict.get("role") if role in {"assistant", "ai"}: msg_type = "ai" elif role in {"user", "human"}: msg_type = "human" elif role == "tool": msg_type = "tool" msg_id = getattr(msg, "id", None) or msg_dict.get("id") if msg_type == "human" or msg_id in existing_ids: continue if msg_type == "ai": await _save_ai_message(conv_repo, thread_id, msg_dict, trace_info=trace_info) elif msg_type == "tool": await _save_tool_message(conv_repo, msg_dict) def _extract_interrupt_info(state) -> Any | None: """从 LangGraph state 中提取中断信息""" if hasattr(state, "tasks") and state.tasks: for task in state.tasks: if hasattr(task, "interrupts") and task.interrupts: return task.interrupts[0] interrupt_data = state.values.get("__interrupt__") if isinstance(interrupt_data, list) and interrupt_data: return interrupt_data[0] return None def _coerce_interrupt_payload(info: Any) -> dict: """将 LangGraph interrupt 对象转换为 dict 结构。""" if isinstance(info, dict): return info payload = getattr(info, "value", None) if isinstance(payload, dict): return payload questions = getattr(info, "questions", None) source = getattr(info, "source", None) result: dict[str, Any] = {} if isinstance(questions, list): result["questions"] = questions if isinstance(source, str) and source.strip(): result["source"] = source return result def _build_ask_user_question_payload(info: Any, thread_id: str) -> dict[str, Any]: """将 interrupt 信息标准化为 ask_user_question_required 载荷。""" payload = _coerce_interrupt_payload(info) questions = _normalize_interrupt_questions(payload.get("questions")) if not questions: questions = [ { "question_id": str(uuid.uuid4()), "question": "请选择一个选项", "options": [], "multi_select": False, "allow_other": True, } ] source = str(payload.get("source") or payload.get("tool_name") or "interrupt") return { "questions": questions, "source": source, "thread_id": thread_id, } def _ensure_full_msg(full_msg: AIMessage | None, accumulated_content: list[str]) -> AIMessage | None: """如果 full_msg 为空且有累积内容，构建 AIMessage""" if not full_msg and accumulated_content: return AIMessage(content="".join(accumulated_content)) return full_msg def _extract_ai_message(messages: list[Any] | None) -> AIMessage | None: """从消息列表中提取最后一条 AIMessage。""" if not isinstance(messages, list): return None for msg in reversed(messages): if isinstance(msg, AIMessage): return msg msg_dict = msg.model_dump() if hasattr(msg, "model_dump") else {} if msg_dict.get("type") == "ai": content = msg_dict.get("content", "") return msg if hasattr(msg, "content") else AIMessage(content=content) return None async def _resolve_agent_runtime( *, db, user: User, requested_agent_id: str | None, thread_id: str | None, ) -> tuple[Agent, Any, dict]: agent_repo = AgentRepository(db) conv_repo = ConversationRepository(db) bound_agent_id = requested_agent_id if thread_id: conversation = await conv_repo.get_conversation_by_thread_id(thread_id) if conversation: if conversation.uid != str(user.uid) or conversation.status == "deleted": raise ValueError("对话线程不存在") if requested_agent_id and requested_agent_id != conversation.agent_id: raise ValueError("已有线程已绑定智能体，不能切换") bound_agent_id = conversation.agent_id if not bound_agent_id: raise ValueError("缺少必需的 agent_id 字段") agent_item = await agent_repo.get_visible_by_slug(slug=bound_agent_id, user=user) if not agent_item: raise ValueError("智能体不存在或无权限访问") backend = agent_manager.get_agent(agent_item.backend_id) if not backend: raise ValueError(f"智能体后端 {agent_item.backend_id} 不存在") agent_config = await normalize_agent_context_config( (agent_item.config_json or {}).get("context", {}), db=db, user=user, context_schema=backend.context_schema, ) return agent_item, backend, agent_config async def check_and_handle_interrupts( agent, langgraph_config: dict, make_chunk, meta: dict, thread_id: str, ) -> AsyncIterator[bytes]: try: graph = await agent.get_graph() state = await graph.aget_state(langgraph_config) if not state or not state.values: return interrupt_info = _extract_interrupt_info(state) if interrupt_info: question_payload = _build_ask_user_question_payload(interrupt_info, thread_id) meta["interrupt"] = question_payload yield make_chunk(status="ask_user_question_required", meta=meta, **question_payload) except Exception as e: logger.exception(f"Error checking interrupts: {e}") async def _ensure_thread_bound_agent( *, conv_repo: ConversationRepository, thread_id: str, uid: str, agent_item: Agent, ) -> None: conversation = await conv_repo.get_conversation_by_thread_id(thread_id) if not conversation: await conv_repo.create_conversation( uid=uid, agent_id=agent_item.slug, thread_id=thread_id, metadata={"backend_id": agent_item.backend_id}, ) return if conversation.agent_id != agent_item.slug: raise ValueError("已有线程已绑定智能体，不能切换") def _normalize_attachment_file_ids(meta: dict | None) -> list[str]: file_ids = (meta or {}).get("attachment_file_ids") or [] if not isinstance(file_ids, list): return [] normalized = [] seen = set() for file_id in file_ids: value = str(file_id).strip() if not value or value in seen: continue seen.add(value) normalized.append(value) return normalized async def _bind_request_attachments( *, conv_repo: ConversationRepository, thread_id: str, request_id: str, attachment_file_ids: list[str], ) -> list[dict]: conversation = await conv_repo.get_conversation_by_thread_id(thread_id) if not conversation: return [] if attachment_file_ids: attachments = await conv_repo.bind_attachments_to_request(conversation.id, request_id, attachment_file_ids) else: attachments = await conv_repo.get_attachments_by_request_id(conversation.id, request_id) return [serialize_attachment(attachment) for attachment in attachments] async def agent_chat( *, query: str, agent_id: str, thread_id: str | None, meta: dict, image_content: str | None, current_user, db, ) -> dict: """非流式对话，返回完整响应""" start_time = asyncio.get_event_loop().time() if image_content: human_message = HumanMessage( content=[ {"type": "text", "text": query}, {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_content}"}}, ] ) message_type = "multimodal_image" else: human_message = HumanMessage(content=query) message_type = "text" if conf.enable_content_guard and await content_guard.check(query): return { "status": "error", "error_type": "content_guard_blocked", "error_message": "输入内容包含敏感词", "request_id": meta.get("request_id"), } uid = str(current_user.uid) meta = dict(meta or {}) if "request_id" not in meta or not meta.get("request_id"): logger.warning("请求缺少 request_id，已自动生成一个新的 request_id") meta["request_id"] = str(uuid.uuid4()) if not thread_id: thread_id = str(uuid.uuid4()) logger.warning(f"No thread_id provided, generated new thread_id: {thread_id}") try: agent_item, agent, agent_config = await _resolve_agent_runtime( db=db, user=current_user, requested_agent_id=agent_id, thread_id=thread_id, ) except ValueError as e: return { "status": "error", "error_type": "invalid_agent", "error_message": str(e), "request_id": meta.get("request_id"), } meta.update( { "query": query, "agent_id": agent_item.slug, "backend_id": agent_item.backend_id, "server_model_name": agent_item.backend_id, "thread_id": thread_id, "uid": current_user.uid, "has_image": bool(image_content), } ) messages = [human_message] input_context = await build_agent_input_context( agent_config, thread_id=thread_id, uid=uid, run_id=meta.get("run_id"), request_id=meta.get("request_id"), ) langfuse_run = _build_langfuse_run_context( current_user=current_user, thread_id=thread_id, agent_id=agent_item.slug, backend_id=agent_item.backend_id, request_id=meta["request_id"], operation="agent_chat_sync", message_type=message_type, ) trace_info: dict[str, Any] = {} try: conv_repo = ConversationRepository(db) await _ensure_thread_bound_agent( conv_repo=conv_repo, thread_id=thread_id, uid=uid, agent_item=agent_item, ) request_attachments = await _bind_request_attachments( conv_repo=conv_repo, thread_id=thread_id, request_id=meta["request_id"], attachment_file_ids=_normalize_attachment_file_ids(meta), ) try: await conv_repo.add_message_by_thread_id( thread_id=thread_id, role="user", content=query, message_type=message_type, image_content=image_content, extra_metadata={ "raw_message": human_message.model_dump(), "request_id": meta.get("request_id"), "attachments": request_attachments, }, ) except Exception as e: logger.error(f"Error saving user message: {e}") langgraph_config = {"configurable": {"thread_id": thread_id, "uid": uid}} invoke_result = await agent.invoke_messages( messages, input_context=input_context, callbacks=langfuse_run.callbacks, metadata=langfuse_run.metadata, tags=langfuse_run.tags, ) full_msg = _extract_ai_message(invoke_result.get("messages") if isinstance(invoke_result, dict) else None) trace_info = get_trace_info(langfuse_run) if full_msg is None: try: graph = await agent.get_graph() state = await graph.aget_state(langgraph_config) full_msg = _extract_ai_message(getattr(state, "values", {}).get("messages", [])) if state else None except Exception: full_msg = None full_content = full_msg.content if full_msg else "" if conf.enable_content_guard and await content_guard.check(full_content): await save_partial_message( conv_repo, thread_id, full_msg, "content_guard_blocked", trace_info=trace_info, ) return { "status": "interrupted", "message": "检测到敏感内容，已中断输出", "request_id": meta.get("request_id"), "time_cost": asyncio.get_event_loop().time() - start_time, } try: graph = await agent.get_graph() state = await graph.aget_state(langgraph_config) agent_state = extract_agent_state(getattr(state, "values", {})) if state else {} except Exception: agent_state = {} try: await save_messages_from_langgraph_state( agent_instance=agent, thread_id=thread_id, conv_repo=conv_repo, config_dict=langgraph_config, trace_info=trace_info, ) except Exception as e: logger.exception(f"Error saving messages from LangGraph state: {e}") return { "status": "error", "error_type": "save_message_error", "error_message": f"消息保存失败: {e}", "request_id": meta.get("request_id"), } return { "status": "finished", "response": full_content, "request_id": meta.get("request_id"), "thread_id": thread_id, "agent_state": agent_state, "time_cost": asyncio.get_event_loop().time() - start_time, } except Exception as e: logger.exception(f"Error in agent_chat: {e}") return { "status": "error", "error_type": "unexpected_error", "error_message": str(e), "request_id": meta.get("request_id"), } finally: flush_langfuse() async def stream_agent_chat( *, query: str, agent_id: str, thread_id: str | None, meta: dict, image_content: str | None, current_user, db, save_user_message: bool = True, ) -> AsyncIterator[bytes]: start_time = asyncio.get_event_loop().time() def make_chunk(content=None, **kwargs): chunk_thread_id = kwargs.pop("thread_id", None) or meta.get("thread_id") or thread_id return ( json.dumps( {"request_id": meta.get("request_id"), "response": content, "thread_id": chunk_thread_id, **kwargs}, ensure_ascii=False, ).encode("utf-8") + b"\n" ) meta = dict(meta or {}) if "request_id" not in meta or not meta.get("request_id"): logger.warning("请求缺少 request_id，已自动生成一个新的 request_id") meta["request_id"] = str(uuid.uuid4()) uid = str(current_user.uid) if not thread_id: thread_id = str(uuid.uuid4()) logger.warning(f"No thread_id provided, generated new thread_id: {thread_id}") if image_content: human_message = HumanMessage( content=[ {"type": "text", "text": query}, {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_content}"}}, ] ) message_type = "multimodal_image" else: human_message = HumanMessage(content=query) message_type = "text" if conf.enable_content_guard and await content_guard.check(query): yield make_chunk( status="error", error_type="content_guard_blocked", error_message="输入内容包含敏感词", meta=meta ) return try: agent_item, agent, agent_config = await _resolve_agent_runtime( db=db, user=current_user, requested_agent_id=agent_id, thread_id=thread_id, ) except ValueError as e: yield make_chunk(status="error", error_type="invalid_agent", error_message=str(e), meta=meta) return meta.update( { "query": query, "agent_id": agent_item.slug, "backend_id": agent_item.backend_id, "server_model_name": agent_item.backend_id, "thread_id": thread_id, "uid": current_user.uid, "has_image": bool(image_content), } ) messages = [human_message] input_context = await build_agent_input_context( agent_config, thread_id=thread_id, uid=uid, run_id=meta.get("run_id"), request_id=meta.get("request_id"), ) langfuse_run = _build_langfuse_run_context( current_user=current_user, thread_id=thread_id, agent_id=agent_item.slug, backend_id=agent_item.backend_id, request_id=meta["request_id"], operation="agent_chat_stream", message_type=message_type, ) full_msg = None accumulated_content: list[str] = [] trace_info: dict[str, Any] = {} last_agent_state_signature = "" try: conv_repo = ConversationRepository(db) await _ensure_thread_bound_agent( conv_repo=conv_repo, thread_id=thread_id, uid=uid, agent_item=agent_item, ) request_attachments = await _bind_request_attachments( conv_repo=conv_repo, thread_id=thread_id, request_id=meta["request_id"], attachment_file_ids=_normalize_attachment_file_ids(meta), ) init_msg = { "role": "user", "content": query, "type": "human", "message_type": message_type, "extra_metadata": { "request_id": meta.get("request_id"), "attachments": request_attachments, }, } if image_content: init_msg["image_content"] = image_content yield make_chunk(status="init", meta=meta, msg=init_msg) if save_user_message: try: await conv_repo.add_message_by_thread_id( thread_id=thread_id, role="user", content=query, message_type=message_type, image_content=image_content, extra_metadata={ "raw_message": human_message.model_dump(), "request_id": meta.get("request_id"), "attachments": request_attachments, }, ) except Exception as e: logger.error(f"Error saving user message: {e}") # 先构建 langgraph_config langgraph_config = {"configurable": {"thread_id": thread_id, "uid": uid}} # LangGraph 会自动从 checkpointer 恢复 state（包括 uploads） # 无需手动加载或传递 full_msg = None accumulated_content = [] protocol_message_ids: dict[tuple[str, str], str] = {} async for mode, payload in _stream_agent_events( agent, messages, input_context=input_context, callbacks=langfuse_run.callbacks, metadata=langfuse_run.metadata, tags=langfuse_run.tags, ): if mode == "values": agent_state = extract_agent_state(payload if isinstance(payload, dict) else {}) signature = _agent_state_signature(agent_state) if signature and signature != last_agent_state_signature: last_agent_state_signature = signature yield make_chunk(status="agent_state", agent_state=agent_state, meta=meta) continue if mode == "stream_event": yield make_chunk( status="stream_event", event=payload, namespace=payload.get("namespace") if isinstance(payload, dict) else [], meta=meta, thread_id=payload.get("thread_id") if isinstance(payload, dict) else None, ) continue msg, metadata = payload namespace = _metadata_namespace(metadata) chunk_thread_id = _metadata_thread_id(metadata, thread_id if not namespace else None) if namespace and not chunk_thread_id: continue is_subagent_chunk = bool(chunk_thread_id and chunk_thread_id != thread_id) stream_events = _message_payload_yuxi_events( msg, metadata=metadata, namespace=namespace, thread_id=chunk_thread_id, protocol_message_ids=protocol_message_ids, ) for stream_event in stream_events: content = _stream_event_response(stream_event) if not is_subagent_chunk and content: trace_info = get_trace_info(langfuse_run) accumulated_content.append(content) content_for_check = "".join(accumulated_content[-10:]) if conf.enable_content_guard and await content_guard.check_with_keywords(content_for_check): full_msg = AIMessage(content="".join(accumulated_content)) await save_partial_message( conv_repo, thread_id, full_msg, "content_guard_blocked", trace_info=trace_info, ) meta["time_cost"] = asyncio.get_event_loop().time() - start_time yield make_chunk(status="interrupted", message="检测到敏感内容，已中断输出", meta=meta) return yield make_chunk( content=content, stream_event=stream_event, metadata=metadata, status="loading", thread_id=chunk_thread_id, ) full_msg = _ensure_full_msg(full_msg, accumulated_content) trace_info = get_trace_info(langfuse_run) if conf.enable_content_guard and hasattr(full_msg, "content") and await content_guard.check(full_msg.content): await save_partial_message( conv_repo, thread_id, full_msg, "content_guard_blocked", trace_info=trace_info, ) meta["time_cost"] = asyncio.get_event_loop().time() - start_time yield make_chunk(status="interrupted", message="检测到敏感内容，已中断输出", meta=meta) return interrupted = False async for chunk in check_and_handle_interrupts(agent, langgraph_config, make_chunk, meta, thread_id): interrupted = True yield chunk meta["time_cost"] = asyncio.get_event_loop().time() - start_time try: graph = await agent.get_graph() state = await graph.aget_state(langgraph_config) agent_state = extract_agent_state(getattr(state, "values", {})) if state else {} except Exception: agent_state = {} final_signature = _agent_state_signature(agent_state) if final_signature and final_signature != last_agent_state_signature: last_agent_state_signature = final_signature yield make_chunk(status="agent_state", agent_state=agent_state, meta=meta) # 先存储数据库，再返回 finished，避免前端查询时数据未落库 try: await save_messages_from_langgraph_state( agent_instance=agent, thread_id=thread_id, conv_repo=conv_repo, config_dict=langgraph_config, trace_info=trace_info, ) except Exception as e: logger.exception(f"Error saving messages from LangGraph state: {e}") yield make_chunk(status="warning", message=f"消息保存失败: {e}", meta=meta) if interrupted: return yield make_chunk(status="finished", meta=meta) except (asyncio.CancelledError, ConnectionError) as e: logger.warning(f"Client disconnected, cancelling stream: {e}") async def save_cleanup(): nonlocal full_msg full_msg = _ensure_full_msg(full_msg, accumulated_content) async with pg_manager.get_async_session_context() as new_db: new_conv_repo = ConversationRepository(new_db) await save_partial_message( new_conv_repo, thread_id, full_msg=full_msg, error_message="对话已中断" if not full_msg else None, error_type="interrupted", trace_info=trace_info, ) cleanup_task = asyncio.create_task(save_cleanup()) try: await asyncio.shield(cleanup_task) except asyncio.CancelledError: pass except Exception as exc: logger.error(f"Error during cleanup save: {exc}") yield make_chunk(status="interrupted", message="对话已中断", meta=meta) except Exception as e: logger.exception(f"Error streaming messages: {e}") error_msg = f"Error streaming messages: {e}" error_type = "unexpected_error" full_msg = _ensure_full_msg(full_msg, accumulated_content) async with pg_manager.get_async_session_context() as new_db: new_conv_repo = ConversationRepository(new_db) await save_partial_message( new_conv_repo, thread_id, full_msg=full_msg, error_message=error_msg, error_type=error_type, trace_info=trace_info, ) yield make_chunk(status="error", error_type=error_type, error_message=error_msg, meta=meta) finally: flush_langfuse() async def stream_agent_resume( *, thread_id: str, resume_input: Any, meta: dict, current_user, db, ) -> AsyncIterator[bytes]: start_time = asyncio.get_event_loop().time() def make_resume_chunk(content=None, **kwargs): chunk_thread_id = kwargs.pop("thread_id", None) or meta.get("thread_id") or thread_id return ( json.dumps( {"request_id": meta.get("request_id"), "response": content, "thread_id": chunk_thread_id, **kwargs}, ensure_ascii=False, ).encode("utf-8") + b"\n" ) yield make_resume_chunk(status="init", meta=meta) resume_command = Command(resume=resume_input) uid = str(current_user.uid) try: agent_item, agent, agent_config = await _resolve_agent_runtime( db=db, user=current_user, requested_agent_id=None, thread_id=thread_id, ) except ValueError as e: yield make_resume_chunk(status="error", error_type="invalid_agent", error_message=str(e), meta=meta) return meta["agent_id"] = agent_item.slug meta["backend_id"] = agent_item.backend_id input_context = await build_agent_input_context( agent_config or {}, thread_id=thread_id, uid=uid, run_id=meta.get("run_id"), request_id=meta.get("request_id"), ) context = agent.context_schema() context.update(input_context) langfuse_run = _build_langfuse_run_context( current_user=current_user, thread_id=thread_id, agent_id=agent_item.slug, backend_id=agent_item.backend_id, request_id=meta.get("request_id") or str(uuid.uuid4()), operation="agent_chat_resume", message_type="resume", ) trace_info: dict[str, Any] = {} last_agent_state_signature = "" stream_source = agent.stream_resume_with_state( resume_command, input_context=input_context, callbacks=langfuse_run.callbacks, metadata=langfuse_run.metadata, tags=langfuse_run.tags, ) protocol_message_ids: dict[tuple[str, str], str] = {} try: async for mode, payload in stream_source: if mode == "values": agent_state = extract_agent_state(payload if isinstance(payload, dict) else {}) signature = _agent_state_signature(agent_state) if signature and signature != last_agent_state_signature: last_agent_state_signature = signature yield make_resume_chunk(status="agent_state", agent_state=agent_state, meta=meta) continue if mode == "stream_event": event_payload = payload if isinstance(payload, dict) else {} yield make_resume_chunk( status="stream_event", event=event_payload, namespace=event_payload.get("namespace") or [], meta=meta, thread_id=event_payload.get("thread_id"), ) continue if mode != "messages": continue msg, metadata = payload metadata = dict(metadata or {}) namespace = _metadata_namespace(metadata) chunk_thread_id = _metadata_thread_id(metadata, thread_id if not namespace else None) if namespace and not chunk_thread_id: continue if chunk_thread_id == thread_id: trace_info = get_trace_info(langfuse_run) stream_events = _message_payload_yuxi_events( msg, metadata=metadata, namespace=namespace, thread_id=chunk_thread_id, protocol_message_ids=protocol_message_ids, ) for stream_event in stream_events: content = _stream_event_response(stream_event) yield make_resume_chunk( content=content, stream_event=stream_event, metadata=metadata, status="loading", thread_id=chunk_thread_id, ) langgraph_config = {"configurable": {"thread_id": thread_id, "uid": uid}} interrupted = False async for chunk in check_and_handle_interrupts(agent, langgraph_config, make_resume_chunk, meta, thread_id): interrupted = True yield chunk meta["time_cost"] = asyncio.get_event_loop().time() - start_time try: graph = await agent.get_graph(context=context) state = await graph.aget_state(langgraph_config) agent_state = extract_agent_state(getattr(state, "values", {})) if state else {} except Exception: agent_state = {} final_signature = _agent_state_signature(agent_state) if final_signature and final_signature != last_agent_state_signature: yield make_resume_chunk(status="agent_state", agent_state=agent_state, meta=meta) # 先存储数据库，再返回 finished，避免前端查询时数据未落库 conv_repo = ConversationRepository(db) try: await save_messages_from_langgraph_state( agent_instance=agent, thread_id=thread_id, conv_repo=conv_repo, config_dict=langgraph_config, trace_info=trace_info, ) except Exception as e: logger.exception(f"Error saving messages from LangGraph state: {e}") yield make_resume_chunk(status="warning", message=f"消息保存失败: {e}", meta=meta) if interrupted: return yield make_resume_chunk(status="finished", meta=meta) except (asyncio.CancelledError, ConnectionError) as e: logger.warning(f"Client disconnected during resume: {e}") async with pg_manager.get_async_session_context() as new_db: new_conv_repo = ConversationRepository(new_db) await save_partial_message( new_conv_repo, thread_id, error_message="对话恢复已中断", error_type="resume_interrupted", trace_info=trace_info, ) yield make_resume_chunk(status="interrupted", message="对话恢复已中断", meta=meta) except Exception as e: logger.exception(f"Error during resume: {e}") async with pg_manager.get_async_session_context() as new_db: new_conv_repo = ConversationRepository(new_db) await save_partial_message( new_conv_repo, thread_id, error_message=f"Error during resume: {e}", error_type="resume_error", trace_info=trace_info, ) yield make_resume_chunk(message=f"Error during resume: {e}", status="error") finally: flush_langfuse() def _serialize_state_messages(values: dict[str, Any]) -> list[dict[str, Any]]: messages = values.get("messages") if isinstance(values, dict) else None if not isinstance(messages, list): return [] serialized = [] for message in messages: if hasattr(message, "model_dump"): serialized.append(message.model_dump()) elif isinstance(message, dict): serialized.append(dict(message)) else: serialized.append({"type": "unknown", "content": str(message)}) return serialized async def _read_checkpoint_state(agent, *, uid: str, thread_id: str): graph = await agent.get_graph() langgraph_config = {"configurable": {"uid": uid, "thread_id": thread_id}} return await graph.aget_state(langgraph_config) def _serialize_subagent_run(run) -> dict[str, Any]: payload = run.input_payload if isinstance(run.input_payload, dict) else {} return { "id": payload.get("tool_call_id") or run.id, "run_id": run.id, "subagent_type": payload.get("subagent_type") or run.agent_id, "subagent_name": payload.get("subagent_name"), "child_thread_id": payload.get("child_thread_id") or run.thread_id, "description": payload.get("description"), "status": run.status, "created_at": run.to_dict().get("created_at"), "completed_at": run.to_dict().get("finished_at"), "result_preview": payload.get("result_preview"), "error": run.error_message, "parent_agent_run_id": run.parent_agent_run_id, } async def get_agent_state_view( *, thread_id: str, current_uid: str, db, include_messages: bool = False, ) -> dict: from fastapi import HTTPException conv_repo = ConversationRepository(db) agent_repo = AgentRepository(db) conversation = await conv_repo.get_conversation_by_thread_id(thread_id) if conversation: if conversation.uid != str(current_uid) or conversation.status == "deleted": raise HTTPException(status_code=404, detail="对话线程不存在") agent_item = await agent_repo.get_by_slug(conversation.agent_id) if not agent_item: raise HTTPException(status_code=404, detail="智能体不存在") agent = agent_manager.get_agent(agent_item.backend_id) if not agent: raise HTTPException(status_code=404, detail="智能体后端不存在") state = await _read_checkpoint_state(agent, uid=str(current_uid), thread_id=thread_id) values = getattr(state, "values", {}) if state else {} response = {"agent_state": extract_agent_state(values)} if include_messages: response["messages"] = _serialize_state_messages(values) return response run_repo = AgentRunRepository(db) subagent_run = await run_repo.get_latest_subagent_run_by_thread_for_user(thread_id, str(current_uid)) if not subagent_run or not subagent_run.parent_agent_run_id: raise HTTPException(status_code=404, detail="对话线程不存在") parent_run = await run_repo.get_run_for_user(subagent_run.parent_agent_run_id, str(current_uid)) if not parent_run: raise HTTPException(status_code=404, detail="对话线程不存在") parent_conversation = await conv_repo.get_conversation_by_thread_id(parent_run.thread_id) if ( not parent_conversation or parent_conversation.id != parent_run.conversation_id or parent_conversation.uid != str(current_uid) or parent_conversation.status == "deleted" ): raise HTTPException(status_code=404, detail="对话线程不存在") child_agent_item = await agent_repo.get_by_slug(subagent_run.agent_id) if not child_agent_item: raise HTTPException(status_code=404, detail="智能体不存在") child_agent = agent_manager.get_agent(child_agent_item.backend_id) if not child_agent: raise HTTPException(status_code=404, detail="智能体后端不存在") checkpoint_thread_id = subagent_run.checkpoint_thread_id or subagent_run.thread_id child_state = await _read_checkpoint_state(child_agent, uid=str(current_uid), thread_id=checkpoint_thread_id) child_values = getattr(child_state, "values", {}) if child_state else {} response = { "agent_state": extract_agent_state(child_values), "parent_thread_id": parent_run.thread_id, "subagent_run": _serialize_subagent_run(subagent_run), } if include_messages: response["messages"] = _serialize_state_messages(child_values) return response