tw/run.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI 客服系统 - 统一启动器

用法:
  python run.py                        # WebSocket 客服模式（默认）
  python run.py --tianwang             # 完整版（HTTP API + WebSocket + AI Agent）
  python run.py --tianwang-multi -w 4  # 天网 + 多进程（HTTP API + 多进程 WebSocket）
  python run.py --api-only             # 仅 HTTP API（不含 WebSocket / AI Agent）
  python run.py --no-agent             # 不启用 AI Agent
  python run.py --multi -w 4           # 多进程模式，4 个 Worker
  python run.py --port 6060            # 指定 HTTP API 端口
"""
import sys
import os
import signal
import logging
import argparse
from pathlib import Path

_root = Path(__file__).resolve().parent
if str(_root) not in sys.path:
    sys.path.insert(0, str(_root))

from core.websocket_logger_setup import setup_logger

logger = setup_logger()

DEFAULT_HTTP_PORT = 6060
DEFAULT_HTTP_HOST = "127.0.0.1"


def _print_api_info(host: str, port: int):
    logger.info(f"HTTP API 服务器已启动：http://{host}:{port}")
    logger.info("")
    logger.info("天网任务接口:")
    logger.info("  POST /api/task/receive    - 接收任务")
    logger.info("  POST /api/task/cancel     - 取消任务")
    logger.info("  GET  /api/task/status/:id - 查询任务状态")
    logger.info("  GET  /api/task/list       - 任务列表")
    logger.info("  GET  /api/health          - 健康检查")


def run_websocket(enable_agent: bool):
    """WebSocket 客服模式（默认）"""
    import asyncio
    from core.websocket_client_v2 import QingjianAPIClient

    logger.info("=" * 60)
    logger.info("AI 客服系统 - WebSocket 模式 (新架构 v2)")
    logger.info(f"AI Agent: {'已启用' if enable_agent else '未启用'}")
    logger.info("=" * 60)

    client = QingjianAPIClient(enable_agent=enable_agent)
    try:
        asyncio.run(client.run())
    except KeyboardInterrupt:
        logger.info("已停止")


def run_tianwang(enable_agent: bool, host: str, port: int):
    """完整版: HTTP API + WebSocket + AI Agent"""
    import asyncio
    from api.http_server import start_http_server
    from core.websocket_client_v2 import QingjianAPIClient

    logger.info("=" * 60)
    logger.info("AI 客服系统 - 天网协作版（完整）")
    logger.info("=" * 60)

    start_http_server(host=host, port=port)
    _print_api_info(host, port)
    logger.info("=" * 60)

    logger.info("正在连接轻简 API...")
    client = QingjianAPIClient(enable_agent=enable_agent)
    logger.info(f"AI Agent: {'已启用' if enable_agent else '未启用'}")
    logger.info("=" * 60)
    logger.info("系统已就绪，等待消息和任务...")

    def _signal_handler(signum, frame):
        logger.info("收到退出信号，正在停止...")
        if hasattr(client, 'task_scheduler'):
            asyncio.run(client.task_scheduler.stop())
        sys.exit(0)

    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    try:
        asyncio.run(client.connect())
    except KeyboardInterrupt:
        logger.info("已停止")
    except Exception as e:
        logger.error(f"运行异常：{e}")
        sys.exit(1)


def run_api_only(host: str, port: int):
    """仅 HTTP API（不含 WebSocket / AI Agent）"""
    import time
    from api.http_server import start_http_server

    logger.info("=" * 60)
    logger.info("AI 客服系统 - 天网协作版（仅 API）")
    logger.info("=" * 60)

    start_http_server(host=host, port=port)
    _print_api_info(host, port)
    logger.info("=" * 60)
    logger.info("系统已就绪，等待天网任务...")

    def _signal_handler(signum, frame):
        logger.info("收到退出信号")
        sys.exit(0)

    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        logger.info("已停止")


def run_multi_process(num_workers: int, enable_agent: bool):
    """多进程模式"""
    from scripts.multi_process_launcher import Coordinator

    logger.info("=" * 60)
    logger.info("AI 客服系统 - 多进程异步并行模式")
    logger.info(f"工作进程数：{num_workers}")
    logger.info(f"AI Agent: {'已启用' if enable_agent else '未启用'}")
    logger.info("=" * 60)

    coordinator = Coordinator(num_workers=num_workers, enable_agent=enable_agent)
    try:
        coordinator.start()
    except KeyboardInterrupt:
        logger.info("已停止")
        coordinator.stop()

def run_tianwang_multi(num_workers: int, enable_agent: bool, host: str, port: int):
    """天网 + 多进程：HTTP API + 多进程 WebSocket 客户端"""
    from api.http_server import start_http_server
    from scripts.multi_process_launcher import Coordinator

    logger.info("=" * 60)
    logger.info("AI 客服系统 - 天网协作版（HTTP API） + 多进程模式")
    logger.info("=" * 60)

    start_http_server(host=host, port=port)
    _print_api_info(host, port)
    logger.info("=" * 60)

    logger.info(f"工作进程数：{num_workers}")
    logger.info(f"AI Agent: {'已启用' if enable_agent else '未启用'}")
    logger.info("=" * 60)

    coordinator = Coordinator(num_workers=num_workers or 0, enable_agent=enable_agent)

    def _signal_handler(signum, frame):
        logger.info("收到退出信号，正在停止多进程协调器...")
        coordinator.stop()
        sys.exit(0)

    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    try:
        coordinator.start()
    except KeyboardInterrupt:
        logger.info("已停止")
        coordinator.stop()

def main():
    parser = argparse.ArgumentParser(
        description='AI 客服系统启动器',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
启动模式:
  (默认)       WebSocket 客服模式
  --tianwang   HTTP API + WebSocket（天网完整版）
  --tianwang-multi HTTP API + 多进程 WebSocket
  --api-only   仅 HTTP API（不含 WebSocket / AI Agent）
  --multi      多进程模式
        """
    )

    mode = parser.add_mutually_exclusive_group()
    mode.add_argument('--tianwang', action='store_true', help='天网完整版（HTTP API + WebSocket）')
    mode.add_argument('--tianwang-multi', action='store_true', help='天网 + 多进程（HTTP API + 多进程 WebSocket）')
    mode.add_argument('--api-only', action='store_true', help='仅 HTTP API（不含 WebSocket / AI Agent）')
    mode.add_argument('--multi', action='store_true', help='多进程模式')

    parser.add_argument('--no-agent', action='store_true', help='不启用 AI Agent')
    parser.add_argument('--host', type=str, default=DEFAULT_HTTP_HOST, help=f'HTTP API 监听地址（默认 {DEFAULT_HTTP_HOST}）')
    parser.add_argument('--port', '-p', type=int, default=DEFAULT_HTTP_PORT, help=f'HTTP API 端口（默认 {DEFAULT_HTTP_PORT}）')
    parser.add_argument('--workers', '-w', type=int, default=None, help='工作进程数（仅多进程模式，默认 CPU 核心数）')

    args = parser.parse_args()
    enable_agent = not args.no_agent

    if args.api_only:
        run_api_only(host=args.host, port=args.port)
    elif args.tianwang:
        run_tianwang(enable_agent=enable_agent, host=args.host, port=args.port)
    elif args.tianwang_multi:
        run_tianwang_multi(num_workers=args.workers, enable_agent=enable_agent, host=args.host, port=args.port)
    elif args.multi:
        run_multi_process(num_workers=args.workers, enable_agent=enable_agent)
    else:
        run_websocket(enable_agent=enable_agent)


if __name__ == "__main__":
    main()