langgraph_message_queue_nodes.py

"""
LangGraph Nodes with Message Queue Integration
将LangGraph节点与消息队列集成，实现异步通信
"""

import json
import asyncio
import logging
from typing import Dict, Any, List, Optional, TypedDict
from datetime import datetime
import time

from langgraph.graph import StateGraph, END, START
from langgraph.prebuilt import ToolNode
from message_queue_client import Message, create_message_queue_client

logger = logging.getLogger(__name__)


class GraphState(TypedDict):
    """图状态定义"""
    error: str
    messages: List
    generation: Dict[str, Any]
    iterations: int
    session_id: str
    validation_result: Dict[str, Any]
    code_solution: Dict[str, Any]
    debug_iterations: int  # debugger-tool 循环次数
    checker_iterations: int  # checker 审查次数
    has_execution_error: bool  # 是否有执行错误（用于 debugger-tool 循环）


class MessageQueueNode:
    """消息队列节点基类"""

    def __init__(self, node_name: str, queue_client, timeout: int = 30, max_retries: int = 3):
        self.node_name = node_name
        self.queue_client = queue_client
        self.timeout = timeout
        self.max_retries = max_retries
        self.pending_requests = {}  # 存储等待响应的请求

    async def publish_request(self, session_id: str, request_data: Dict[str, Any]) -> str:
        """发布请求到消息队列"""
        request_id = f"{session_id}_{self.node_name}_{int(time.time() * 1000)}"

        message = Message(
            agent_name=self.node_name,
            session_id=session_id,
            payload=request_data,
            message_id=request_id
        )

        self.queue_client.publish(f"{self.node_name}.requests", message)
        logger.info(f"Published request {request_id} to {self.node_name}.requests")

        return request_id

    async def wait_for_response(self, request_id: str, session_id: str) -> Optional[Dict[str, Any]]:
        """等待响应消息"""
        start_time = time.time()
        response_topic = f"{self.node_name}.responses"
        error_topic = f"{self.node_name}.errors"

        # 设置临时订阅
        response_received = None

        def handle_response(message: Message):
            nonlocal response_received
            if message.payload.get("request_id") == request_id:
                response_received = message.payload

        # 订阅响应和错误主题
        self.queue_client.subscribe(response_topic, handle_response)
        self.queue_client.subscribe(error_topic, handle_response)

        try:
            # 等待响应
            while response_received is None and (time.time() - start_time) < self.timeout:
                await asyncio.sleep(0.1)

            if response_received is None:
                raise TimeoutError(f"Timeout waiting for response to request {request_id}")

            return response_received

        finally:
            pass

    async def process_with_retry(self, session_id: str, request_data: Dict[str, Any]) -> Dict[str, Any]:
        """带重试的处理逻辑"""
        for attempt in range(self.max_retries):
            try:
                request_id = await self.publish_request(session_id, request_data)
                response = await self.wait_for_response(request_id, session_id)

                if response and response.get("success"):
                    return response.get("data", {})
                else:
                    error_msg = response.get("error", "Unknown error") if response else "No response"
                    logger.warning(f"Attempt {attempt + 1} failed: {error_msg}")

                    if attempt < self.max_retries - 1:
                        await asyncio.sleep(2 ** attempt)  # 指数退避
                    else:
                        raise Exception(f"All retry attempts failed. Last error: {error_msg}")

            except Exception as e:
                logger.error(f"Attempt {attempt + 1} exception: {e}")
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(2 ** attempt)
                else:
                    raise

        raise Exception("Max retries exceeded")


class CoderNode(MessageQueueNode):
    """Coder节点 - 生成代码"""

    def __init__(self, queue_client, **kwargs):
        super().__init__("coder", queue_client, **kwargs)

    async def generate_code(self, state: GraphState) -> GraphState:
        """生成代码解决方案"""
        logger.info(f"---GENERATING CODE SOLUTION for session {state['session_id']}---")

        messages = state["messages"]
        iterations = state["iterations"]
        error = state["error"]
        session_id = state["session_id"]

        # 如果有错误，添加重试提示
        if error == "yes":
            messages += [
                (
                    "user",
                    "Now, try again. Invoke the code tool to structure the output with a prefix, imports, and code block:",
                )
            ]

        # 构建请求数据
        request_data = {
            "context": "LCEL documentation and examples",  # 可以从外部获取
            "messages": messages,
            "session_id": session_id,
            "iterations": iterations
        }

        try:
            # 通过消息队列调用Coder服务
            code_solution = await self.process_with_retry(session_id, request_data)

            # 更新状态
            state["code_solution"] = code_solution
            state["iterations"] = iterations + 1
            state["error"] = "no"

            # 添加助手消息
            messages += [
                (
                    "assistant",
                    f"{code_solution.get('prefix', '')} \n Imports: {code_solution.get('imports', '')} \n Code: {code_solution.get('code', '')}",
                )
            ]
            state["messages"] = messages

            logger.info(f"Code generation completed for session {session_id}")

        except Exception as e:
            logger.error(f"Code generation failed for session {session_id}: {e}")
            state["error"] = "yes"
            messages += [("user", f"Code generation failed: {str(e)}")]
            state["messages"] = messages

        return state


