""" 语义搜索服务 支持 PostgreSQL（向量+全文混合）和 SQLite（LIKE 关键词搜索） """ from __future__ import annotations import json import logging import time from typing import List, Optional from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession from app.config import settings from app.database import IS_SQLITE from app.schemas.search import SearchRequest, SearchResult, SearchResponse logger = logging.getLogger(__name__) class SearchService: """语义搜索服务""" def __init__(self, db: AsyncSession): self.db = db async def search(self, request: SearchRequest) -> SearchResponse: """ 执行搜索。 - PostgreSQL: 向量搜索 + 中文全文搜索混合排序 - SQLite: LIKE 关键词搜索（无向量能力） """ start_time = time.time() if IS_SQLITE: return await self._search_sqlite(request, start_time) else: return await self._search_postgres(request, start_time) # ──────────────────────────── SQLite 搜索 ──────────────────────────── async def _search_sqlite(self, request: SearchRequest, start_time: float) -> SearchResponse: """SQLite 模式：使用 LIKE 关键词搜索""" # LLM 查询扩展（可选） expanded_queries = [request.query] try: from app.services.llm_service import LLMService expanded_queries = await LLMService.expand_query(request.query) except Exception as exc: logger.debug("查询扩展跳过: %s", exc) results: list[SearchResult] = [] seen_chunk_ids = set() for q in expanded_queries: like_pattern = f"%{q}%" sql = text(""" SELECT kc.id AS chunk_id, kp.id AS page_id, kp.title AS page_title, kc.content, kp.course_name, kp.teacher_name, kp.live_date FROM knowledge_chunks kc JOIN knowledge_pages kp ON kc.page_id = kp.id WHERE kc.content LIKE :pattern AND (:course IS NULL OR kp.course_name = :course) AND (:teacher IS NULL OR kp.teacher_name = :teacher) ORDER BY kc.id LIMIT :limit """) params = { "pattern": like_pattern, "course": request.course_name, "teacher": request.teacher_name, "limit": request.top_k, } result = await self.db.execute(sql, params) rows = result.fetchall() for row in rows: if row.chunk_id not in seen_chunk_ids: seen_chunk_ids.add(row.chunk_id) # 简单的相关性评分：关键词出现次数 score = min(1.0, row.content.lower().count(q.lower()) * 0.2 + 0.5) results.append(SearchResult( chunk_id=row.chunk_id, page_id=row.page_id, page_title=row.page_title, content=row.content, score=round(score, 4), course_name=row.course_name, teacher_name=row.teacher_name, live_date=row.live_date, highlight=None, )) results.sort(key=lambda x: x.score, reverse=True) results = results[:request.top_k] elapsed_ms = (time.time() - start_time) * 1000 return SearchResponse( query=request.query, total=len(results), results=results, elapsed_ms=round(elapsed_ms, 2), ) # ──────────────────────────── PostgreSQL 搜索 ──────────────────────────── async def _search_postgres(self, request: SearchRequest, start_time: float) -> SearchResponse: """PostgreSQL 模式：向量搜索 + 中文全文搜索混合""" # LLM 查询扩展 expanded_queries = [request.query] if request.use_fulltext: try: from app.services.llm_service import LLMService expanded_queries = await LLMService.expand_query(request.query) except Exception as exc: logger.debug("查询扩展跳过: %s", exc) # 生成查询向量 from app.services.embedding_service import EmbeddingService all_query_embeddings = await EmbeddingService.embed_batch(expanded_queries) # 构建过滤条件 where_clauses = [] params: dict = {} if request.course_name: where_clauses.append("kp.course_name = :course_name") params["course_name"] = request.course_name if request.teacher_name: where_clauses.append("kp.teacher_name = :teacher_name") params["teacher_name"] = request.teacher_name if request.live_date_from: where_clauses.append("kp.live_date >= :live_date_from") params["live_date_from"] = request.live_date_from if request.live_date_to: where_clauses.append("kp.live_date <= :live_date_to") params["live_date_to"] = request.live_date_to where_sql = "" if where_clauses: where_sql = "AND " + " AND ".join(where_clauses) # 向量搜索 vector_results = [] params["limit"] = request.top_k * 2 for q_embedding in all_query_embeddings: embedding_str = "[" + ",".join(str(x) for x in q_embedding) + "]" vector_sql = f""" SELECT kc.id AS chunk_id, kp.id AS page_id, kp.title AS page_title, kc.content, 1 - (kc.embedding <=> :embedding::vector) AS vector_score, 0.0 AS text_score, kp.course_name, kp.teacher_name, kp.live_date FROM knowledge_chunks kc JOIN knowledge_pages kp ON kc.page_id = kp.id WHERE kc.embedding IS NOT NULL {where_sql} ORDER BY kc.embedding <=> :embedding::vector LIMIT :limit """ v_params = {**params, "embedding": embedding_str} result = await self.db.execute(text(vector_sql), v_params) vector_results.extend(result.fetchall()) # 全文搜索 text_results = [] if request.use_fulltext: for q in expanded_queries: clean_query = q.replace("'", "''") fulltext_sql = f""" SELECT kc.id AS chunk_id, kp.id AS page_id, kp.title AS page_title, kc.content, 0.0 AS vector_score, ts_rank_cd(kc.search_vector, to_tsquery('chinese_zh', :query)) AS text_score, kp.course_name, kp.teacher_name, kp.live_date, ts_headline('chinese_zh', kc.content, to_tsquery('chinese_zh', :query), 'MaxWords=50, MinWords=20, ShortWord=2') AS highlight FROM knowledge_chunks kc JOIN knowledge_pages kp ON kc.page_id = kp.id WHERE kc.search_vector @@ to_tsquery('chinese_zh', :query) {where_sql} ORDER BY text_score DESC LIMIT :limit """ ft_params = {**params, "query": clean_query} ft_result = await self.db.execute(text(fulltext_sql), ft_params) text_results.extend(ft_result.fetchall()) # 合并结果 merged: dict[int, dict] = {} for row in vector_results: chunk_id = row.chunk_id if chunk_id not in merged: merged[chunk_id] = { "chunk_id": chunk_id, "page_id": row.page_id, "page_title": row.page_title, "content": row.content, "vector_score": row.vector_score, "text_score": 0.0, "course_name": row.course_name, "teacher_name": row.teacher_name, "live_date": row.live_date, "highlight": None, } else: merged[chunk_id]["vector_score"] = max( merged[chunk_id]["vector_score"], row.vector_score ) for row in text_results: chunk_id = row.chunk_id if chunk_id not in merged: merged[chunk_id] = { "chunk_id": chunk_id, "page_id": row.page_id, "page_title": row.page_title, "content": row.content, "vector_score": 0.0, "text_score": row.text_score or 0.0, "course_name": row.course_name, "teacher_name": row.teacher_name, "live_date": row.live_date, "highlight": row.highlight, } else: merged[chunk_id]["text_score"] = max( merged[chunk_id]["text_score"], row.text_score or 0.0 ) if row.highlight: merged[chunk_id]["highlight"] = row.highlight # 计算综合得分 max_text_score = max( (r["text_score"] for r in merged.values()), default=1.0 ) or 1.0 scored_results = [] for item in merged.values(): normalized_text = item["text_score"] / max_text_score if max_text_score > 0 else 0 combined_score = 0.7 * item["vector_score"] + 0.3 * normalized_text if combined_score >= request.threshold: scored_results.append( SearchResult( chunk_id=item["chunk_id"], page_id=item["page_id"], page_title=item["page_title"], content=item["content"], score=round(combined_score, 4), course_name=item["course_name"], teacher_name=item["teacher_name"], live_date=item["live_date"], highlight=item["highlight"], ) ) scored_results.sort(key=lambda x: x.score, reverse=True) scored_results = scored_results[:request.top_k] elapsed_ms = (time.time() - start_time) * 1000 return SearchResponse( query=request.query, total=len(scored_results), results=scored_results, elapsed_ms=round(elapsed_ms, 2), )