diff --git a/search/main.py b/search/main.py index c5063aa..6a0df49 100644 --- a/search/main.py +++ b/search/main.py @@ -270,6 +270,111 @@ async def qdrant_search( return response.points +async def qdrant_search_dense_only( + client: AsyncQdrantClient, + dense_vector: list[float], + question_data: Question, +) -> Any | None: + must_conditions: [] + + if question_data.date_range: + must_conditions.append( + models.FieldCondition( + key="metadata.start", + range=models.Range( + gte=question_data.date_range.from_, + lte=question_data.date_range.to_, + ), + ) + ) + + if question_data.asker: + must_conditions.append( + models.FieldCondition( + key="metadata.participants", + match=models.MatchValue(value=question_data.asker), + ) + ) + + search_filter = models.Filter(must=must_conditions) if must_conditions else None + + response = await client.query_points( + collection_name=QDRANT_COLLECTION_NAME, + prefetch=[ + models.Prefetch( + query=dense_vector, + using=QDRANT_DENSE_VECTOR_NAME, + limit=DENSE_PREFETCH_K, + filter=search_filter, + ), + ], + query=models.FusionQuery(fusion=models.Fusion.RRF), + limit=RETRIEVE_K, + with_payload=True, + ) + + if not response.points: + return None + + return response.points + + +def collect_query_variants(question: Question) -> list[str]: + variants: list[str] = [] + seen: set[str] = set() + + def add_query(text: str | None) -> None: + if text is None: + return + normalized = text.strip() + if not normalized: + return + if normalized in seen: + return + seen.add(normalized) + variants.append(normalized) + + add_query(question.search_text) + add_query(question.text) + + for variant in question.variants or []: + add_query(variant) + + return variants + + +def collect_hyde_queries(question: Question, base_queries: list[str]) -> list[str]: + hyde_queries: list[str] = [] + seen: set[str] = set(base_queries) + + for hyde_query in question.hyde or []: + normalized = hyde_query.strip() + if not normalized: + continue + if normalized in seen: + continue + seen.add(normalized) + hyde_queries.append(normalized) + + return hyde_queries + + +def deduplicate_points(points: list[Any]) -> list[Any]: + unique_points: list[Any] = [] + seen_ids: set[str] = set() + + for point in points: + point_id = str(getattr(point, "id", "")) + if not point_id: + continue + if point_id in seen_ids: + continue + seen_ids.add(point_id) + unique_points.append(point) + + return unique_points + + def extract_message_ids(point: Any) -> list[str]: payload = point.payload or {} metadata = payload.get("metadata") or {} @@ -334,18 +439,31 @@ async def health() -> dict[str, str]: @app.post("/search", response_model=SearchAPIResponse) async def search(payload: SearchAPIRequest) -> SearchAPIResponse: - query = payload.question.text.strip() - if not query: - raise HTTPException(status_code=400, detail="question.text is required") + queries = collect_query_variants(payload.question) + if not queries: + raise HTTPException(status_code=400, detail="question.search_text or question.text is required") + hyde_queries = collect_hyde_queries(payload.question, queries) + query = queries[0] client: httpx.AsyncClient = app.state.http qdrant: AsyncQdrantClient = app.state.qdrant - dense_vector = await embed_dense(client, query) - sparse_vector = await embed_sparse(query) - best_points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question) + all_points: list[Any] = [] + for query_variant in queries: + dense_vector = await embed_dense(client, query_variant) + sparse_vector = await embed_sparse(query_variant) + points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question) + if points: + all_points.extend(list(points)) - if best_points is None: + for hyde_query in hyde_queries: + hyde_dense_vector = await embed_dense(client, hyde_query) + hyde_points = await qdrant_search_dense_only(qdrant, hyde_dense_vector, payload.question) + if hyde_points: + all_points.extend(list(hyde_points)) + + best_points = deduplicate_points(all_points) + if not best_points: return SearchAPIResponse(results=[]) best_points = await rerank_points(client, query, list(best_points))