main_langgraph.py

"""
Main LangGraph Application Entry Point
LangGraph应用主入口，集成消息队列架构
"""

import asyncio
import logging
import argparse
import json
from datetime import datetime
from typing import Dict, Any, Optional

from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from langgraph_message_queue_nodes import (
    create_message_queue_workflow,
    AsyncWorkflowExecutor,
    GraphState
)
from message_queue_client import create_message_queue_client
from monitoring import SystemMonitor, create_monitoring_app
from error_handling import CircuitBreakerConfig, RetryConfig

# 配置日志
import os
log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(os.path.join(log_dir, 'langgraph-app.log')),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class CodeGenerationRequest(BaseModel):
    """代码生成请求"""
    question: str
    context: str = ""
    session_id: str = None
    max_iterations: int = 3


class CodeGenerationResponse(BaseModel):
    """代码生成响应"""
    success: bool
    session_id: str
    code_solution: Dict[str, Any]
    iterations: int
    execution_time: float
    error: Optional[str] = None


class LangGraphApplication:
    """LangGraph应用主类"""

    def __init__(self, queue_type: str = "redis", queue_config: Dict[str, Any] = None):
        self.queue_type = queue_type
        self.queue_config = queue_config or {}
        self.queue_client = None
        self.workflow = None
        self.executor = None
        self.monitor = None
        self.app = None

    async def initialize(self):
        """初始化应用"""
        logger.info("Initializing LangGraph application...")

        # 创建消息队列客户端
        self.queue_client = create_message_queue_client(self.queue_type, **self.queue_config)

        # 创建工作流
        self.workflow = create_message_queue_workflow(self.queue_type, self.queue_config)

        # 创建执行器
        self.executor = AsyncWorkflowExecutor(self.workflow, self.queue_client)

        # 创建监控器
        self.monitor = SystemMonitor("langgraph_app")

        # 注册健康检查
        self.monitor.register_health_check("langgraph_workflow", self._health_check)

        # 创建FastAPI应用
        self.app = self._create_fastapi_app()

        logger.info("LangGraph application initialized successfully")

    async def _health_check(self) -> bool:
        """健康检查"""
        try:
            # 检查消息队列连接
            if not self.queue_client:
                return False

            # 检查工作流状态
            if not self.workflow:
                return False

            return True
        except Exception as e:
            logger.error(f"Health check failed: {e}")
            return False

    def _create_fastapi_app(self) -> FastAPI:
        """创建FastAPI应用"""
        app = FastAPI(
            title="Code Assistant with Message Queue Architecture",
            version="1.0.0",
            description="A code generation system using LangGraph and message queues"
        )

        # 监控API
        monitoring_app = create_monitoring_app(self.monitor)
        app.mount("/monitoring", monitoring_app)

        @app.on_event("startup")
        async def startup_event():
            await self.start()

        @app.on_event("shutdown")
        async def shutdown_event():
            await self.stop()

        @app.get("/health")
        async def health_check():
            """健康检查端点"""
            is_healthy = await self._health_check()
            return {
                "status": "healthy" if is_healthy else "unhealthy",
                "timestamp": datetime.utcnow().isoformat(),
                "service": "langgraph_app"
            }

        @app.get("/")
        async def root():
            """根端点"""
            return {
                "message": "Code Assistant with Message Queue Architecture",
                "version": "1.0.0",
                "endpoints": {
                    "generate": "/generate",
                    "health": "/health",
                    "monitoring": "/monitoring"
                }
            }

        @app.post("/generate", response_model=CodeGenerationResponse)
        async def generate_code(request: CodeGenerationRequest):
            """生成代码端点"""
            start_time = datetime.utcnow()

            try:
                # 生成会话ID
                session_id = request.session_id or f"session_{datetime.utcnow().timestamp()}"

                # 创建初始状态
                initial_state = GraphState(
                    error="no",
                    messages=[("user", request.question)],
                    generation={},
                    iterations=0,
                    session_id=session_id,
                    validation_result={},
                    code_solution={}
                )

                # 记录请求
                self.monitor.record_request("generate", "started", 0)

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

                # 计算执行时间
                execution_time = (datetime.utcnow() - start_time).total_seconds()

                # 检查是否成功生成代码
                success = bool(result.get("code_solution", {}))

                # 记录完成
                self.monitor.record_request("generate", "success" if success else "failed", execution_time)

                return CodeGenerationResponse(
                    success=success,
                    session_id=session_id,
                    code_solution=result.get("code_solution", {}),
                    iterations=result.get("iterations", 0),
                    execution_time=execution_time,
                    error=None if success else "Failed to generate code"
                )

            except Exception as e:
                execution_time = (datetime.utcnow() - start_time).total_seconds()
                self.monitor.record_request("generate", "error", execution_time)

                logger.error(f"Code generation failed: {e}")
                return CodeGenerationResponse(
                    success=False,
                    session_id=request.session_id or "unknown",
                    code_solution={},
                    iterations=0,
                    execution_time=execution_time,
                    error=str(e)
                )

        @app.get("/status")
        async def get_status():
            """获取系统状态"""
            try:
                dashboard_data = self.monitor.get_dashboard_data()
                return JSONResponse(content=dashboard_data)
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))

        return app

    async def start(self):
        """启动应用"""
        logger.info("Starting LangGraph application...")

        # 启动监控
        await self.monitor.start_monitoring()

        # 启动消息消费
        if hasattr(self.executor, 'start_consuming'):
            asyncio.create_task(self.executor.start_consuming())

        logger.info("LangGraph application started successfully")

    async def stop(self):
        """停止应用"""
        logger.info("Stopping LangGraph application...")

        # 停止监控
        await self.monitor.stop_monitoring()

        # 关闭消息队列连接
        if self.queue_client:
            self.queue_client.close()

        logger.info("LangGraph application stopped")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="LangGraph Code Assistant Application")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
    parser.add_argument("--port", type=int, default=8004, help="Port to bind to")
    parser.add_argument("--queue-type", default="redis", choices=["redis", "rabbitmq"], help="Message queue type")
    parser.add_argument("--redis-host", default="localhost", help="Redis host")
    parser.add_argument("--redis-port", type=int, default=6379, help="Redis port")
    parser.add_argument("--rabbitmq-host", default="localhost", help="RabbitMQ host")
    parser.add_argument("--rabbitmq-port", type=int, default=5672, help="RabbitMQ port")

    args = parser.parse_args()

    # 配置消息队列
    queue_config = {}
    if args.queue_type == "redis":
        queue_config = {
            "host": args.redis_host,
            "port": args.redis_port
        }
    elif args.queue_type == "rabbitmq":
        queue_config = {
            "host": args.rabbitmq_host,
            "port": args.rabbitmq_port
        }

    # 创建应用
    app = LangGraphApplication(args.queue_type, queue_config)

    # 运行应用
    import uvicorn

    async def run():
        await app.initialize()
        config = uvicorn.Config(app.app, host=args.host, port=args.port, log_level="info")
        server = uvicorn.Server(config)
        await server.serve()

    asyncio.run(run())


if __name__ == "__main__":
    main()