class DebuggerNode(MessageQueueNode):
    """Debugger节点 - 调试代码"""

    def __init__(self, queue_client, **kwargs):
        super().__init__("debugger", queue_client, **kwargs)

    async def debug_code(self, state: GraphState) -> GraphState:
        """调试代码"""
        logger.info(f"---DEBUGGING CODE for session {state['session_id']} (debug iteration: {state.get('debug_iterations', 0)})---")

        messages = state["messages"]
        session_id = state["session_id"]
        code_solution = state.get("code_solution", {})

        # 构建调试请求数据
        request_data = {
            "code_solution": code_solution,
            "messages": messages,
            "session_id": session_id,
            "execution_error": state.get("generation", {}).get("error")  # 传递执行错误信息
        }

        try:
            # 通过消息队列调用Debugger服务
            debug_result = await self.process_with_retry(session_id, request_data)

            # 更新状态
            state["generation"] = debug_result
            
            # 如果发现错误，需要继续与 tool 迭代
            if debug_result.get("has_errors", False):
                state["has_execution_error"] = True
                error_info = debug_result.get("error_info", "Unknown debugging error")
                messages += [
                    ("assistant", f"Debug found issues: {error_info}"),
                ]
                logger.warning(f"Debug found errors for session {session_id}: {error_info}")
                
                # 更新代码解决方案（debugger 修复后的代码）
                if debug_result.get("fixed_code"):
                    state["code_solution"] = debug_result["fixed_code"]
            else:
                state["has_execution_error"] = False
                messages += [("assistant", "Debug completed successfully, no issues found.")]
                logger.info(f"Debug completed successfully for session {session_id}")

            state["messages"] = messages
            # 增加 debug 迭代次数
            state["debug_iterations"] = state.get("debug_iterations", 0) + 1

        except Exception as e:
            logger.error(f"Debug failed for session {session_id}: {e}")
            state["has_execution_error"] = True
            messages += [("user", f"Debug failed: {str(e)}")]
            state["messages"] = messages

        return state


class CheckerNode(MessageQueueNode):
    """Checker节点 - 检查代码"""

    def __init__(self, queue_client, **kwargs):
        super().__init__("checker", queue_client, **kwargs)

    async def check_code(self, state: GraphState) -> GraphState:
        """检查代码质量和正确性"""
        logger.info(f"---CHECKING CODE for session {state['session_id']} (checker iteration: {state.get('checker_iterations', 0)})---")

        messages = state["messages"]
        session_id = state["session_id"]
        code_solution = state.get("code_solution", {})
        generation = state.get("generation", {})

        # 构建检查请求数据
        request_data = {
            "code_solution": code_solution,
            "generation": generation,
            "messages": messages,
            "session_id": session_id
        }

        try:
            # 通过消息队列调用Checker服务
            validation_result = await self.process_with_retry(session_id, request_data)

            # 更新状态
            state["validation_result"] = validation_result
            
            # 检查验证结果
            if validation_result.get("is_valid", False):
                state["error"] = "no"
                messages += [
                    ("assistant", f"✅ Code validation passed: {validation_result.get('message', 'Code is valid')}")
                ]
                logger.info(f"Code validation passed for session {session_id}")
            else:
                state["error"] = "yes"
                validation_errors = validation_result.get("errors", "Validation failed")
                suggestions = validation_result.get("suggestions", "")
                
                messages += [
                    ("assistant", f"❌ Code validation failed: {validation_errors}"),
                    ("assistant", f"Suggestions: {suggestions}"),
                    ("user", "Please fix the code according to the validation feedback and ensure it meets the original requirements.")
                ]
                logger.warning(f"Code validation failed for session {session_id}: {validation_errors}")
                
                # 如果 checker 提供了修改建议的代码，更新 code_solution
                if validation_result.get("suggested_code"):
                    state["code_solution"] = validation_result["suggested_code"]
                    messages += [("assistant", "Code has been modified according to suggestions.")]
                
                # 重置 debug_iterations，因为要重新开始 debugger-tool 循环
                state["debug_iterations"] = 0
                state["has_execution_error"] = True  # 标记需要重新执行

            state["messages"] = messages
            # 增加 checker 迭代次数
            state["checker_iterations"] = state.get("checker_iterations", 0) + 1

        except Exception as e:
            logger.error(f"Code check failed for session {session_id}: {e}")
            state["error"] = "yes"
            messages += [("user", f"Code check failed: {str(e)}")]
            state["messages"] = messages

        return state


