"""하이브리드 검색 API — FTS + 트리그램 + 벡터""" from typing import Annotated from fastapi import APIRouter, Depends, Query from pydantic import BaseModel from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession from ai.client import AIClient from core.auth import get_current_user from core.database import get_session from models.user import User router = APIRouter() # 가중치 (초기값, 튜닝 가능) W_FTS = 0.4 W_TRGM = 0.2 W_VECTOR = 0.4 class SearchResult(BaseModel): id: int title: str | None ai_domain: str | None ai_summary: str | None file_format: str score: float snippet: str | None class SearchResponse(BaseModel): results: list[SearchResult] total: int query: str mode: str @router.get("/", response_model=SearchResponse) async def search( q: str, user: Annotated[User, Depends(get_current_user)], session: Annotated[AsyncSession, Depends(get_session)], mode: str = Query("hybrid", pattern="^(fts|trgm|vector|hybrid)$"), limit: int = Query(20, ge=1, le=100), ): """문서 검색 mode: - fts: PostgreSQL 전문검색 (GIN 인덱스) - trgm: 트리그램 부분매칭 (한국어 지원) - vector: 벡터 유사도 검색 (의미검색) - hybrid: FTS + 트리그램 + 벡터 결합 (기본) """ if mode == "fts": results = await _search_fts(session, q, limit) elif mode == "trgm": results = await _search_trgm(session, q, limit) elif mode == "vector": results = await _search_vector(session, q, limit) else: results = await _search_hybrid(session, q, limit) return SearchResponse( results=results, total=len(results), query=q, mode=mode, ) async def _search_fts(session: AsyncSession, query: str, limit: int) -> list[SearchResult]: """PostgreSQL 전문검색 (GIN 인덱스)""" # simple 설정으로 한국어 토큰화 없이 공백 기반 분리 result = await session.execute( text(""" SELECT id, title, ai_domain, ai_summary, file_format, ts_rank( to_tsvector('simple', coalesce(title, '') || ' ' || coalesce(extracted_text, '')), plainto_tsquery('simple', :query) ) AS score, left(extracted_text, 200) AS snippet FROM documents WHERE to_tsvector('simple', coalesce(title, '') || ' ' || coalesce(extracted_text, '')) @@ plainto_tsquery('simple', :query) ORDER BY score DESC LIMIT :limit """), {"query": query, "limit": limit}, ) return [SearchResult(**row._mapping) for row in result] async def _search_trgm(session: AsyncSession, query: str, limit: int) -> list[SearchResult]: """트리그램 부분매칭 + ILIKE fallback (한국어 지원)""" # threshold 낮춰서 한국어 매칭 향상 await session.execute(text("SET pg_trgm.similarity_threshold = 0.1")) result = await session.execute( text(""" SELECT id, title, ai_domain, ai_summary, file_format, GREATEST( similarity(coalesce(title, '') || ' ' || coalesce(extracted_text, ''), :query), CASE WHEN (coalesce(title, '') || ' ' || coalesce(extracted_text, '')) ILIKE '%%' || :query || '%%' THEN 0.5 ELSE 0 END ) AS score, left(extracted_text, 200) AS snippet FROM documents WHERE (coalesce(title, '') || ' ' || coalesce(extracted_text, '')) %% :query OR (coalesce(title, '') || ' ' || coalesce(extracted_text, '')) ILIKE '%%' || :query || '%%' ORDER BY score DESC LIMIT :limit """), {"query": query, "limit": limit}, ) return [SearchResult(**row._mapping) for row in result] async def _search_vector(session: AsyncSession, query: str, limit: int) -> list[SearchResult]: """벡터 유사도 검색 (코사인 거리)""" client = AIClient() try: query_embedding = await client.embed(query) except Exception: return [] # GPU 서버 불가 시 빈 결과 finally: await client.close() # pgvector 코사인 거리 (0=동일, 2=반대) result = await session.execute( text(""" SELECT id, title, ai_domain, ai_summary, file_format, (1 - (embedding <=> :embedding::vector)) AS score, left(extracted_text, 200) AS snippet FROM documents WHERE embedding IS NOT NULL ORDER BY embedding <=> :embedding::vector LIMIT :limit """), {"embedding": str(query_embedding), "limit": limit}, ) return [SearchResult(**row._mapping) for row in result] async def _search_hybrid(session: AsyncSession, query: str, limit: int) -> list[SearchResult]: """하이브리드 검색 — FTS + ILIKE (안정적 한국어 지원)""" result = await session.execute( text(""" SELECT id, title, ai_domain, ai_summary, file_format, GREATEST( coalesce(ts_rank( to_tsvector('simple', coalesce(title, '') || ' ' || coalesce(extracted_text, '')), plainto_tsquery('simple', :query) ), 0), CASE WHEN (coalesce(title, '') || ' ' || coalesce(extracted_text, '')) ILIKE '%%' || :query || '%%' THEN 0.5 ELSE 0 END ) AS score, left(extracted_text, 200) AS snippet FROM documents WHERE to_tsvector('simple', coalesce(title, '') || ' ' || coalesce(extracted_text, '')) @@ plainto_tsquery('simple', :query) OR (coalesce(title, '') || ' ' || coalesce(extracted_text, '')) ILIKE '%%' || :query || '%%' ORDER BY score DESC LIMIT :limit """), {"query": query, "limit": limit}, ) return [SearchResult(**row._mapping) for row in result]