main.py

import asyncio
import json
import time
import os
import aio_pika

from mminfo_extractor.mminfo_extractor import extract_template, extract_knowledge_and_innovation, ExtractionConfig
from config import settings
from mminfo_extractor.utils.rabbitmq_manager import RabbitMQManager
from mminfo_extractor.utils.logging_config import setup_logging, get_logger

# 配置日志系统，只显示关键信息
setup_logging(level="INFO")  # 保持INFO级别，但减少冗余日志
logger = get_logger(__name__)

extraction_config = ExtractionConfig(
    local_file_dir="local_files",
    supported_template_types=["化学", "机械", "电学"],
    llm_model_name=settings.LLM_MODEL_NAME,
    llm_base_url=settings.LLM_BASE_URL,
    llm_api_key=settings.LLM_API_KEY,
    vllm_model_name=settings.VLLM_MODEL_NAME,
    vllm_base_url=settings.VLLM_BASE_URL,
    vllm_api_key=settings.VLLM_API_KEY,
    # 并发控制配置
    max_concurrent_requests=settings.MAX_CONCURRENT_REQUESTS,
    requests_per_minute=settings.REQUESTS_PER_MINUTE,
    request_delay=settings.REQUEST_DELAY,
    max_retries=settings.MAX_RETRIES,
    retry_delay=settings.RETRY_DELAY,
    request_timeout=settings.REQUEST_TIMEOUT,
    # 下载配置
    download_max_concurrent=settings.DOWNLOAD_MAX_CONCURRENT,
    download_requests_per_minute=settings.DOWNLOAD_REQUESTS_PER_MINUTE,
    download_limit=settings.DOWNLOAD_LIMIT,
    download_limit_per_host=settings.DOWNLOAD_LIMIT_PER_HOST,
    host_check_ttl_seconds=settings.HOST_CHECK_TTL_SECONDS,

    enable_token_counting=settings.ENABLE_TOKEN_COUNTING,
    enable_enhance_image=settings.ENABLE_ENHANCE_IMAGE,
    enable_enhance_formula=settings.ENABLE_ENHANCE_FORMULA,
    enable_enhance_table=settings.ENABLE_ENHANCE_TABLE,

    mineru_server_url=settings.MINERU_SERVER_URL,
    mineru_api_key=settings.MINERU_API_KEY,
)

async def send_to_dead_letter_queue(rabbitmq_manager: RabbitMQManager, message_body: dict, error: Exception, message_type: str):
    """发送失败的消息到死信队列（如果启用）"""
    # 检查是否启用了死信队列
    if not getattr(settings, 'ENABLE_DEAD_LETTER_QUEUE', False):
        logger.warning(f"死信队列未启用，仅记录错误日志: {message_body.get('task_id', 'unknown')}")
        logger.error(f"任务失败详情: {error}")
        return
    
    try:
        dead_letter_message = {
            "original_message": message_body,
            "error": str(error),
            "error_type": type(error).__name__,
            "message_type": message_type,
            "timestamp": time.time(),
            "task_id": message_body.get("task_id", "unknown")
        }
        
        dead_letter_json = json.dumps(dead_letter_message, ensure_ascii=False)
        
        # 发送到死信队列
        await rabbitmq_manager.publish_message(
            message_body=dead_letter_json,
            routing_key=settings.DEAD_LETTER_ROUTING_KEY,
            correlation_id=f"dlq_{message_body.get('task_id', 'unknown')}"
        )
        
        logger.warning(f"消息已发送到死信队列: {message_body.get('task_id', 'unknown')}")
        
    except Exception as dlq_error:
        logger.error(f"发送到死信队列失败: {dlq_error}")
        # 死信队列发送失败，记录到日志
        logger.error(f"原始失败消息: {message_body}, 错误: {error}")


async def on_message_template(
    rabbitmq_manager: RabbitMQManager,
    message: aio_pika.IncomingMessage,
):
    # 延迟ACK，确保prefetch_count生效
    async with message.process(requeue=True):
        try:
            # 解析消息内容
            message_body = json.loads(message.body.decode("utf-8"))
            files = message_body.get("files")
            template_type = message_body.get("template_type")
            task_id = message_body.get("task_id")
            logger.info(f"Received template message: {message_body}")

            # 将长耗时任务移到线程中，避免阻塞事件循环
            def _run_extract_template():
                import asyncio as _asyncio
                return _asyncio.run(extract_template(extraction_config, task_id, files, template_type))
            
            response = await asyncio.to_thread(_run_extract_template)
            response_json = response.model_dump_json()

            # 任务完成后立即发布
            await rabbitmq_manager.publish_message(
                message_body=response_json,
                routing_key=settings.PUB_TEMPLATE_ROUTING_KEY,
                correlation_id=message.correlation_id
            )

        except Exception as e:
            logger.error(f"处理模板消息时出错: {e}")
            await send_to_dead_letter_queue(rabbitmq_manager, message_body if 'message_body' in locals() else {}, e, "template")

async def on_message_invention(
    rabbitmq_manager: RabbitMQManager,
    message: aio_pika.IncomingMessage,
):
    # 延迟ACK，确保prefetch_count生效
    async with message.process(requeue=True):
        try:
            # 解析消息内容
            message_body = json.loads(message.body.decode("utf-8"))
            logger.info(f"Received invention message: {message_body}")
            files = message_body.get("files")
            task_id = message_body.get("task_id")

            # 将长耗时任务移到线程中，避免阻塞事件循环
            def _run_extract_invention():
                import asyncio as _asyncio
                return _asyncio.run(extract_knowledge_and_innovation(extraction_config, task_id, files))
            
            response = await asyncio.to_thread(_run_extract_invention)
            response_json = response.model_dump_json()

            # 任务完成后立即发布
            await rabbitmq_manager.publish_message(
                message_body=response_json,
                routing_key=settings.PUB_INVENTION_ROUTING_KEY,
                correlation_id=message.correlation_id
            )

        except Exception as e:
            logger.error(f"处理发明点消息时出错: {e}")
            await send_to_dead_letter_queue(rabbitmq_manager, message_body if 'message_body' in locals() else {}, e, "invention")

async def main():
    # 连接 RabbitMQ
    rabbitmq_manager = RabbitMQManager(settings)
    try:
        await rabbitmq_manager.connect()
        # 读取消费者角色：both | invention | template
        consumer_role = settings.CONSUMER_ROLE.lower()
        if consumer_role not in ("both", "invention", "template"):
            consumer_role = "both"

        # 根据角色只注册对应的消费者，便于多进程部署
        if consumer_role in ("both", "invention"):
            await rabbitmq_manager.register_consumer(
                'invention', lambda msg: on_message_invention(rabbitmq_manager, msg)
            )
        if consumer_role in ("both", "template"):
            await rabbitmq_manager.register_consumer(
                'template', lambda msg: on_message_template(rabbitmq_manager, msg)
            )

        logger.info(f"Message consumption started (role={consumer_role})")
        await asyncio.Future()
    except KeyboardInterrupt:
        logger.info("Shutting down gracefully...")
    except Exception as e:
        logger.error(f"Unexpected error in main: {e}")
    finally:
        await rabbitmq_manager.close()



if __name__ == "__main__":
    asyncio.run(main())