class ToolNode(MessageQueueNode):
    """Tool节点 - 执行代码工具"""

    def __init__(self, queue_client, **kwargs):
        super().__init__("tool", queue_client, **kwargs)

    async def execute_code(self, state: GraphState) -> GraphState:
        """执行代码并返回结果"""
        logger.info(f"---EXECUTING CODE for session {state['session_id']} (debug iteration: {state.get('debug_iterations', 0)})---")

        messages = state["messages"]
        session_id = state["session_id"]
        code_solution = state.get("code_solution", {})

        # 构建执行请求数据
        request_data = {
            "code": code_solution.get("code", ""),
            "imports": code_solution.get("imports", ""),
            "prefix": code_solution.get("prefix", ""),
            "session_id": session_id
        }

        try:
            # 通过消息队列调用Tool服务执行代码
            execution_result = await self.process_with_retry(session_id, request_data)

            # 更新状态
            state["generation"] = execution_result
            
            # 检查执行结果
            if execution_result.get("success", False):
                state["has_execution_error"] = False
                output = execution_result.get("output", "")
                messages += [
                    ("assistant", f"✅ Code executed successfully. Output:\n{output}")
                ]
                logger.info(f"Code execution succeeded for session {session_id}")
            else:
                state["has_execution_error"] = True
                error_msg = execution_result.get("error", "Execution failed")
                messages += [
                    ("assistant", f"❌ Code execution failed: {error_msg}"),
                    ("user", "The code has execution errors. Please debug and fix them.")
                ]
                logger.warning(f"Code execution failed for session {session_id}: {error_msg}")

            state["messages"] = messages

        except Exception as e:
            logger.error(f"Code execution failed for session {session_id}: {e}")
            state["has_execution_error"] = True
            messages += [("user", f"Code execution failed: {str(e)}")]
            state["messages"] = messages

        return state



# 决策函数
def decide_after_tool(state: GraphState) -> str:
    """
    Tool 执行后的决策：
    - 如果有执行错误且未超过最大调试迭代次数 -> 回到 debugger
    - 如果无执行错误 -> 去 checker 验证
    - 如果超过最大迭代次数 -> 强制去 checker
    """
    has_error = state.get("has_execution_error", False)
    debug_iterations = state.get("debug_iterations", 0)
    max_debug_iterations = 5  # debugger-tool 最大循环次数
    
    if has_error and debug_iterations < max_debug_iterations:
        logger.info(f"---DECISION: Tool found errors, return to DEBUGGER (iteration {debug_iterations}/{max_debug_iterations})---")
        return "debug"
    else:
        if has_error:
            logger.warning(f"---DECISION: Max debug iterations reached, forcing to CHECKER---")
        else:
            logger.info(f"---DECISION: No execution errors, proceed to CHECKER---")
        return "check_code"


def decide_after_checker(state: GraphState) -> str:
    """
    Checker 验证后的决策：
    - 如果验证通过 -> END
    - 如果验证不通过且未超过最大检查迭代次数 -> 回到 debugger
    - 如果超过最大迭代次数 -> 强制结束
    """
    is_valid = state.get("validation_result", {}).get("is_valid", False)
    checker_iterations = state.get("checker_iterations", 0)
    max_checker_iterations = 3  # checker 最大迭代次数
    
    if is_valid:
        logger.info(f"---DECISION: Validation PASSED, workflow END---")
        return "end"
    elif checker_iterations < max_checker_iterations:
        logger.warning(f"---DECISION: Validation FAILED, return to DEBUGGER (checker iteration {checker_iterations}/{max_checker_iterations})---")
        return "debug"
    else:
        logger.error(f"---DECISION: Max checker iterations reached, forcing END---")
        return "end"


