From dbb5bc49b02a86d3780837c82941254dd90bda56 Mon Sep 17 00:00:00 2001 From: chentaiyan-fullive Date: Thu, 27 Aug 2026 11:57:14 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BC=98=E5=8C=96=E5=94=A4=E9=86=92?= =?UTF-8?q?=E7=BC=93=E5=AD=98=E4=B8=8E=E5=90=91=E9=87=8F=E7=9B=B8=E4=BC=BC?= =?UTF-8?q?=E5=BA=A6=E8=AE=A1=E7=AE=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/cache/sleep_stage_cache.py | 55 ++++-- app/services/retrieval.py | 339 ++++++++++++++++++++++++++------ tests/test_retrieval.py | 119 ++++++++--- tests/test_sleep_stage_cache.py | 128 +++++++++++- 4 files changed, 540 insertions(+), 101 deletions(-) diff --git a/app/cache/sleep_stage_cache.py b/app/cache/sleep_stage_cache.py index 9b2954e..e1df956 100644 --- a/app/cache/sleep_stage_cache.py +++ b/app/cache/sleep_stage_cache.py @@ -6,6 +6,7 @@ from __future__ import annotations +import asyncio import hashlib import json from collections.abc import Awaitable, Callable @@ -15,7 +16,8 @@ from app.core.config import Settings -SLEEP_STAGES = ("放松", "入睡", "守护", "清醒") +SLEEP_STAGES = ("放松", "入睡", "守护", "唤醒") +_CACHE_CLEANUP_STAGES = (*SLEEP_STAGES, "清醒") SLEEP_STAGE_INDEX_KEY_PREFIX = "sleep_stage_v2_index:" SLEEP_STAGE_DOC_KEY_PREFIX = "sleep_stage_v2_doc:" # 兼容旧测试/调用名 @@ -92,6 +94,7 @@ def __init__(self, redis: RedisLike, *, ttl_sec: int = _WEEK_SECONDS) -> None: raise ValueError("ttl_sec must be >= 1") self._redis = redis self._ttl_sec = ttl_sec + self._mutation_lock = asyncio.Lock() async def get(self, stages: list[str]) -> list[dict[str, Any]] | None: """全部阶段索引命中且文档齐全才返回；否则 None。""" @@ -158,18 +161,50 @@ async def set_stage(self, stage: str, docs: list[dict[str, Any]]) -> None: len(urls), ) + async def get_or_load( + self, + stages: list[str], + loader: StageLoader, + ) -> list[dict[str, Any]]: + """读取缓存；未命中时仅加载缺失阶段，并合并同进程并发 miss。""" + normalized = _normalize_stages(stages) + cached = await self.get(normalized) + if cached is not None: + return cached + + async with self._mutation_lock: + cached = await self.get(normalized) + if cached is not None: + return cached + + for stage in normalized: + if await self.get([stage]) is not None: + continue + docs = await loader(stage) + await self.set_stage(stage, docs) + + loaded = await self.get(normalized) + if loaded is None: + raise RuntimeError("sleep stage cache load completed without a readable value") + return loaded + async def warm(self, loader: StageLoader) -> None: """清空后按四个阶段从数据源重建索引与文档。""" - await self.clear_all() - for stage in SLEEP_STAGES: - docs = await loader(stage) - await self.set_stage(stage, docs) + async with self._mutation_lock: + await self._clear_all_unlocked() + for stage in SLEEP_STAGES: + docs = await loader(stage) + await self.set_stage(stage, docs) logger.info("睡眠阶段候选缓存预热完成，stages={}", list(SLEEP_STAGES)) async def clear_all(self) -> None: """删除四个阶段索引及其引用的全部文档。""" + async with self._mutation_lock: + await self._clear_all_unlocked() + + async def _clear_all_unlocked(self) -> None: urls: set[str] = set() - index_keys = [build_sleep_stage_index_key(stage) for stage in SLEEP_STAGES] + index_keys = [build_sleep_stage_index_key(stage) for stage in _CACHE_CLEANUP_STAGES] for key in index_keys: raw = await self._redis.get(key) if raw is None: @@ -180,7 +215,7 @@ async def clear_all(self) -> None: legacy_urls: set[str] = set() legacy_index_keys = [ - f"{_LEGACY_SLEEP_STAGE_INDEX_KEY_PREFIX}{stage}" for stage in SLEEP_STAGES + f"{_LEGACY_SLEEP_STAGE_INDEX_KEY_PREFIX}{stage}" for stage in _CACHE_CLEANUP_STAGES ] for key in legacy_index_keys: raw = await self._redis.get(key) @@ -191,11 +226,9 @@ async def clear_all(self) -> None: legacy_urls.update(str(item) for item in items if item) doc_keys = [build_sleep_stage_doc_key(url) for url in urls] - legacy_doc_keys = [ - _build_legacy_sleep_stage_doc_key(url) for url in legacy_urls - ] + legacy_doc_keys = [_build_legacy_sleep_stage_doc_key(url) for url in legacy_urls] legacy_candidate_keys = [ - f"{_LEGACY_SLEEP_STAGE_CANDIDATE_KEY_PREFIX}{stage}" for stage in SLEEP_STAGES + f"{_LEGACY_SLEEP_STAGE_CANDIDATE_KEY_PREFIX}{stage}" for stage in _CACHE_CLEANUP_STAGES ] to_delete = [ *index_keys, diff --git a/app/services/retrieval.py b/app/services/retrieval.py index 0cb0431..cbf3c12 100644 --- a/app/services/retrieval.py +++ b/app/services/retrieval.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any +import numpy as np from loguru import logger from app.core.config import Settings @@ -95,6 +96,78 @@ class ExtractedQueryTags: disliked_tags: list[str] +@dataclass(frozen=True) +class VectorSimilaritySnapshot: + """请求向量与候选标签向量的一次性批量余弦结果。""" + + scores: np.ndarray + tag_columns: dict[str, int] + + @classmethod + def build( + cls, + request_vectors: list[list[float]], + dictionary_vectors: DictionaryVectors, + ) -> VectorSimilaritySnapshot: + dictionary_items = [ + (tag_id, vector) for tag_id, vector in dictionary_vectors.items() if vector + ] + tag_columns = {tag_id: column for column, (tag_id, _vector) in enumerate(dictionary_items)} + if not request_vectors or not dictionary_items: + return cls( + scores=np.zeros( + (len(request_vectors), len(dictionary_items)), + dtype=np.float64, + ), + tag_columns=tag_columns, + ) + scores = _cosine_similarity_matrix( + request_vectors, + [vector for _tag_id, vector in dictionary_items], + ) + return cls(scores=scores, tag_columns=tag_columns) + + def count_matches( + self, + tag_ids: list[str], + request_tags: list[str], + *, + row_offset: int, + threshold: float, + ) -> int: + columns = [self.tag_columns[tag_id] for tag_id in tag_ids if tag_id in self.tag_columns] + if not columns or not request_tags: + return 0 + end = row_offset + len(request_tags) + if row_offset < 0 or end > self.scores.shape[0]: + raise ValueError("similarity snapshot row range is invalid") + block = self.scores[row_offset:end, columns] + return sum( + request_tag not in MUTUALLY_EXCLUSIVE_CONTENT_TAGS + and bool(np.any(block[row] >= threshold)) + for row, request_tag in enumerate(request_tags) + ) + + def max_similarity( + self, + tag_ids: list[str], + request_tags: list[str], + *, + row_offset: int, + ) -> float: + columns = [self.tag_columns[tag_id] for tag_id in tag_ids if tag_id in self.tag_columns] + rows = [ + row_offset + row + for row, request_tag in enumerate(request_tags) + if request_tag not in MUTUALLY_EXCLUSIVE_CONTENT_TAGS + ] + if not columns or not rows: + return 0.0 + if rows[-1] >= self.scores.shape[0] or row_offset < 0: + raise ValueError("similarity snapshot row range is invalid") + return float(np.max(self.scores[np.ix_(rows, columns)])) + + class RetrievalService: """三维度检索：睡眠阶段 → 内容形态 → 厌恶剔除 → 粗排 → 精排。""" @@ -147,18 +220,27 @@ async def _search_tag_only(self, request: SearchAudioRequest) -> list[dict[str, candidates_raw, _normalize_color_noise_aliases(request.content_tags), ) - dictionary_vectors, content_tags, disliked_tags, content_vectors, dislike_vectors = ( - await self._prepare_content_admission_inputs( - candidates_raw, - request, - need_dictionary=need_dictionary, - ) + ( + dictionary_vectors, + content_tags, + disliked_tags, + content_vectors, + dislike_vectors, + ) = await self._prepare_content_admission_inputs( + candidates_raw, + request, + need_dictionary=need_dictionary, + ) + similarity_snapshot = VectorSimilaritySnapshot.build( + [*content_vectors, *dislike_vectors], + dictionary_vectors, ) admitted = await self._apply_content_admission( candidates_raw, content_tags, dictionary_vectors, request_vectors=content_vectors, + similarity_snapshot=similarity_snapshot, ) step2_ms = _elapsed_ms(step2_started) logger.info( @@ -178,6 +260,8 @@ async def _search_tag_only(self, request: SearchAudioRequest) -> list[dict[str, vector_disliked, dictionary_vectors, dislike_vectors=dislike_vectors, + similarity_snapshot=similarity_snapshot, + similarity_row_offset=len(content_vectors), ) filtered = _apply_voice_code_filter( filtered, @@ -296,6 +380,11 @@ async def _search_text_multi_route( await asyncio.gather(dict_task, encode_task, return_exceptions=True) raise + similarity_snapshot = VectorSimilaritySnapshot.build( + [*content_vectors, *dislike_vectors], + dictionary_vectors, + ) + tag_candidates: list[ScoredCandidate] = [] if content_tags: tag_candidates = await self._score_content_candidates( @@ -303,6 +392,7 @@ async def _search_text_multi_route( content_tags, dictionary_vectors, request_vectors=content_vectors, + similarity_snapshot=similarity_snapshot, ) merged = await self._merge_and_rank_text_candidates( @@ -313,6 +403,8 @@ async def _search_text_multi_route( voice_filter_tags=disliked_tags, dictionary_vectors=dictionary_vectors, dislike_vectors=dislike_vectors, + similarity_snapshot=similarity_snapshot, + similarity_row_offset=len(content_vectors), top_k=request.top_k, ) rank_ms = _elapsed_ms(rank_started) @@ -360,8 +452,7 @@ async def _fetch_step1_candidates(self, sleep_stage_tags: list[str]) -> list[dic ) return cached - candidates = await self._es_search.filter_by_sleep_stage(sleep_stage_tags) - await self._backfill_sleep_stage_cache() + candidates = await self._load_sleep_stage_candidates_on_miss(sleep_stage_tags) logger.info( "检索步骤1/4 睡眠阶段过滤：候选数={}，耗时={:.1f}毫秒", len(candidates), @@ -369,6 +460,22 @@ async def _fetch_step1_candidates(self, sleep_stage_tags: list[str]) -> list[dic ) return candidates + async def _load_sleep_stage_candidates_on_miss( + self, + sleep_stage_tags: list[str], + ) -> list[dict[str, Any]]: + """缓存 miss 时只回填请求阶段；缓存故障则直接回退 ES。""" + if self._sleep_stage_cache is None: + return await self._es_search.filter_by_sleep_stage(sleep_stage_tags) + try: + return await self._sleep_stage_cache.get_or_load( + sleep_stage_tags, + self._load_sleep_stage_candidates, + ) + except Exception as exc: + logger.warning("按需回填睡眠阶段候选缓存失败，回退 ES：{}", exc) + return await self._es_search.filter_by_sleep_stage(sleep_stage_tags) + async def _get_sleep_stage_cached( self, sleep_stage_tags: list[str], @@ -401,11 +508,7 @@ async def warm_query_tag_vectors(self) -> None: started = time.perf_counter() tags = await self._es_search.list_content_tag_vectors() labels = _unique_preserve_order( - [ - label - for tag in tags - if (label := str(tag.get("label", "")).strip()) - ] + [label for tag in tags if (label := str(tag.get("label", "")).strip())] ) if not labels: logger.info("查询标签向量缓存预热跳过：内容词典为空") @@ -604,6 +707,7 @@ async def _apply_content_admission( dictionary_vectors: DictionaryVectors, *, request_vectors: list[list[float]] | None = None, + similarity_snapshot: VectorSimilaritySnapshot | None = None, ) -> list[ScoredCandidate]: """步骤 2：无 content_tags 时跳过准入，保留睡眠阶段候选全集。""" if not content_tags: @@ -622,6 +726,10 @@ async def _apply_content_admission( request_vectors = ( await self._encode_texts(content_tags) if request_vectors is None else request_vectors ) + similarity_snapshot = similarity_snapshot or VectorSimilaritySnapshot.build( + request_vectors, + dictionary_vectors, + ) admitted: list[ScoredCandidate] = [] for doc in candidates: @@ -646,6 +754,7 @@ async def _apply_content_admission( content_tags, request_vectors, dictionary_vectors, + similarity_snapshot=similarity_snapshot, ) if vector_hits > 0: admitted.append( @@ -667,6 +776,7 @@ async def _score_content_candidates( dictionary_vectors: DictionaryVectors, *, request_vectors: list[list[float]] | None = None, + similarity_snapshot: VectorSimilaritySnapshot | None = None, ) -> list[ScoredCandidate]: """文本多路检索里的标签路：产出分数，不决定整体短路。""" if not candidates or not content_tags: @@ -675,6 +785,10 @@ async def _score_content_candidates( request_vectors = ( await self._encode_texts(content_tags) if request_vectors is None else request_vectors ) + similarity_snapshot = similarity_snapshot or VectorSimilaritySnapshot.build( + request_vectors, + dictionary_vectors, + ) scored: list[ScoredCandidate] = [] tag_count = max(len(content_tags), 1) for doc in candidates: @@ -688,6 +802,7 @@ async def _score_content_candidates( content_tags, request_vectors, dictionary_vectors, + similarity_snapshot=similarity_snapshot, ) match_count = len(exact_hits) if exact_hits else vector_hits if match_count <= 0: @@ -713,25 +828,25 @@ def _count_fuzzy_vector_matches( request_tags: list[str], request_vectors: list[list[float]], dictionary_vectors: DictionaryVectors, + *, + similarity_snapshot: VectorSimilaritySnapshot | None = None, + similarity_row_offset: int = 0, ) -> int: """使用请求级向量快照计分；互斥标签只允许精确命中。""" tag_ids = EsSearch.content_tag_ids(tags) if not tag_ids or not request_vectors: return 0 - threshold = self._settings.sim_threshold - matched = 0 - - for request_tag, req_vec in zip(request_tags, request_vectors, strict=True): - if request_tag in MUTUALLY_EXCLUSIVE_CONTENT_TAGS: - continue - for tid in tag_ids: - doc_vec = dictionary_vectors.get(tid) - if doc_vec and _cosine_similarity(req_vec, doc_vec) >= threshold: - matched += 1 - break - - return matched + snapshot = similarity_snapshot or VectorSimilaritySnapshot.build( + request_vectors, + dictionary_vectors, + ) + return snapshot.count_matches( + tag_ids, + request_tags, + row_offset=similarity_row_offset, + threshold=self._settings.sim_threshold, + ) async def _apply_dislike_filter( self, @@ -740,15 +855,19 @@ async def _apply_dislike_filter( dictionary_vectors: DictionaryVectors, *, dislike_vectors: list[list[float]] | None = None, + similarity_snapshot: VectorSimilaritySnapshot | None = None, + similarity_row_offset: int = 0, ) -> list[ScoredCandidate]: """步骤 3 前半：厌恶标签向量 vs 文档内容标签向量，余弦 ≥ SIM_THRESHOLD 则剔除。""" if not disliked_tags: return candidates dislike_vectors = ( - await self._encode_texts(disliked_tags) - if dislike_vectors is None - else dislike_vectors + await self._encode_texts(disliked_tags) if dislike_vectors is None else dislike_vectors + ) + similarity_snapshot = similarity_snapshot or VectorSimilaritySnapshot.build( + dislike_vectors, + dictionary_vectors, ) result: list[ScoredCandidate] = [] for candidate in candidates: @@ -758,6 +877,8 @@ async def _apply_dislike_filter( disliked_tags, dislike_vectors, dictionary_vectors, + similarity_snapshot=similarity_snapshot, + similarity_row_offset=similarity_row_offset, ) > 0 ): @@ -782,6 +903,8 @@ async def _merge_and_rank_text_candidates( dictionary_vectors: DictionaryVectors, top_k: int | None, dislike_vectors: list[list[float]] | None = None, + similarity_snapshot: VectorSimilaritySnapshot | None = None, + similarity_row_offset: int = 0, voice_filter_tags: list[str] | None = None, ) -> list[ScoredCandidate]: merged: dict[str, ScoredCandidate] = {} @@ -800,6 +923,10 @@ async def _merge_and_rank_text_candidates( if disliked_tags and dislike_vectors is None else (dislike_vectors or []) ) + similarity_snapshot = similarity_snapshot or VectorSimilaritySnapshot.build( + dislike_vectors, + dictionary_vectors, + ) ranked: list[ScoredCandidate] = [] for candidate in merged.values(): penalty = self._dislike_penalty( @@ -807,6 +934,8 @@ async def _merge_and_rank_text_candidates( disliked_tags, dislike_vectors, dictionary_vectors, + similarity_snapshot=similarity_snapshot, + similarity_row_offset=similarity_row_offset, ) if penalty >= 1.0: continue @@ -841,11 +970,18 @@ async def _extract_query_tags( exact_content_tags = _unique_preserve_order( [*color_intents, *_match_labels_from_text(positive_text, tag_vectors)] ) + fragment_vectors = ( + await self._encode_texts(negative_fragments) if negative_fragments else [] + ) + vector_labels, similarity_scores = _tag_vector_similarity_scores( + [query_vector, *fragment_vectors], + tag_vectors, + ) content_tags = list(exact_content_tags) content_tags.extend( - _similar_labels_from_vector( - query_vector, - tag_vectors, + _similar_labels_from_scores( + similarity_scores[0] if similarity_scores.shape[0] else np.zeros(0), + vector_labels, threshold=AUTO_TAG_SIM_THRESHOLD, exclude=set(content_tags), limit=AUTO_TAG_TOP_K - len(content_tags), @@ -858,13 +994,12 @@ async def _extract_query_tags( exact_disliked_tags = _match_labels_from_text(" ".join(negative_fragments), tag_vectors) disliked_tags = list(exact_disliked_tags) - if negative_fragments: - fragment_vectors = await self._encode_texts(negative_fragments) - for vector in fragment_vectors: + if fragment_vectors: + for row in range(1, len(fragment_vectors) + 1): disliked_tags.extend( - _similar_labels_from_vector( - vector, - tag_vectors, + _similar_labels_from_scores( + similarity_scores[row], + vector_labels, threshold=AUTO_DISLIKE_SIM_THRESHOLD, exclude=set(disliked_tags), limit=AUTO_TAG_TOP_K - len(disliked_tags), @@ -890,6 +1025,9 @@ def _dislike_penalty( disliked_tags: list[str], dislike_vectors: list[list[float]], dictionary_vectors: DictionaryVectors, + *, + similarity_snapshot: VectorSimilaritySnapshot | None = None, + similarity_row_offset: int = 0, ) -> float: if not disliked_tags: return 0.0 @@ -902,6 +1040,8 @@ def _dislike_penalty( disliked_tags, dislike_vectors, dictionary_vectors, + similarity_snapshot=similarity_snapshot, + similarity_row_offset=similarity_row_offset, ) if max_similarity >= self._settings.strong_dislike_sim_threshold: return 1.0 @@ -915,20 +1055,23 @@ def _max_fuzzy_vector_similarity( request_tags: list[str], request_vectors: list[list[float]], dictionary_vectors: DictionaryVectors, + *, + similarity_snapshot: VectorSimilaritySnapshot | None = None, + similarity_row_offset: int = 0, ) -> float: tag_ids = EsSearch.content_tag_ids(tags) if not tag_ids or not request_vectors: return 0.0 - max_similarity = 0.0 - for request_tag, req_vec in zip(request_tags, request_vectors, strict=True): - if request_tag in MUTUALLY_EXCLUSIVE_CONTENT_TAGS: - continue - for tag_id in tag_ids: - doc_vec = dictionary_vectors.get(tag_id) - if doc_vec: - max_similarity = max(max_similarity, _cosine_similarity(req_vec, doc_vec)) - return max_similarity + snapshot = similarity_snapshot or VectorSimilaritySnapshot.build( + request_vectors, + dictionary_vectors, + ) + return snapshot.max_similarity( + tag_ids, + request_tags, + row_offset=similarity_row_offset, + ) def _candidate_from_doc( self, @@ -1020,6 +1163,41 @@ def _cosine_similarity(a: list[float], b: list[float]) -> float: return dot / (norm_a * norm_b) +def _cosine_similarity_matrix( + request_vectors: list[list[float]], + dictionary_vectors: list[list[float]], +) -> np.ndarray: + """批量计算二维余弦矩阵，行对应请求向量、列对应词典向量。""" + try: + requests = np.asarray(request_vectors, dtype=np.float64) + dictionary = np.asarray(dictionary_vectors, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise ValueError("similarity vectors must form rectangular numeric matrices") from exc + if requests.ndim != 2 or dictionary.ndim != 2: + raise ValueError("similarity vectors must be two-dimensional") + if requests.shape[1] != dictionary.shape[1]: + raise ValueError( + "similarity vector dimension mismatch: " + f"request={requests.shape[1]}, dictionary={dictionary.shape[1]}" + ) + + request_norms = np.linalg.norm(requests, axis=1, keepdims=True) + dictionary_norms = np.linalg.norm(dictionary, axis=1, keepdims=True) + normalized_requests = np.divide( + requests, + request_norms, + out=np.zeros_like(requests), + where=request_norms != 0, + ) + normalized_dictionary = np.divide( + dictionary, + dictionary_norms, + out=np.zeros_like(dictionary), + where=dictionary_norms != 0, + ) + return normalized_requests @ normalized_dictionary.T + + def _candidate_key(source: dict[str, Any]) -> str: return str( source.get("_id") or source.get("id") or source.get("audio_url") or source.get("audio_name") @@ -1185,9 +1363,7 @@ def _normalize_to_dictionary_labels( ] dictionary = { - str(item["label"]).strip() - for item in tag_vectors - if str(item.get("label", "")).strip() + str(item["label"]).strip() for item in tag_vectors if str(item.get("label", "")).strip() } normalized: list[str] = [] for raw in raw_tags: @@ -1212,23 +1388,62 @@ def _similar_labels_from_vector( threshold: float, exclude: set[str], limit: int, +) -> list[str]: + labels, scores = _tag_vector_similarity_scores([query_vector], tag_vectors) + row = scores[0] if scores.shape[0] else np.zeros(0) + return _similar_labels_from_scores( + row, + labels, + threshold=threshold, + exclude=exclude, + limit=limit, + ) + + +def _tag_vector_similarity_scores( + query_vectors: list[list[float]], + tag_vectors: list[dict[str, Any]], +) -> tuple[list[str], np.ndarray]: + """过滤无效词典项并一次计算全部查询与标签的余弦分数。""" + eligible: list[tuple[str, list[float]]] = [] + for item in tag_vectors: + label = str(item.get("label", "")).strip() + vector = item.get("vector") + if not label or len(label) < MIN_AUTO_TAG_LABEL_LEN or not vector: + continue + eligible.append((label, vector)) + if not query_vectors or not eligible: + return ( + [label for label, _vector in eligible], + np.zeros((len(query_vectors), len(eligible)), dtype=np.float64), + ) + return ( + [label for label, _vector in eligible], + _cosine_similarity_matrix( + query_vectors, + [vector for _label, vector in eligible], + ), + ) + + +def _similar_labels_from_scores( + scores: np.ndarray, + labels: list[str], + *, + threshold: float, + exclude: set[str], + limit: int, ) -> list[str]: if limit <= 0: return [] + if scores.ndim != 1 or scores.shape[0] != len(labels): + raise ValueError("tag similarity scores and labels must have matching lengths") scored: list[tuple[float, str]] = [] - for item in tag_vectors: - label = str(item["label"]).strip() - vector = item.get("vector") - if ( - not label - or len(label) < MIN_AUTO_TAG_LABEL_LEN - or label in exclude - or not vector - ): + for score, label in zip(scores, labels, strict=True): + if label in exclude: continue - similarity = _cosine_similarity(query_vector, vector) - if similarity >= threshold: - scored.append((similarity, label)) + if score >= threshold: + scored.append((float(score), label)) scored.sort(key=lambda pair: pair[0], reverse=True) return _prefer_longer_labels([label for _, label in scored])[:limit] diff --git a/tests/test_retrieval.py b/tests/test_retrieval.py index 77c25ce..f152e9e 100644 --- a/tests/test_retrieval.py +++ b/tests/test_retrieval.py @@ -16,7 +16,10 @@ VOICE_MARKER, RetrievalService, ScoredCandidate, + VectorSimilaritySnapshot, _apply_voice_code_filter, + _cosine_similarity, + _cosine_similarity_matrix, _match_labels_from_text, _normalize_to_dictionary_labels, _strip_voice_mention_tags, @@ -28,6 +31,42 @@ _VECTOR_DIM = 512 +def test_vectorized_cosine_matches_scalar_reference() -> None: + requests = [[1.0, 2.0, 3.0], [0.0, 0.0, 0.0], [-1.0, 0.5, 2.0]] + dictionary = [[3.0, 2.0, 1.0], [0.0, 0.0, 0.0], [1.0, -2.0, 0.5]] + + actual = _cosine_similarity_matrix(requests, dictionary) + + for row, request in enumerate(requests): + for column, candidate in enumerate(dictionary): + assert actual[row, column] == pytest.approx( + _cosine_similarity(request, candidate), + abs=1e-12, + ) + + +def test_vectorized_cosine_rejects_dimension_mismatch() -> None: + with pytest.raises(ValueError, match="dimension mismatch"): + _cosine_similarity_matrix([[1.0, 0.0]], [[1.0, 0.0, 0.0]]) + + +def test_similarity_snapshot_preserves_exclusive_tag_semantics() -> None: + snapshot = VectorSimilaritySnapshot.build( + [[1.0, 0.0], [1.0, 0.0]], + {"rain": [1.0, 0.0]}, + ) + + assert ( + snapshot.count_matches( + ["rain"], + ["白噪音", "雨声"], + row_offset=0, + threshold=0.8, + ) + == 1 + ) + + def _tag_entries(prefix: str, labels: list[str]) -> list[dict[str, str]]: return [{"tag_id": f"{prefix}_{label}", "code": label, "name": label} for label in labels] @@ -171,6 +210,47 @@ async def test_search_uses_sleep_stage_cache_when_available() -> None: es_search.filter_by_sleep_stage.assert_not_called() +@pytest.mark.asyncio +async def test_search_cache_miss_uses_on_demand_stage_loader() -> None: + cached_doc = _audio_doc("唤醒音乐", sleep_stage=["唤醒"], content_form=["音乐"]) + sleep_cache = MagicMock() + sleep_cache.get = AsyncMock(return_value=None) + sleep_cache.get_or_load = AsyncMock(return_value=[cached_doc]) + service, es_search, _encoder = _build_service() + service._sleep_stage_cache = sleep_cache + es_search.parse_tags = EsSearch.parse_tags + + results = await service.search( + SearchAudioRequest(sleep_stage_tags=["唤醒"], content_tags=["音乐"], top_k=5) + ) + + assert [result["audio_name"] for result in results] == ["唤醒音乐"] + sleep_cache.get_or_load.assert_awaited_once_with( + ["唤醒"], + service._load_sleep_stage_candidates, + ) + es_search.filter_by_sleep_stage.assert_not_called() + + +@pytest.mark.asyncio +async def test_search_cache_loader_error_falls_back_to_es() -> None: + es_doc = _audio_doc("ES 唤醒音乐", sleep_stage=["唤醒"], content_form=["音乐"]) + sleep_cache = MagicMock() + sleep_cache.get = AsyncMock(return_value=None) + sleep_cache.get_or_load = AsyncMock(side_effect=RuntimeError("Redis unavailable")) + service, es_search, _encoder = _build_service() + service._sleep_stage_cache = sleep_cache + es_search.parse_tags = EsSearch.parse_tags + es_search.filter_by_sleep_stage = AsyncMock(return_value=[es_doc]) + + results = await service.search( + SearchAudioRequest(sleep_stage_tags=["唤醒"], content_tags=["音乐"], top_k=5) + ) + + assert [result["audio_name"] for result in results] == ["ES 唤醒音乐"] + es_search.filter_by_sleep_stage.assert_awaited_once_with(["唤醒"]) + + @pytest.mark.asyncio async def test_search_merges_multi_stage_cache_via_service_path() -> None: """多睡眠阶段请求走缓存 get，由缓存侧按 audio_url 去重。""" @@ -367,9 +447,7 @@ async def test_text_query_normalizes_color_noise_alias_and_excludes_siblings( canonical_doc = _audio_doc( canonical, content_form=["颜色噪音", canonical], description_score=1.0 ) - sibling_doc = _audio_doc( - sibling, content_form=["颜色噪音", sibling], description_score=1.0 - ) + sibling_doc = _audio_doc(sibling, content_form=["颜色噪音", sibling], description_score=1.0) service, es_search, encoder = _build_service( settings=Settings(search_sleep_stage_filter_enabled=False) ) @@ -468,9 +546,7 @@ async def test_search_keeps_candidate_when_disliked_vector_below_threshold() -> ) es_search.parse_tags = EsSearch.parse_tags encoder.encode = AsyncMock( - side_effect=lambda texts: [ - unit_vec if t == "白噪音" else orthogonal_vec for t in texts - ], + side_effect=lambda texts: [unit_vec if t == "白噪音" else orthogonal_vec for t in texts], ) es_search.get_dictionary_vectors = AsyncMock(return_value={"cf_白噪音": unit_vec}) request = SearchAudioRequest( @@ -607,7 +683,6 @@ def track_coarse(candidates, *args, **kwargs): assert [r["audio_name"] for r in results] == ["保留少命中"] - @pytest.mark.asyncio async def test_search_returns_all_when_top_k_omitted() -> None: """未传 top_k 时返回全部候选，不截断。""" @@ -798,9 +873,7 @@ async def test_text_query_extracts_positive_and_negative_tags() -> None: ) encoder.encode_one = AsyncMock(return_value=unit_rain) encoder.encode = AsyncMock( - side_effect=lambda texts: [ - unit_noise if t == "嘈杂" else unit_rain for t in texts - ] + side_effect=lambda texts: [unit_noise if t == "嘈杂" else unit_rain for t in texts] ) es_search.get_dictionary_vectors = AsyncMock( return_value={"cf_雨声": unit_rain, "cf_嘈杂": unit_noise} @@ -1044,7 +1117,9 @@ async def test_warm_query_tag_vectors_encodes_dictionary_labels() -> None: ] ) encoder.encode = AsyncMock( - side_effect=lambda texts: [[float(i)] + [0.0] * (_VECTOR_DIM - 1) for i, _ in enumerate(texts)], + side_effect=lambda texts: [ + [float(i)] + [0.0] * (_VECTOR_DIM - 1) for i, _ in enumerate(texts) + ], ) await service.warm_query_tag_vectors() @@ -1076,9 +1151,7 @@ def test_dislike_penalty_respects_strong_threshold_from_settings() -> None: tags = EsSearch.parse_tags(_audio_doc("候选", content_form=["人声出现位置"])) dictionary = {"cf_人声出现位置": tag_vec} - soft_service, _, _ = _build_service( - settings=Settings(strong_dislike_sim_threshold=0.85) - ) + soft_service, _, _ = _build_service(settings=Settings(strong_dislike_sim_threshold=0.85)) soft_penalty = soft_service._dislike_penalty( tags, disliked_tags=["人声"], @@ -1087,9 +1160,7 @@ def test_dislike_penalty_respects_strong_threshold_from_settings() -> None: ) assert soft_penalty == 0.2 - hard_service, _, _ = _build_service( - settings=Settings(strong_dislike_sim_threshold=0.78) - ) + hard_service, _, _ = _build_service(settings=Settings(strong_dislike_sim_threshold=0.78)) hard_penalty = hard_service._dislike_penalty( tags, disliked_tags=["人声"], @@ -1106,15 +1177,17 @@ def test_tags_mention_voice_detects_substring() -> None: def test_strip_voice_mention_tags_removes_voice_related() -> None: - assert _strip_voice_mention_tags( - ["人声", "避免突发", "避免人声和语言引导", "节奏"] - ) == ["避免突发", "节奏"] + assert _strip_voice_mention_tags(["人声", "避免突发", "避免人声和语言引导", "节奏"]) == [ + "避免突发", + "节奏", + ] def test_usable_dislike_tags_ignores_event_density() -> None: - assert _usable_dislike_tags( - ["人声", "声音事件密度", "避免突发", "声音事件密度"] - ) == ["人声", "避免突发"] + assert _usable_dislike_tags(["人声", "声音事件密度", "避免突发", "声音事件密度"]) == [ + "人声", + "避免突发", + ] def test_voice_value_code_reads_nested_value() -> None: diff --git a/tests/test_sleep_stage_cache.py b/tests/test_sleep_stage_cache.py index 01ff44f..79218a3 100644 --- a/tests/test_sleep_stage_cache.py +++ b/tests/test_sleep_stage_cache.py @@ -2,8 +2,9 @@ from __future__ import annotations -import json +import asyncio import hashlib +import json from typing import Any import pytest @@ -64,13 +65,15 @@ def test_key_builders() -> None: assert SLEEP_STAGE_DOC_KEY_PREFIX == "sleep_stage_v2_doc:" assert build_sleep_stage_index_key("放松") == f"{SLEEP_STAGE_INDEX_KEY_PREFIX}放松" assert build_sleep_stage_doc_key("https://cdn/a.mp3").startswith(SLEEP_STAGE_DOC_KEY_PREFIX) - assert SLEEP_STAGES == ("放松", "入睡", "守护", "清醒") + assert SLEEP_STAGES == ("放松", "入睡", "守护", "唤醒") def test_merge_urls_preserve_order() -> None: - assert merge_urls_preserve_order( - [["https://a", "https://b"], ["https://a", "https://c"]] - ) == ["https://a", "https://b", "https://c"] + assert merge_urls_preserve_order([["https://a", "https://b"], ["https://a", "https://c"]]) == [ + "https://a", + "https://b", + "https://c", + ] @pytest.mark.asyncio @@ -163,6 +166,110 @@ async def loader(stage: str) -> list[dict[str, Any]]: assert json.loads(raw) == [f"https://cdn/{stage}.mp3"] +@pytest.mark.asyncio +async def test_get_or_load_only_populates_missing_stage() -> None: + redis = _FakeRedis() + cache = SleepStageCandidateCache(redis, ttl_sec=60) + relax = _doc("放松", "https://cdn/relax.mp3", ["放松"]) + wake = _doc("唤醒", "https://cdn/wake.mp3", ["唤醒"]) + await cache.set_stage("放松", [relax]) + calls: list[str] = [] + + async def loader(stage: str) -> list[dict[str, Any]]: + calls.append(stage) + return [wake] + + got = await cache.get_or_load(["放松", "唤醒"], loader) + + assert calls == ["唤醒"] + assert [doc["audio_url"] for doc in got] == [ + "https://cdn/relax.mp3", + "https://cdn/wake.mp3", + ] + assert await cache.get(["放松"]) == [relax] + + +@pytest.mark.asyncio +async def test_get_or_load_coalesces_concurrent_stage_misses() -> None: + redis = _FakeRedis() + cache = SleepStageCandidateCache(redis, ttl_sec=60) + wake = _doc("唤醒", "https://cdn/wake.mp3", ["唤醒"]) + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def loader(stage: str) -> list[dict[str, Any]]: + nonlocal calls + assert stage == "唤醒" + calls += 1 + started.set() + await release.wait() + return [wake] + + first = asyncio.create_task(cache.get_or_load(["唤醒"], loader)) + await started.wait() + second = asyncio.create_task(cache.get_or_load(["唤醒"], loader)) + release.set() + + assert await asyncio.gather(first, second) == [[wake], [wake]] + assert calls == 1 + + +@pytest.mark.asyncio +async def test_get_or_load_caches_empty_stage() -> None: + redis = _FakeRedis() + cache = SleepStageCandidateCache(redis, ttl_sec=60) + calls = 0 + + async def loader(stage: str) -> list[dict[str, Any]]: + nonlocal calls + calls += 1 + return [] + + assert await cache.get_or_load(["唤醒"], loader) == [] + assert await cache.get_or_load(["唤醒"], loader) == [] + assert calls == 1 + + +@pytest.mark.asyncio +async def test_get_or_load_does_not_cache_loader_error() -> None: + redis = _FakeRedis() + cache = SleepStageCandidateCache(redis, ttl_sec=60) + + async def failing_loader(stage: str) -> list[dict[str, Any]]: + raise RuntimeError("ES unavailable") + + with pytest.raises(RuntimeError, match="ES unavailable"): + await cache.get_or_load(["唤醒"], failing_loader) + + assert await cache.get(["唤醒"]) is None + + +@pytest.mark.asyncio +async def test_clear_waits_for_inflight_stage_load() -> None: + redis = _FakeRedis() + cache = SleepStageCandidateCache(redis, ttl_sec=60) + started = asyncio.Event() + release = asyncio.Event() + + async def loader(stage: str) -> list[dict[str, Any]]: + started.set() + await release.wait() + return [_doc(stage, "https://cdn/wake.mp3", [stage])] + + load_task = asyncio.create_task(cache.get_or_load(["唤醒"], loader)) + await started.wait() + clear_task = asyncio.create_task(cache.clear_all()) + await asyncio.sleep(0) + assert not clear_task.done() + + release.set() + await load_task + await clear_task + + assert redis.store == {} + + @pytest.mark.asyncio async def test_clear_all_removes_indexes_and_docs() -> None: redis = _FakeRedis() @@ -191,3 +298,14 @@ async def test_clear_all_removes_legacy_indexes_docs_and_candidate_keys() -> Non assert await redis.get("sleep_stage_index:放松") is None assert await redis.get(legacy_doc_key) is None assert await redis.get("sleep_stage_candidates:放松") is None + + +@pytest.mark.asyncio +async def test_clear_all_removes_stale_awake_stage_index() -> None: + redis = _FakeRedis() + cache = SleepStageCandidateCache(redis, ttl_sec=60) + await cache.set_stage("清醒", [_doc("旧清醒", "https://cdn/awake.mp3", ["清醒"])]) + + await cache.clear_all() + + assert await redis.get(build_sleep_stage_index_key("清醒")) is None