agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/shared/message.py

@@ -0,0 +1,419 @@
+import os
+import logging
+import traceback
+from datetime import datetime
+import json
+from typing import Dict, Any, AsyncGenerator, Union, List, Tuple
+from agent_api_server.shared.util_func import load_graph_config, load_graph, get_env
+from langgraph.types import StateSnapshot
+from agent_api_server.cache.redis_cache import ThreadState
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    ToolMessage, AIMessageChunk)
+from .detect_message import detect_content_type
+from langgraph.types import Interrupt, Command
+from agent_api_server.shared.base_model import ChatMessage
+from fastapi import Body, status, HTTPException
+from langfuse.langchain import CallbackHandler
+from agent_api_server.shared.decode_token import decode_jwt
+from agent_api_server.cache.redis_cache import AsyncRedisThreadStorage
+from langfuse import propagate_attributes
+
+
+logger = logging.getLogger(__name__)
+
+# Type aliases for better readability
+MessageContent = Union[str, List[Any]]
+EventData = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
+StreamEvent = Union[Tuple[str, Any], Tuple[str, str, Any]]
+SerializableData = Union[Dict[str, Any], List[Any], str, int, float, bool, None, datetime]
+
+
+def format_state_snapshot(snapshot: StateSnapshot) -> Dict[str, Any]:
+    """Format state snapshot into a serializable dictionary."""
+    interrupts = []
+    if hasattr(snapshot, 'tasks'):
+        for task in snapshot.tasks:
+            if hasattr(task, 'interrupts'):
+                for i in task.interrupts:
+                    if hasattr(i, 'value'):
+                        interrupts.append({"value": i.value})
+
+    return {
+        "values": getattr(snapshot, 'values', None),
+        "next": getattr(snapshot, 'next', None),
+        "config": getattr(snapshot, 'config', None),
+        "interrupts": interrupts,
+        "parent_config": getattr(snapshot, 'parent_config', None),
+        "metadata": getattr(snapshot, 'metadata', None)
+    }
+
+
+def convert_message_content(content: MessageContent) -> str:
+    """Convert message content to string representation."""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        return "".join(
+            item.get("text") if isinstance(item, dict) and item.get("type") == "text"
+            else str(item)
+            for item in content
+        )
+    return str(content)
+
+
+def langchain_to_chat_message(message: BaseMessage) -> ChatMessage:
+    """Convert LangChain message to ChatMessage with proper type handling."""
+    content = convert_message_content(message.content)
+    content_type = detect_content_type(content)
+
+    # Handle specific message types
+    if isinstance(message, HumanMessage):
+        message_type = "human"
+        tool_calls = []
+        response_metadata = {}
+        references = []
+    elif isinstance(message, ToolMessage):
+        message_type = "tool"
+        tool_calls = []
+        response_metadata = {}
+        references = []
+    else:
+        message_type = "ai"
+        tool_calls = getattr(message, "tool_calls", [])
+        response_metadata = getattr(message, "response_metadata", {})
+        additional_kwargs = getattr(message, "additional_kwargs", {})
+        references = additional_kwargs.get("references", [])
+    return ChatMessage(
+        type=message_type,
+        content_type=content_type,
+        content=content,
+        tool_calls=tool_calls,  # 确保是 list
+        response_metadata=response_metadata,  # 确保是 dict
+        references=references
+    )
+
+
+def serialize_data(data: Any) -> SerializableData:
+    """Recursively serialize data to JSON-compatible format."""
+    if isinstance(data, (str, int, float, bool)) or data is None:
+        return data
+    if isinstance(data, dict):
+        return {k: serialize_data(v) for k, v in data.items()}
+    if isinstance(data, (list, tuple, set)):
+        return [serialize_data(item) for item in data]
+    if isinstance(data, datetime):
+        return data.isoformat()
+
+    # Handle objects with serialization methods
+    if hasattr(data, 'model_dump'):
+        return serialize_data(data.model_dump())
+    if hasattr(data, 'dict'):
+        return serialize_data(data.dict())
+    if hasattr(data, '__dict__'):
+        return serialize_data(data.__dict__)
+
+    return str(data)
+
+
+def format_sse_event(event_type: str, data: Any) -> str:
+    """Format data as Server-Sent Event (SSE)."""
+    try:
+        # 确保 event_data 是字典，且 data 能被安全合并
+        event_data = {"event": event_type}
+
+        # 如果 data 是字典，直接合并；否则作为 raw_data 字段
+        if isinstance(data, dict):
+            event_data.update(data)
+        else:
+            event_data["data"] = str(data)  # 非字典类型转为字符串存储
+
+        payload = json.dumps(event_data, ensure_ascii=False)
+        return f"data: {payload}\n\n"
+    except (TypeError, ValueError) as e:
+        # 错误处理：返回带错误信息的 SSE
+        error_payload = {
+            "event": "error",
+            "exception": str(e),
+            "original_data": str(data)[:200]  # 截断避免过长
+        }
+        payload = json.dumps(error_payload, ensure_ascii=False)
+        return f"data: {payload}\n\n"
+
+
+async def process_node_updates(node: str, updates: Dict[str, Any]) -> AsyncGenerator[
+    str, None]:
+    try:
+        # Handle agent nodes specially
+        if node == "agent":
+            messages = updates.get("messages", [])
+            if messages:
+                message = messages[-1]
+                chat_msg = langchain_to_chat_message(message)
+                serialized_content = serialize_data(chat_msg)
+            else:
+                node_message = AIMessage(content=str(updates))
+                chat_msg = langchain_to_chat_message(node_message)
+                serialized_content = serialize_data(chat_msg)
+
+            yield format_sse_event("agent_message", {
+                "node": node,
+                "update_content": serialized_content,
+                "update_type": "complete",
+                "timestamp": datetime.now().isoformat()
+            })
+        elif node == "__interrupt__":
+            interrupt: Interrupt
+            for interrupt in updates:
+                node_message = AIMessage(content=str(interrupt.value))
+                chat_msg = langchain_to_chat_message(node_message)
+                serialized_content = serialize_data(chat_msg)
+
+                yield format_sse_event("interrupt_message", {
+                    "node": node,
+                    "update_content": serialized_content,
+                    "update_type": "complete",
+                    "timestamp": datetime.now().isoformat()
+                })
+        elif node == "tools":
+            messages = updates.get("messages", [])
+            if messages:
+                message = messages[-1]
+                chat_msg = langchain_to_chat_message(message)
+                serialized_content = serialize_data(chat_msg)
+            else:
+                node_message = AIMessage(content=str(updates))
+                chat_msg = langchain_to_chat_message(node_message)
+                serialized_content = serialize_data(chat_msg)
+
+            yield format_sse_event("tools_message", {
+                "node": node,
+                "update_content": serialized_content,
+                "update_type": "complete",
+                "timestamp": datetime.now().isoformat()
+            })
+        elif node == "supervisor":
+            messages = updates.get("messages", [])
+            if messages:
+                message = messages[-1]
+                chat_msg = langchain_to_chat_message(message)
+                serialized_content = serialize_data(chat_msg)
+            else:
+                node_message = AIMessage(content=str(updates))
+                chat_msg = langchain_to_chat_message(node_message)
+                serialized_content = serialize_data(chat_msg)
+
+            yield format_sse_event("supervisor_message", {
+                "node": node,
+                "update_content": serialized_content,
+                "update_type": "complete",
+                "timestamp": datetime.now().isoformat()
+            })
+        else:
+            messages = updates.get("messages", [])
+            if messages:
+                message = messages[-1]
+                chat_msg = langchain_to_chat_message(message)
+                serialized_content = serialize_data(chat_msg)
+            else:
+                node_message = AIMessage(content=str(updates))
+                chat_msg = langchain_to_chat_message(node_message)
+                serialized_content = serialize_data(chat_msg)
+
+            yield format_sse_event("node_message", {
+                "node": node,
+                "update_content": serialized_content,
+                "update_type": "complete",
+                "timestamp": datetime.now().isoformat()
+            })
+    except Exception as e:
+        logger.error(f"Failed to process {node} updates: {str(e)}", exc_info=True)
+        yield format_sse_event("error", {
+            "node": node,
+            "error": f"Message processing error: {str(e)}",
+            "timestamp": datetime.now().isoformat()
+        })
+
+
+async def handle_stream_event(stream_event: StreamEvent) -> AsyncGenerator[str, None]:
+    """Process stream events and yield formatted SSE events."""
+    try:
+        if isinstance(stream_event, tuple) and len(stream_event) == 2:
+            stream_mode, event = stream_event
+
+            if stream_mode == "updates":
+                if isinstance(event, dict):
+                    for node, updates in event.items():
+                        async for chunk in process_node_updates(node, updates):
+                            yield chunk
+            elif stream_mode == "messages":
+                if isinstance(event, tuple) and len(event) == 2:
+                    message_chunk, metadata = event
+
+                    is_ai_chunk = False
+                    if isinstance(message_chunk, AIMessageChunk):
+                        is_ai_chunk = True
+
+                    if is_ai_chunk:
+                        if message_chunk.content and message_chunk.content.strip():
+                            yield format_sse_event("token_stream", {
+                                "node": metadata.get('langgraph_node'),
+                                "update_content": message_chunk.content,
+                                "update_type": "start",
+                                "timestamp": datetime.now().isoformat()
+                            })
+
+                        finish_reason = getattr(message_chunk, 'response_metadata', {}).get('finish_reason')
+                        if finish_reason:
+                            complete_message = AIMessage(
+                                content=getattr(message_chunk, 'content', ''),
+                                tool_calls=getattr(message_chunk, 'tool_calls', []),
+                                response_metadata=getattr(message_chunk, 'response_metadata', {})
+                            )
+
+                            chat_msg = langchain_to_chat_message(complete_message)
+                            serialized_content = serialize_data(chat_msg)
+
+                            yield format_sse_event("token_stream", {
+                                "node": metadata.get('langgraph_node'),
+                                "update_type": "complete",
+                                "update_content": message_chunk.content if hasattr(message_chunk, 'content') else "",
+                                "timestamp": datetime.now().isoformat()
+                            })
+                else:
+                    logger.warning(f"Unexpected messages event format: {type(event)}")
+            elif stream_mode == "custom":
+                logger.info(f"receive custom stream event, event message is {event}")
+                yield event
+
+        elif isinstance(stream_event, tuple) and len(stream_event) == 3:
+            _, stream_mode, event = stream_event
+            async for chunk in handle_stream_event((stream_mode, event)):
+                yield chunk
+        else:
+            logger.warning(f"Unexpected event format: {type(stream_event)}")
+            yield format_sse_event("error", {
+                "error": f"Unexpected event format: {type(stream_event)}",
+                "event_data": str(stream_event)[:200]
+            })
+    except Exception as e:
+        logger.error(f"Error processing event: {str(e)}", exc_info=True)
+        yield format_sse_event("error", {
+            "error": f"Event processing error: {str(e)}",
+            "event_type": type(stream_event).__name__ if 'stream_event' in locals() else "unknown"
+        })
+
+
+async def message_generator(
+        state: ThreadState,
+        ts_tenant: str,
+        ei_token: str,
+        inputs: Dict[str, Any] = Body(..., embed=True),
+        files: List[Dict[str, Any]] = Body(default=[], embed=True)
+) -> AsyncGenerator[str, None]:
+    """Generate message stream for the given thread."""
+    storage = AsyncRedisThreadStorage.get_worker_instance()
+
+    user_id = None
+    if ei_token:
+        try:
+            header, payload, user_info = decode_jwt(ei_token)
+            if user_info and 'id' in user_info:
+                user_id = user_info['id']
+                logger.info(f"Extracted user_id from token: {user_id}")
+            else:
+                logger.warning("No user_id found in JWT token")
+        except Exception as e:
+            logger.warning(f"Failed to decode JWT token: {str(e)}")
+
+    try:
+        _, graph_instance, checkpointer = await load_graph(state.graph_name, await load_graph_config(), True)
+        data = await storage.update_thread(state.thread_id, status="running")
+        if not data:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail={
+                    "error": "update_failed",
+                    "message": "Failed to update thread status",
+                    "thread_id": state.thread_id
+                }
+            )
+
+        subgraph_names = {name for name, _ in graph_instance.get_subgraphs()}
+        logger.info(f"Found graphs: {subgraph_names}, resume message is: {inputs.get('resume', '')}")
+
+        configurable_params = {
+            **dict(get_env(ts_tenant=ts_tenant)),
+            "graph_name": state.graph_name,
+            "thread_id": state.thread_id,
+            "TSTenant": ts_tenant,
+            "EIToken": ei_token,
+            "files": files or []
+        }
+
+        langfuse_keys = [
+            os.getenv("LANGFUSE_SECRET_KEY"),
+            os.getenv("LANGFUSE_PUBLIC_KEY"),
+            os.getenv("LANGFUSE_BASE_URL")
+        ]
+
+        callbacks_config = {}
+        if all(langfuse_keys):
+            langfuse_handler = CallbackHandler()
+            callbacks_config = {"callbacks": [langfuse_handler], "run_name": state.graph_name}
+
+        config = {"configurable": configurable_params, **callbacks_config}
+
+        state_get = await graph_instance.aget_state(
+            config={"configurable": {**dict(get_env(ts_tenant=ts_tenant)), "thread_id": state.thread_id}}
+        )
+        interrupted_tasks = [
+            task for task in state_get.tasks if hasattr(task, "interrupts") and task.interrupts
+        ]
+
+        async def process_stream():
+            if interrupted_tasks:
+                logger.info(f"find resume message, so run {state.thread_id} for graph {state.graph_name} with resume")
+                with propagate_attributes(session_id=state.thread_id,
+                                          user_id=user_id, trace_name=f"{state.graph_name}") if user_id else propagate_attributes(
+                    session_id=state.thread_id, trace_name=f"{state.graph_name}"):
+                    async for stream_event in graph_instance.astream(
+                            Command(resume=inputs.get("resume", "")),
+                            config=config,
+                            stream_mode=["updates", "messages", "custom"],
+                            subgraphs=True
+                    ):
+                        async for chunk in handle_stream_event(stream_event):
+                            yield chunk
+
+            else:
+                logger.info(f"do not find resume message, so first run {state.thread_id} for graph {state.graph_name}")
+                with propagate_attributes(session_id=state.thread_id,
+                                          user_id=user_id, trace_name=f"{state.graph_name}") if user_id else propagate_attributes(
+                    session_id=state.thread_id, trace_name=f"{state.graph_name}"):
+                    async for stream_event in graph_instance.astream(
+                            inputs or {},
+                            config=config,
+                            stream_mode=["updates", "messages", "custom"],
+                            subgraphs=True
+                    ):
+                        async for chunk in handle_stream_event(stream_event):
+                            yield chunk
+
+        async for chunk in process_stream():
+            yield chunk
+
+    except HTTPException as e:
+        yield format_sse_event("error", f"Stream processing failed: {e.detail}")
+    except (RuntimeError, ValueError) as e:
+        yield format_sse_event("error", f"Stream processing failed: {e}")
+    except Exception as e:
+        logging.error("Stream processing failed:\n%s", traceback.format_exc())
+        yield format_sse_event("error", f"Stream processing failed: {str(e)}")
+    finally:
+        data = await storage.update_thread(state.thread_id, status="complete")
+        if not data:
+            yield format_sse_event("error", f"Failed to update thread status")
+        yield "data: [DONE]\n\n"