def create_message_queue_workflow(queue_type: str = "redis", queue_config: Dict[str, Any] = None):
    """
    创建消息队列工作流
    
    工作流结构：
    START -> Coder -> Debugger -> Tool -> 决策1
                        ↑           |
                        |           v
                        +------[有错误]
                                   |
                              [无错误] -> Checker -> 决策2
                                            ↑        |
                                            |        v
                                            +---[不通过]
                                                    |
                                               [通过] -> END
    
    决策1 (decide_after_tool):
      - 有执行错误 -> 返回 Debugger (内层循环)
      - 无执行错误 -> 去 Checker
    
    决策2 (decide_after_checker):
      - 验证通过 -> END
      - 验证不通过 -> 返回 Debugger (外层循环，重置 debug_iterations)
    """
    # 创建消息队列客户端
    queue_client = create_message_queue_client(queue_type, **(queue_config or {}))

    # 创建节点
    coder_node = CoderNode(queue_client)
    checker_node = CheckerNode(queue_client)
    debugger_node = DebuggerNode(queue_client)
    tool_node = ToolNode(queue_client)
    
    # 创建工作流
    workflow = StateGraph(GraphState)

    # 定义节点
    workflow.add_node("generate", coder_node.generate_code)
    workflow.add_node("debug", debugger_node.debug_code)
    workflow.add_node("check_code", checker_node.check_code)
    workflow.add_node("tool", tool_node.execute_code)
    
    # 构建图的边
    # 1. START -> Coder
    workflow.add_edge(START, "generate")
    
    # 2. Coder -> Debugger
    workflow.add_edge("generate", "debug")
    
    # 3. Debugger -> Tool (debugger 总是将代码交给 tool 执行)
    workflow.add_edge("debug", "tool")
    
    # 4. Tool -> 条件分支 (decide_after_tool)
    #    - 有执行错误 -> 回到 Debugger
    #    - 无执行错误 -> 去 Checker
    workflow.add_conditional_edges(
        "tool",
        decide_after_tool,
        {
            "debug": "debug",           # 有错误，回到 debugger
            "check_code": "check_code"  # 无错误，去 checker
        }
    )
    
    # 5. Checker -> 条件分支 (decide_after_checker)
    #    - 验证通过 -> END
    #    - 验证不通过 -> 回到 Debugger
    workflow.add_conditional_edges(
        "check_code",
        decide_after_checker,
        {
            "debug": "debug",  # 验证失败，回到 debugger
            "end": END         # 验证通过，结束
        }
    )

    return workflow.compile()


# 异步工作流执行器
class AsyncWorkflowExecutor:
    """异步工作流执行器"""

    def __init__(self, workflow, queue_client):
        self.workflow = workflow
        self.queue_client = queue_client

    async def invoke(self, initial_state: GraphState) -> GraphState:
        """异步执行工作流"""
        try:
            logger.info(f"Starting workflow execution for session {initial_state['session_id']}")

            # 执行工作流，设置递归限制
            # 最坏情况：checker_iterations (3) * debug_iterations (5) * 2 = 30，设置为 50 足够
            result = await self.workflow.ainvoke(
                initial_state,
                config={"recursion_limit": 50}
            )

            logger.info(f"Workflow completed for session {initial_state['session_id']}")
            return result

        except Exception as e:
            logger.error(f"Workflow execution failed for session {initial_state['session_id']}: {e}")
            raise

    async def start_consuming(self):
        """启动消息消费"""
        if hasattr(self.queue_client, 'start_consuming'):
            self.queue_client.start_consuming()


# 使用示例
async def example_usage():
    """使用示例"""
    # 配置消息队列
    queue_config = {
        "host": "localhost",
        "port": 6379
    }

    # 创建工作流
    workflow = create_message_queue_workflow("redis", queue_config)

    # 创建执行器
    queue_client = create_message_queue_client("redis", **queue_config)
    executor = AsyncWorkflowExecutor(workflow, queue_client)

    # 初始状态
    initial_state = GraphState(
        error="no",
        messages=[("user", "How do I build a RAG chain in LCEL?")],
        generation={},
        iterations=0,
        session_id=f"test_session_{int(time.time())}",
        validation_result={},
        code_solution={},
        debug_iterations=0,
        checker_iterations=0,
        has_execution_error=False
    )

    try:
        # 执行工作流
        result = await executor.invoke(initial_state)

        # 输出结果
        if result.get("code_solution"):
            solution = result["code_solution"]
            print("Generated Code Solution:")
            print(f"Prefix: {solution.get('prefix', '')}")
            print(f"Imports: {solution.get('imports', '')}")
            print(f"Code: {solution.get('code', '')}")
        else:
            print("No code solution generated")

    except Exception as e:
        print(f"Workflow failed: {e}")

    finally:
        queue_client.close()


if __name__ == "__main__":
    # 运行示例
    asyncio.run(example_usage())