"""하이브리드 검색 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]: """트리그램 부분매칭 (한국어 지원)""" result = await session.execute( text(""" SELECT id, title, ai_domain, ai_summary, file_format, similarity( coalesce(title, '') || ' ' || coalesce(extracted_text, ''), :query ) AS score, left(extracted_text, 200) AS snippet FROM documents WHERE (coalesce(title, '') || ' ' || coalesce(extracted_text, '')) %% :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 + 트리그램 + 벡터 가중 합산""" # 벡터 임베딩 생성 (실패 시 FTS+트리그램만) query_embedding = None try: client = AIClient() query_embedding = await client.embed(query) await client.close() except Exception: pass vector_clause = "" vector_score = "0" params = {"query": query, "limit": limit, "w_fts": W_FTS, "w_trgm": W_TRGM, "w_vector": W_VECTOR} if query_embedding: vector_clause = "LEFT JOIN LATERAL (SELECT 1 - (d.embedding <=> :embedding::vector) AS vscore) v ON true" vector_score = "coalesce(v.vscore, 0)" params["embedding"] = str(query_embedding) else: # 벡터 없으면 FTS+트리그램만 사용 params["w_fts"] = 0.6 params["w_trgm"] = 0.4 params["w_vector"] = 0.0 result = await session.execute( text(f""" SELECT d.id, d.title, d.ai_domain, d.ai_summary, d.file_format, ( :w_fts * coalesce(ts_rank( to_tsvector('simple', coalesce(d.title, '') || ' ' || coalesce(d.extracted_text, '')), plainto_tsquery('simple', :query) ), 0) + :w_trgm * coalesce(similarity( coalesce(d.title, '') || ' ' || coalesce(d.extracted_text, ''), :query ), 0) + :w_vector * {vector_score} ) AS score, left(d.extracted_text, 200) AS snippet FROM documents d {vector_clause} WHERE coalesce(d.extracted_text, '') != '' ORDER BY score DESC LIMIT :limit """), params, ) return [SearchResult(**row._mapping) for row in result]