""" 文档搜索相关 API """ from fastapi import APIRouter, Depends, HTTPException, Query, BackgroundTasks from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, or_ from typing import Optional, List from pathlib import Path import logging from app.core.database import get_db from app.core.deps import get_current_user from app.models.user import User from app.models.project import Project, ProjectMember from app.services.search_service import search_service from app.services.storage import storage_service from app.schemas.response import success_response router = APIRouter() logger = logging.getLogger(__name__) @router.get("/documents", response_model=dict) async def search_documents( keyword: str = Query(..., min_length=1, description="搜索关键词"), project_id: Optional[int] = Query(None, description="限制在指定项目中搜索"), current_user: User = Depends(get_current_user), db: AsyncSession = Depends(get_db) ): """ 文档搜索 (混合模式：Whoosh 全文检索 + 数据库项目搜索 + 文件系统文件名搜索 fallback) """ try: if not keyword: return success_response(data=[]) # 1. 确定搜索范围 (项目ID列表) allowed_project_ids = [] if project_id: # 检查指定项目的访问权限 result = await db.execute(select(Project).where(Project.id == project_id)) project = result.scalar_one_or_none() if not project: raise HTTPException(status_code=404, detail="项目不存在") # 检查权限 if project.owner_id != current_user.id and project.is_public != 1: member_result = await db.execute( select(ProjectMember).where( ProjectMember.project_id == project_id, ProjectMember.user_id == current_user.id ) ) if not member_result.scalar_one_or_none(): raise HTTPException(status_code=403, detail="无权访问该项目") allowed_project_ids.append(str(project_id)) else: # 获取所有可访问的项目 # 1. 用户创建的项目 owned_result = await db.execute( select(Project.id).where(Project.owner_id == current_user.id, Project.status == 1) ) allowed_project_ids.extend([str(pid) for pid in owned_result.scalars().all()]) # 2. 用户参与的项目 member_result = await db.execute( select(ProjectMember.project_id) .join(Project, ProjectMember.project_id == Project.id) .where( ProjectMember.user_id == current_user.id, Project.status == 1 ) ) allowed_project_ids.extend([str(pid) for pid in member_result.scalars().all()]) # 去重 allowed_project_ids = list(set(allowed_project_ids)) if not allowed_project_ids: return success_response(data=[]) # 2. 执行搜索 search_results = [] # A. 数据库项目搜索 (仅当未指定 project_id 时，或者需要搜项目本身) # 如果前端指定了 project_id，通常是在项目内搜文件，不需要搜项目本身 if not project_id: projects_query = select(Project).where( Project.id.in_(allowed_project_ids), or_( Project.name.ilike(f"%{keyword}%"), Project.description.ilike(f"%{keyword}%") ) ) project_res = await db.execute(projects_query) matched_projects = project_res.scalars().all() for proj in matched_projects: search_results.append({ "type": "project", "project_id": proj.id, "project_name": proj.name, "project_description": proj.description or "", "match_type": "项目名称/描述", }) # B. Whoosh 全文检索 whoosh_results = [] try: if project_id: whoosh_results = await search_service.search(keyword, str(project_id)) else: # 全局搜索 whoosh_results = await search_service.search(keyword, limit=50) # 过滤权限 whoosh_results = [r for r in whoosh_results if str(r['project_id']) in allowed_project_ids] except Exception as e: logger.warning(f"Whoosh search failed: {e}") pass # 获取 Whoosh 结果涉及的项目 ID whoosh_project_ids = set(res['project_id'] for res in whoosh_results if res.get('project_id')) # 查询项目名称映射 project_name_map = {} if whoosh_project_ids: p_res = await db.execute(select(Project.id, Project.name).where(Project.id.in_(whoosh_project_ids))) for pid, pname in p_res.all(): project_name_map[str(pid)] = pname # 添加 Whoosh 结果 for res in whoosh_results: pid_str = str(res['project_id']) search_results.append({ "type": "file", "project_id": res['project_id'], "project_name": project_name_map.get(pid_str, "未知项目"), "file_path": res['path'], "file_name": res['title'], "highlights": res.get('highlights'), "match_type": "全文检索" }) # C. 文件系统文件名搜索 (Fallback / Complementary) # 为了保证未索引的文件也能通过文件名搜到 # 获取需要扫描的项目 projects_to_scan = [] if project_id: # 单项目扫描 res = await db.execute(select(Project).where(Project.id == project_id)) p = res.scalar_one_or_none() if p: projects_to_scan = [p] elif len(search_results) < 20: # 全局扫描：仅当结果较少时才进行全盘扫描，避免性能问题 # 这是一个简单的启发式策略 res = await db.execute(select(Project).where(Project.id.in_(allowed_project_ids))) projects_to_scan = res.scalars().all() # 已存在的文件路径集合 (用于去重) existing_paths = set() for res in search_results: if res.get('type') == 'file': # 统一 key 格式 existing_paths.add(f"{res['project_id']}:{res['file_path']}") keyword_lower = keyword.lower() for project in projects_to_scan: try: project_path = storage_service.get_secure_path(project.storage_key) if not project_path.exists(): continue # 查找文件名匹配 md_files = list(project_path.rglob("*.md")) pdf_files = list(project_path.rglob("*.pdf")) for file_path in md_files + pdf_files: if "_assets" in file_path.parts: continue if keyword_lower in file_path.name.lower(): rel_path = str(file_path.relative_to(project_path)) unique_key = f"{project.id}:{rel_path}" if unique_key not in existing_paths: search_results.append({ "type": "file", "project_id": project.id, "project_name": project.name, "file_path": rel_path, "file_name": file_path.name, "match_type": "文件名匹配" }) existing_paths.add(unique_key) except Exception: continue return success_response(data=search_results[:100]) except Exception as e: logger.error(f"Search API error: {e}") return success_response(data=[], message="搜索服务暂时不可用") async def rebuild_index_task(db: AsyncSession): """后台任务：重建索引""" logger.info("Starting index rebuild...") try: # 获取所有项目 result = await db.execute(select(Project).where(Project.status == 1)) projects = result.scalars().all() documents = [] for project in projects: try: # 遍历项目文件 project_root = storage_service.get_secure_path(project.storage_key) if not project_root.exists(): continue # 查找所有 .md 文件 md_files = list(project_root.rglob("*.md")) for file_path in md_files: if "_assets" in file_path.parts: continue try: content = await storage_service.read_file(file_path) relative_path = str(file_path.relative_to(project_root)) documents.append({ "project_id": project.id, "path": relative_path, "title": file_path.stem, "content": content }) except Exception: continue except Exception as e: logger.error(f"Error processing project {project.id}: {e}") continue # 批量写入索引 import asyncio loop = asyncio.get_running_loop() await loop.run_in_executor(None, search_service.rebuild_index_sync, documents) logger.info(f"Index rebuild completed. Indexed {len(documents)} documents.") except Exception as e: logger.error(f"Index rebuild failed: {e}") @router.post("/rebuild-index", response_model=dict) async def rebuild_index( background_tasks: BackgroundTasks, current_user: User = Depends(get_current_user), db: AsyncSession = Depends(get_db) ): """ 重建搜索索引 (仅限超级管理员) """ if not current_user.is_superuser: raise HTTPException(status_code=403, detail="权限不足") background_tasks.add_task(rebuild_index_task, db) return success_response(message="索引重建任务已启动")