course-agent-od/backend/scripts/vectorizer_service.py

#!/usr/bin/env python3
"""
向量化服务 - 独立处理文档向量化
"""

import os
import sys
import time
import logging
from pathlib import Path
from datetime import datetime, timedelta

# 添加src目录到Python路径
sys.path.append(str(Path(__file__).parent.parent))

from src.core.database import SessionLocal, get_database_info
from src.models.document import Document
from src.services.document_service import DocumentService
from src.core.config import get_settings

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('vectorizer_service.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 减少不必要的日志
logging.getLogger("pypdf").setLevel(logging.ERROR)
logging.getLogger("PIL").setLevel(logging.ERROR)
logging.getLogger("chromadb").setLevel(logging.WARNING)

class VectorizerService:
    """向量化服务"""
    
    def __init__(self):
        self.settings = get_settings()
        self.batch_size = 5  # 每次处理的文档数量
        self.check_interval = 60  # 检查间隔（秒）
        self.max_retries = 3  # 最大重试次数
        self.retry_delay = 30  # 重试延迟（秒）
        
    def get_unprocessed_documents(self, session):
        """获取未处理的文档"""
        try:
            # 查找未处理的文档
            documents = session.query(Document).filter(
                Document.is_processed == False,
                Document.source_type == "knowledge_base"
            ).order_by(Document.id).limit(self.batch_size).all()
            
            # 过滤掉文件不存在的文档
            valid_documents = []
            for doc in documents:
                if Path(doc.file_path).exists():
                    valid_documents.append(doc)
                else:
                    logger.warning(f"文档文件不存在: {doc.file_path}")
                    # 标记为已处理以避免重复检查
                    doc.is_processed = True
                    session.commit()
            
            return valid_documents
        except Exception as e:
            logger.error(f"获取未处理文档失败: {e}")
            return []
    
    def process_document(self, session, document_service, document):
        """处理单个文档"""
        logger.info(f"开始处理文档: {document.file_path} (ID: {document.id})")
        
        for retry in range(self.max_retries):
            try:
                # 使用异步处理
                import asyncio
                
                # 创建新的事件循环
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                
                # 处理文档
                success = loop.run_until_complete(
                    document_service.process_document(document.id)
                )
                
                loop.close()
                
                if success:
                    logger.info(f"✅ 文档处理成功: {document.file_path}")
                    return True
                else:
                    logger.warning(f"❌ 文档处理失败: {document.file_path}")
                    if retry < self.max_retries - 1:
                        logger.info(f"等待 {self.retry_delay} 秒后重试...")
                        time.sleep(self.retry_delay)
                    
            except Exception as e:
                logger.error(f"文档处理异常 (尝试 {retry + 1}/{self.max_retries}): {e}")
                if retry < self.max_retries - 1:
                    logger.info(f"等待 {self.retry_delay} 秒后重试...")
                    time.sleep(self.retry_delay)
        
        logger.error(f"文档处理最终失败: {document.file_path}")
        return False
    
    def run(self):
        """运行向量化服务"""
        logger.info("=" * 60)
        logger.info("向量化服务启动")
        logger.info("=" * 60)
        
        # 显示数据库信息
        db_info = get_database_info()
        logger.info(f"数据库: {db_info.get('type', 'Unknown')}")
        logger.info(f"批量大小: {self.batch_size}")
        logger.info(f"检查间隔: {self.check_interval}秒")
        
        session = SessionLocal()
        document_service = DocumentService(session)
        
        try:
            while True:
                try:
                    # 获取未处理的文档
                    documents = self.get_unprocessed_documents(session)
                    
                    if documents:
                        logger.info(f"找到 {len(documents)} 个待处理文档")
                        
                        success_count = 0
                        for doc in documents:
                            if self.process_document(session, document_service, doc):
                                success_count += 1
                        
                        logger.info(f"批量处理完成: {success_count}/{len(documents)} 成功")
                    else:
                        logger.info("没有待处理的文档，等待中...")
                    
                    # 等待下一次检查
                    time.sleep(self.check_interval)
                    
                except KeyboardInterrupt:
                    logger.info("收到中断信号，停止服务")
                    break
                except Exception as e:
                    logger.error(f"服务运行异常: {e}")
                    logger.info(f"等待 {self.check_interval} 秒后继续...")
                    time.sleep(self.check_interval)
                    
        finally:
            session.close()
            logger.info("向量化服务停止")

def main():
    """主函数"""
    # 检查环境
    if not os.path.exists("vectorizer_service.log"):
        with open("vectorizer_service.log", "w") as f:
            f.write(f"向量化服务日志 - 启动时间: {datetime.now()}\n")
    
    # 运行服务
    service = VectorizerService()
    service.run()

if __name__ == "__main__":
    main()