diff --git a/.ai_update/changes.md b/.ai_update/changes.md index 928a88a..6e65148 100644 --- a/.ai_update/changes.md +++ b/.ai_update/changes.md @@ -1,27 +1,125 @@ -# AI Update Log +# AI Change Log -## Scope -- File: `search/main.py` -- Purpose: fixed and extended retrieval/rerank pipeline according to TODO items. +Дата создания: 2026-04-18 -## Done Changes -- Switched base query selection to `question.search_text` with fallback to `question.text`. -- Added support for `question.variants` as additional query variants in retrieval. -- Added support for `question.hyde` as additional dense-only queries. -- Added support for `question.keywords` as the primary source for sparse query text with fallback to current query variant. -- Stopped losing retrieval candidates after rerank: rerank is applied to head (`RERANK_LIMIT`), tail candidates are preserved. -- Added deduplication of retrieval points by Qdrant point id before rerank. -- Implemented score aggregation by `message_id` (sum of chunk scores mapped to same message). -- Limited final response to `top-50` message ids via `FINAL_TOP_K = 50`. -- Final output message ids are now selected from aggregated scores (sorted by score desc, tie-break by message_id). +## Как пользоваться -## Notes -- Earlier step introduced direct `message_ids` deduplication before response. -- Current logic supersedes this by ranking and selecting unique `message_id` values from aggregated scores. +Этот файл - рабочий журнал изменений для Codex. -## Verification -- Syntax check passed after each main change: `python -m py_compile search/main.py`. +Перед новыми правками нужно прочитать этот файл и учитывать: -## How To Use This Log -- Treat this file as the source of truth for already completed `search/main.py` tasks. -- On next tasks, read this file first to avoid duplicate edits. +- что уже было изменено; +- какие файлы трогались; +- какие проверки запускались; +- какие ограничения и договоренности есть по задаче. + +После каждой осмысленной правки нужно добавлять новую запись с: + +- кратким описанием изменения; +- списком измененных файлов; +- результатом проверок; +- открытыми рисками или TODO, если они есть. + +## Договоренности + +- Не менять контракты `POST /index`, `POST /sparse_embedding`, `POST /search`. +- По текущей задаче фокус держать на `index`, если пользователь не просит иначе. +- Не трогать `docker-compose.yml` и инфраструктуру без отдельной просьбы. +- Не откатывать чужие или пользовательские изменения. + +## Записи + +### 2026-04-18 - разрешены конфликты с upstream/main + +Что сделано: + +- Объединены записи журнала из ветки PR и `upstream/main`. +- В `search/main.py` сохранены расширенные retrieval/rerank доработки из обеих веток. +- Убраны конфликтные маркеры после merge `upstream/main`. +- Сохранены безопасные лимиты retrieval и защита от падения rerank. + +Измененные файлы: + +- `.ai_update/changes.md` +- `search/main.py` + +Проверки: + +- `python3 -m py_compile index/main.py search/main.py` +- Прямой smoke `build_chunks` на `data/Go Nova.json`: 15 чанков, покрыто 25 из 25 сообщений. +- Pure smoke для `search/main.py` helper-функций через stub-модули. + +Открытые риски: + +- Полный интеграционный прогон с настоящими Qdrant/dense/rerank зависит от внешних сервисов. + +### 2026-04-18 - index P1 and search retrieval/rerank + +Что сделано: + +- В `index/main.py` заменен символьный chunking на сборку чанков окнами сообщений. +- Добавлен учет временного разрыва между сообщениями через `INDEX_TIME_GAP_SECONDS`. +- Overlap теперь работает по границам сообщений через `INDEX_CHUNK_OVERLAP_MESSAGES`, а не по хвосту строки. +- Сообщения рендерятся структурно: `author`, `time`, `thread`, `mentions`, флаги, `quote`, `forward`, `file`, `system_event`. +- `file_snippets` парсятся как JSON; в индекс попадают имя файла, mime, url, владелец и дата создания. +- `member_event` превращается в индексируемый системный текст. +- `page_content`, `dense_content`, `sparse_content` разведены. +- В `index/Dockerfile` старый `CHUNK_SIZE=10` заменен на реальные `INDEX_*` настройки chunking. +- В `search/main.py` исправлен runtime-баг с неинициализированным `must_conditions`. +- Основной query в search теперь берется из `question.search_text` с fallback на `question.text`. +- `question.variants`, `question.hyde`, `question.keywords` и entities используются как дополнительные dense/sparse запросы. +- Retrieval делает несколько prefetch-запросов и fusion через Qdrant. +- Rerank больше не выбрасывает retrieval-хвост. +- Retrieval points дедуплицируются по Qdrant point id. +- Финальные `message_ids` агрегируются по score, дедуплицируются и ограничиваются `top-50`. + +Измененные файлы: + +- `index/main.py` +- `index/Dockerfile` +- `search/main.py` +- `.ai_update/changes.md` + +Проверки: + +- `python3 -m py_compile index/main.py search/main.py` +- Прямой smoke `build_chunks` на `data/Go Nova.json`. +- Прямой smoke endpoint-функции `index(...)` на `data/Go Nova.json`. +- Pure smoke для `search` helper-функций через stub-модули, потому в host env нет `qdrant_client` и `httpx`. + +Результаты проверки индекса: + +- Было 29 чанков, стало 15. +- Покрытие сообщений на `data/Go Nova.json`: 25 из 25. +- Системное сообщение с `member_event` больше не выпадает. +- `file_snippets` с `IMG_8471.webp` попадает в `sparse_content`. +- Quote и forward маркеры попадают в `dense_content`. +- `page_content`, `dense_content`, `sparse_content` больше не одинаковые. + +Открытые риски: + +- Полный интеграционный прогон `search` с настоящими Qdrant/dense/rerank локально не выполнялся. +- `date_range` фильтр включается только если установленный `qdrant_client` поддерживает `models.DatetimeRange`. +- Полный docker build локально не запускался. + +### 2026-04-18 - создан журнал изменений + +Что сделано: + +- Создан файл `.ai_update/changes.md`. +- Зафиксировано, что до этого код не менялся, была только разведка репозитория и ТЗ. + +Контекст по текущему состоянию: + +- `index/main.py` тогда использовал символьный chunking. +- `render_message` брал только `message.text` и `parts[*].text`. +- `page_content`, `dense_content`, `sparse_content` были одинаковые. +- В примере `data/Go Nova.json` старый `build_chunks` покрывал 24 из 25 сообщений; системное сообщение с `member_event` выпадало из индекса. + +Измененные файлы: + +- `.ai_update/changes.md` + +Проверки: + +- Код не запускался, потому что создан только журнал. diff --git a/.gitignore b/.gitignore index e15106e..a2a1af2 100644 --- a/.gitignore +++ b/.gitignore @@ -214,3 +214,4 @@ __marimo__/ # Streamlit .streamlit/secrets.toml +.DS_Store diff --git a/index/Dockerfile b/index/Dockerfile index 06878c2..ef80652 100644 --- a/index/Dockerfile +++ b/index/Dockerfile @@ -9,7 +9,11 @@ COPY main.py . ENV HOST=0.0.0.0 ENV PORT=8000 -ENV CHUNK_SIZE=10 +ENV INDEX_CHUNK_MAX_CHARS=2200 +ENV INDEX_MESSAGE_MAX_CHARS=1400 +ENV INDEX_TEXT_SECTION_MAX_CHARS=1000 +ENV INDEX_TIME_GAP_SECONDS=21600 +ENV INDEX_CHUNK_OVERLAP_MESSAGES=1 ENV FASTEMBED_CACHE_PATH=/models/fastembed ENV HF_HOME=/models/huggingface diff --git a/index/main.py b/index/main.py index ae86aec..91c4476 100644 --- a/index/main.py +++ b/index/main.py @@ -1,9 +1,12 @@ import logging import os +import json +import re +from dataclasses import dataclass +from datetime import datetime, timezone from functools import lru_cache from typing import Any import asyncio -import hashlib from fastapi import FastAPI, Request from fastapi.exceptions import RequestValidationError @@ -88,31 +91,441 @@ app = FastAPI(title="Index Service", version="0.1.0") # Ваша внутренняя логика построения чанков. Можете делать всё, что посчитаете нужным. # Текущий код – минимальный пример -CHUNK_SIZE = 512 -OVERLAP_SIZE = 256 +CHUNK_MAX_CHARS = int(os.getenv("INDEX_CHUNK_MAX_CHARS", "2200")) +MESSAGE_MAX_CHARS = int(os.getenv("INDEX_MESSAGE_MAX_CHARS", "1400")) +TEXT_SECTION_MAX_CHARS = int(os.getenv("INDEX_TEXT_SECTION_MAX_CHARS", "1000")) +TIME_GAP_SECONDS = int(os.getenv("INDEX_TIME_GAP_SECONDS", str(6 * 60 * 60))) +CHUNK_OVERLAP_MESSAGES = int(os.getenv("INDEX_CHUNK_OVERLAP_MESSAGES", "1")) SPARSE_MODEL_NAME = "Qdrant/bm25" FASTEMBED_CACHE_PATH = "/models/fastembed" # Важная переманная, которая позволяет вычислять sparse вектор в несколько ядер. Не рекомендуется изменять. -UVICORN_WORKERS=8 +UVICORN_WORKERS = 8 -def render_message(message: Message) -> str: - text = "" - if message.text: - text += message.text +@dataclass(frozen=True) +class RenderedText: + page: str + dense: str + sparse: str - if message.parts: - parts_text: list[str] = [] - for part in message.parts: - # parts различаются по своему типу, см. README.md - part_text = part.get("text") - if isinstance(part_text, str) and part_text: - parts_text.append(part_text) - if parts_text: - text += "\n".join(parts_text) - return text +@dataclass(frozen=True) +class RenderedUnit: + message_id: str + time: int + text: RenderedText + + +@dataclass(frozen=True) +class ChunkUnit: + unit: RenderedUnit + is_new: bool + + +def clean_text(value: Any) -> str: + if not isinstance(value, str): + return "" + + text = ( + value.replace("\r\n", "\n") + .replace("\r", "\n") + .replace("\u200b", " ") + .replace("\xa0", " ") + ) + lines = [re.sub(r"[ \t]+", " ", line).strip() for line in text.split("\n")] + result: list[str] = [] + previous_blank = False + + for line in lines: + if not line: + if result and not previous_blank: + result.append("") + previous_blank = True + continue + + result.append(line) + previous_blank = False + + return "\n".join(result).strip() + + +def unique_preserve_order(values: list[str]) -> list[str]: + seen: set[str] = set() + result: list[str] = [] + + for value in values: + if value and value not in seen: + seen.add(value) + result.append(value) + + return result + + +def format_time(timestamp: int) -> str: + return datetime.fromtimestamp(timestamp, tz=timezone.utc).isoformat().replace("+00:00", "Z") + + +def format_optional_time(value: Any) -> str: + if value is None or value == "": + return "" + + try: + return format_time(int(value)) + except (TypeError, ValueError, OSError, OverflowError): + return clean_text(str(value)) + + +def lexical_terms(value: str) -> str: + return clean_text(re.sub(r"[^0-9A-Za-zА-Яа-яЁё]+", " ", value)) + + +def split_long_text(text: str, max_chars: int) -> list[str]: + text = clean_text(text) + if not text: + return [] + if len(text) <= max_chars: + return [text] + + paragraphs = [item.strip() for item in re.split(r"\n{2,}", text) if item.strip()] + pieces: list[str] = [] + current = "" + + def append_current() -> None: + nonlocal current + if current: + pieces.append(current) + current = "" + + def split_oversized(paragraph: str) -> list[str]: + words = paragraph.split() + result: list[str] = [] + part = "" + + for word in words: + if not part: + part = word + continue + + if len(part) + 1 + len(word) <= max_chars: + part += " " + word + else: + result.append(part) + part = word + + if part: + result.append(part) + + return result + + for paragraph in paragraphs: + candidates = [paragraph] if len(paragraph) <= max_chars else split_oversized(paragraph) + for candidate in candidates: + separator = "\n\n" if current else "" + if current and len(current) + len(separator) + len(candidate) > max_chars: + append_current() + + current = candidate if not current else current + separator + candidate + + append_current() + return pieces + + +def combine_texts(items: list[RenderedText], separator: str = "\n") -> RenderedText: + return RenderedText( + page=separator.join(item.page for item in items if item.page).strip(), + dense=separator.join(item.dense for item in items if item.dense).strip(), + sparse=separator.join(item.sparse for item in items if item.sparse).strip(), + ) + + +def rendered_length(text: RenderedText) -> int: + return max(len(text.page), len(text.dense), len(text.sparse)) + + +def message_header(message: Message) -> RenderedText: + timestamp = format_time(message.time) + mentions = unique_preserve_order(message.mentions or []) + + page_lines = [ + f"author: {message.sender_id}", + f"time: {timestamp}", + ] + dense_lines = [ + f"author: {message.sender_id}", + f"message_time: {timestamp}", + ] + sparse_terms = [ + "author", + message.sender_id, + lexical_terms(message.sender_id), + timestamp, + ] + + if message.thread_sn: + page_lines.append(f"thread: {message.thread_sn}") + dense_lines.append(f"thread: {message.thread_sn}") + sparse_terms.extend(["thread", message.thread_sn, lexical_terms(message.thread_sn)]) + + if mentions: + mentions_text = ", ".join(mentions) + page_lines.append(f"mentions: {mentions_text}") + dense_lines.append(f"mentions: {mentions_text}") + sparse_terms.extend(["mentions", *mentions, *(lexical_terms(item) for item in mentions)]) + + flags = [] + if message.is_system: + flags.append("system") + if message.is_forward: + flags.append("forward") + if message.is_quote: + flags.append("quote") + if message.is_hidden: + flags.append("hidden") + + if flags: + flags_text = ", ".join(flags) + dense_lines.append(f"message_flags: {flags_text}") + sparse_terms.extend(flags) + + return RenderedText( + page="\n".join(page_lines), + dense="\n".join(dense_lines), + sparse=" ".join(term for term in sparse_terms if term), + ) + + +def render_plain_sections(text: str, label: str = "text") -> list[RenderedText]: + sections: list[RenderedText] = [] + pieces = split_long_text(text, TEXT_SECTION_MAX_CHARS) + + for index, piece in enumerate(pieces): + suffix = f" part {index + 1}/{len(pieces)}" if len(pieces) > 1 else "" + sections.append( + RenderedText( + page=piece, + dense=f"{label}{suffix}: {piece}", + sparse=f"{label} {piece}", + ) + ) + + return sections + + +def render_part_sections(part: dict[str, Any]) -> list[RenderedText]: + text = clean_text(part.get("text")) + if not text: + return [] + + media_type = clean_text(part.get("mediaType") or part.get("type") or "text").lower() + source = clean_text(part.get("sn")) + part_time = part.get("time") + source_bits = [] + if source: + source_bits.append(f"source: {source}") + formatted_part_time = format_optional_time(part_time) + if formatted_part_time: + source_bits.append(f"source_time: {formatted_part_time}") + source_text = ", ".join(source_bits) + + if media_type == "quote": + label = f"quote from {source}" if source else "quote" + dense_label = f"quote; {source_text}" if source_text else "quote" + sparse_prefix = f"quote цитата {source} {lexical_terms(source)}" + elif media_type == "forward": + label = f"forwarded from {source}" if source else "forwarded" + dense_label = f"forwarded_message; {source_text}" if source_text else "forwarded_message" + sparse_prefix = f"forward forwarded_message пересланное {source} {lexical_terms(source)}" + else: + label = "text" + dense_label = "text" + sparse_prefix = "text" + + sections: list[RenderedText] = [] + pieces = split_long_text(text, TEXT_SECTION_MAX_CHARS) + + for index, piece in enumerate(pieces): + suffix = f" part {index + 1}/{len(pieces)}" if len(pieces) > 1 else "" + page_prefix = f"{label}{suffix}:" + sections.append( + RenderedText( + page=f"{page_prefix}\n{piece}" if media_type in {"quote", "forward"} else piece, + dense=f"{dense_label}{suffix}: {piece}", + sparse=f"{sparse_prefix} {piece}", + ) + ) + + return sections + + +def render_file_sections(raw_snippets: str) -> list[RenderedText]: + raw_snippets = clean_text(raw_snippets) + if not raw_snippets: + return [] + + try: + parsed = json.loads(raw_snippets) + except json.JSONDecodeError: + return [ + RenderedText( + page=f"file_snippet: {raw_snippets}", + dense=f"file_snippet: {raw_snippets}", + sparse=f"file file_snippet {raw_snippets}", + ) + ] + + snippets = parsed if isinstance(parsed, list) else [parsed] + sections: list[RenderedText] = [] + + for snippet in snippets: + if not isinstance(snippet, dict): + text = clean_text(str(snippet)) + sections.append(RenderedText(page=f"file: {text}", dense=f"file: {text}", sparse=f"file {text}")) + continue + + name = clean_text(snippet.get("name")) + mime = clean_text(snippet.get("mime")) + url = clean_text(snippet.get("original_url") or snippet.get("url")) + owner = clean_text(snippet.get("uid")) + created = clean_text(snippet.get("date_create")) + file_bits = [ + f"name: {name}" if name else "", + f"mime: {mime}" if mime else "", + f"url: {url}" if url else "", + f"owner: {owner}" if owner else "", + f"created: {created}" if created else "", + ] + file_text = ", ".join(bit for bit in file_bits if bit) + sparse_terms = " ".join( + term + for term in [ + "file", + "attachment", + "document", + name, + lexical_terms(name), + mime, + url, + owner, + lexical_terms(owner), + created, + ] + if term + ) + sections.append( + RenderedText( + page=f"file: {file_text}", + dense=f"file: {file_text}", + sparse=sparse_terms, + ) + ) + + return sections + + +def render_member_event(message: Message) -> list[RenderedText]: + event = message.member_event + if not event: + return [] + + event_type = clean_text(event.get("type") or "member_event") + members_raw = event.get("members") + members = [clean_text(item) for item in members_raw] if isinstance(members_raw, list) else [] + members = unique_preserve_order([item for item in members if item]) + + if members: + members_text = ", ".join(members) + else: + members_text = " ".join(clean_text(str(value)) for value in event.values() if value) + + page = f"system_event: {event_type}; actor: {message.sender_id}; members: {members_text}" + dense = ( + f"system_event: {event_type}; action: add or update chat members; " + f"actor: {message.sender_id}; members: {members_text}" + ) + sparse = " ".join( + term + for term in [ + "system_event", + "member_event", + event_type, + "addMembers", + "добавление участников", + message.sender_id, + lexical_terms(message.sender_id), + members_text, + lexical_terms(members_text), + ] + if term + ) + + return [RenderedText(page=page, dense=dense, sparse=sparse)] + + +def render_message_sections(message: Message) -> list[RenderedText]: + sections: list[RenderedText] = [] + sections.extend(render_plain_sections(message.text, "message_text")) + + for part in message.parts or []: + if isinstance(part, dict): + sections.extend(render_part_sections(part)) + + sections.extend(render_file_sections(message.file_snippets)) + sections.extend(render_member_event(message)) + + return sections + + +def render_message_units(message: Message) -> list[RenderedUnit]: + sections = render_message_sections(message) + if not sections: + return [] + + header = message_header(message) + units: list[RenderedUnit] = [] + current: list[RenderedText] = [] + + def flush() -> None: + nonlocal current + if not current: + return + + text = combine_texts([header, *current]) + units.append(RenderedUnit(message_id=message.id, time=message.time, text=text)) + current = [] + + for section in sections: + candidate = combine_texts([header, *current, section]) + if current and rendered_length(candidate) > MESSAGE_MAX_CHARS: + flush() + + current.append(section) + + flush() + return units + + +def chunk_text(items: list[ChunkUnit]) -> RenderedText: + return RenderedText( + page="\n\n".join(item.unit.text.page for item in items if item.unit.text.page).strip(), + dense="\n\n".join(item.unit.text.dense for item in items if item.unit.text.dense).strip(), + sparse="\n\n".join(item.unit.text.sparse for item in items if item.unit.text.sparse).strip(), + ) + + +def chunk_length(items: list[ChunkUnit]) -> int: + return rendered_length(chunk_text(items)) + + +def trim_context(context: list[RenderedUnit], unit: RenderedUnit) -> list[RenderedUnit]: + result = context[-CHUNK_OVERLAP_MESSAGES:] if CHUNK_OVERLAP_MESSAGES > 0 else [] + items = [ChunkUnit(item, False) for item in result] + [ChunkUnit(unit, True)] + + while result and chunk_length(items) > CHUNK_MAX_CHARS: + result = result[1:] + items = [ChunkUnit(item, False) for item in result] + [ChunkUnit(unit, True)] + + return result def build_chunks( @@ -121,71 +534,78 @@ def build_chunks( ) -> list[IndexAPIItem]: result: list[IndexAPIItem] = [] - def build_text_and_ranges(messages: list[Message]) -> tuple[str, list[tuple[int, int, str]]]: - text_parts: list[str] = [] - message_ranges: list[tuple[int, int, str]] = [] - position = 0 + overlap_units = [ + unit + for message in overlap_messages + for unit in render_message_units(message) + ] + new_units = [ + unit + for message in new_messages + for unit in render_message_units(message) + ] + current: list[ChunkUnit] = [] + last_new_time: int | None = None - for index, message in enumerate(messages): - text = render_message(message) - if not text: - continue + def flush_current() -> None: + nonlocal current + if not current: + return - if index > 0 and text_parts: - text_parts.append("\n") - position += 1 - - start = position - text_parts.append(text) - position += len(text) - message_ranges.append((start, position, message.id)) - - return "".join(text_parts), message_ranges - - def slice_tail( - text: str, - tail_size: int, - ) -> str: - if tail_size <= 0: - return "" - - tail_start = max(0, len(text) - tail_size) - return text[tail_start:] - - overlap_text, overlap_message_ranges = build_text_and_ranges(overlap_messages) - previous_chunk_text = slice_tail(overlap_text, OVERLAP_SIZE) - - new_text, new_message_ranges = build_text_and_ranges(new_messages) - - for start in range(0, len(new_text), CHUNK_SIZE): - chunk_body = new_text[start : start + CHUNK_SIZE] - if not chunk_body: - continue - - chunk_body_ranges = [ - ( - max(message_start, start) - start, - min(message_end, start + len(chunk_body)) - start, - message_id, - ) - for message_start, message_end, message_id in new_message_ranges - if message_end > start and message_start < start + len(chunk_body) - ] - chunk_overlap = previous_chunk_text - chunk_text = chunk_overlap - if chunk_text and chunk_body: - chunk_text += "\n" - chunk_text += chunk_body + message_ids = unique_preserve_order( + [item.unit.message_id for item in current if item.is_new] + ) + if not message_ids: + current = [] + return + text = chunk_text(current) result.append( IndexAPIItem( - page_content=chunk_text, - dense_content=chunk_text, - sparse_content=chunk_text, - message_ids=[message_id for _, _, message_id in chunk_body_ranges], + page_content=text.page, + dense_content=text.dense, + sparse_content=text.sparse, + message_ids=message_ids, ) ) - previous_chunk_text = slice_tail(chunk_text, OVERLAP_SIZE) + current = [] + + def request_overlap_context(unit: RenderedUnit) -> list[RenderedUnit]: + close_units = [ + item + for item in overlap_units + if abs(unit.time - item.time) <= TIME_GAP_SECONDS + ] + return trim_context(close_units, unit) + + for unit in new_units: + if not current: + context = request_overlap_context(unit) + current = [ChunkUnit(item, False) for item in context] + current.append(ChunkUnit(unit, True)) + last_new_time = unit.time + continue + + gap = abs(unit.time - last_new_time) if last_new_time is not None else 0 + candidate = [*current, ChunkUnit(unit, True)] + should_split = gap > TIME_GAP_SECONDS or chunk_length(candidate) > CHUNK_MAX_CHARS + + if should_split: + previous_new_units = [item.unit for item in current if item.is_new] + context = ( + trim_context(previous_new_units, unit) + if gap <= TIME_GAP_SECONDS + else request_overlap_context(unit) + ) + flush_current() + current = [ChunkUnit(item, False) for item in context] + current.append(ChunkUnit(unit, True)) + else: + current.append(ChunkUnit(unit, True)) + + last_new_time = unit.time + + flush_current() return result diff --git a/search/main.py b/search/main.py index d7f1c21..e51af05 100644 --- a/search/main.py +++ b/search/main.py @@ -171,11 +171,14 @@ app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan) # Внутри шаблона dense и rerank берутся из внешних HTTP endpoint'ов, # которые предоставляет проверяющая система. # Текущий код ниже — минимальный пример search pipeline. -DENSE_PREFETCH_K = 10 -SPRASE_PREFETCH_K = 30 -RETRIEVE_K = 20 -RERANK_LIMIT = 10 +DENSE_PREFETCH_K = 30 +SPARSE_PREFETCH_K = 40 +RETRIEVE_K = 80 +RERANK_LIMIT = 20 FINAL_TOP_K = 50 +MAX_DENSE_QUERIES = 4 +MAX_SPARSE_QUERIES = 3 + async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]: # Dense endpoint ожидает OpenAI-compatible body с input как списком строк. @@ -207,108 +210,116 @@ async def embed_sparse(text: str) -> SparseVector: values=[float(value) for value in item.values.tolist()], ) -# ПЕРЕПИСАТЬ -async def qdrant_search( - client: AsyncQdrantClient, - dense_vector: list[float], - sparse_vector: SparseVector, - question_data: Question -) -> Any | None: - must_conditions: [] +def unique_non_empty(values: list[str | None]) -> list[str]: + result: list[str] = [] + seen: set[str] = set() - # Фильтр по диапазону дат (поле metadata.start в Qdrant) [cite: 147, 148, 175] - if question_data.date_range: + for value in values: + text = (value or "").strip() + if text and text not in seen: + seen.add(text) + result.append(text) + + return result + + +def question_entity_terms(question: Question) -> list[str]: + entities = question.entities + if entities is None: + return [] + + values: list[str | None] = [] + values.extend(entities.people or []) + values.extend(entities.emails or []) + values.extend(entities.documents or []) + values.extend(entities.names or []) + values.extend(entities.links or []) + return unique_non_empty(values) + + +def build_query_texts(question: Question) -> tuple[str, list[str], list[str]]: + primary_query = (question.search_text or question.text).strip() + dense_queries = unique_non_empty( + [ + primary_query, + *(question.variants or []), + *(question.hyde or []), + ] + )[:MAX_DENSE_QUERIES] + + keyword_query = " ".join(question.keywords or []).strip() + entity_query = " ".join(question_entity_terms(question)).strip() + sparse_queries = unique_non_empty( + [ + keyword_query, + entity_query, + primary_query, + *(question.variants or []), + ] + )[:MAX_SPARSE_QUERIES] + + return primary_query, dense_queries, sparse_queries + + +def build_search_filter(question: Question) -> models.Filter | None: + must_conditions: list[Any] = [] + + if question.date_range and hasattr(models, "DatetimeRange"): must_conditions.append( models.FieldCondition( key="metadata.start", - range=models.Range( - gte=question_data.date_range.from_, - lte=question_data.date_range.to_ - ) + range=models.DatetimeRange( + gte=question.date_range.from_, + lte=question.date_range.to, + ), ) ) - # Фильтр по автору вопроса (поле metadata.participants) [cite: 161, 163] - if question_data.asker: - must_conditions.append( - models.FieldCondition( - key="metadata.participants", - match=models.MatchValue(value=question_data.asker) - ) - ) + return models.Filter(must=must_conditions) if must_conditions else None - # Создаем итоговый объект фильтра, если есть условия - search_filter = models.Filter(must=must_conditions) if must_conditions else None - response = await client.query_points( - collection_name=QDRANT_COLLECTION_NAME, - prefetch=[ +async def qdrant_search( + client: AsyncQdrantClient, + dense_vectors: list[list[float]], + sparse_vectors: list[SparseVector], + question_data: Question, +) -> Any | None: + search_filter = build_search_filter(question_data) + prefetch: list[models.Prefetch] = [] + + for dense_vector in dense_vectors: + prefetch.append( models.Prefetch( query=dense_vector, using=QDRANT_DENSE_VECTOR_NAME, limit=DENSE_PREFETCH_K, filter=search_filter, - ), + ) + ) + + for sparse_vector in sparse_vectors: + if not sparse_vector.indices: + continue + + prefetch.append( models.Prefetch( query=models.SparseVector( indices=sparse_vector.indices, values=sparse_vector.values, ), using=QDRANT_SPARSE_VECTOR_NAME, - limit=SPRASE_PREFETCH_K, + limit=SPARSE_PREFETCH_K, filter=search_filter, - ), - ], - query=models.FusionQuery(fusion=models.Fusion.RRF), - limit=RETRIEVE_K, - with_payload=True, - ) + ) + ) - if not response.points: + if not prefetch: return None - 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, - ), - ], + prefetch=prefetch, query=models.FusionQuery(fusion=models.Fusion.RRF), limit=RETRIEVE_K, with_payload=True, @@ -320,73 +331,17 @@ async def qdrant_search_dense_only( 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 build_sparse_query_text(question: Question, fallback_query: str) -> str: - keywords: list[str] = [] - seen: set[str] = set() - - for keyword in question.keywords or []: - normalized = keyword.strip() - if not normalized: - continue - if normalized in seen: - continue - seen.add(normalized) - keywords.append(normalized) - - if keywords: - return " ".join(keywords) - - return fallback_query - - 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: + point_id = getattr(point, "id", None) + if point_id is None: + unique_points.append(point) continue + + point_id = str(point_id) if point_id in seen_ids: continue seen_ids.add(point_id) @@ -444,44 +399,64 @@ async def rerank_points( ) -> list[tuple[Any, float]]: rerank_candidates = points[:RERANK_LIMIT] tail_candidates = points[RERANK_LIMIT:] - rerank_targets = [point.payload.get("page_content") for point in rerank_candidates] - scores = await get_rerank_scores(client, query, rerank_targets) + rerank_targets = [ + str((point.payload or {}).get("page_content") or "") + for point in rerank_candidates + ] + + try: + scores = await get_rerank_scores(client, query, rerank_targets) + except Exception: + logger.exception("Rerank failed, returning retrieval order") + return [(point, extract_point_score(point)) for point in points] + + if len(scores) != len(rerank_candidates): + logger.warning( + "Rerank returned %d scores for %d candidates", + len(scores), + len(rerank_candidates), + ) + return [(point, extract_point_score(point)) for point in points] reranked_candidates = [ (point, float(score)) for score, point in sorted( - zip(scores, rerank_candidates, strict=True), + zip(scores, rerank_candidates), key=lambda item: item[0], reverse=True, ) ] tail_with_scores = [ (point, extract_point_score(point)) - for _, point in sorted( - [(extract_point_score(point), point) for point in tail_candidates], - key=lambda item: item[0], - reverse=True, - ) + for point in tail_candidates ] return reranked_candidates + tail_with_scores -def aggregate_message_scores(scored_points: list[tuple[Any, float]]) -> dict[str, float]: +def aggregate_message_scores( + scored_points: list[tuple[Any, float]], +) -> tuple[dict[str, float], dict[str, int]]: aggregated_scores: dict[str, float] = {} + first_seen_rank: dict[str, int] = {} - for point, point_score in scored_points: + for rank, (point, point_score) in enumerate(scored_points): point_message_ids = set(extract_message_ids(point)) for message_id in point_message_ids: aggregated_scores[message_id] = aggregated_scores.get(message_id, 0.0) + point_score + first_seen_rank.setdefault(message_id, rank) - return aggregated_scores + return aggregated_scores, first_seen_rank -def select_top_message_ids(aggregated_scores: dict[str, float], limit: int) -> list[str]: +def select_top_message_ids( + aggregated_scores: dict[str, float], + first_seen_rank: dict[str, int], + limit: int, +) -> list[str]: sorted_items = sorted( aggregated_scores.items(), - key=lambda item: (-item[1], item[0]), + key=lambda item: (-item[1], first_seen_rank[item[0]]), ) return [message_id for message_id, _ in sorted_items[:limit]] @@ -494,37 +469,23 @@ async def health() -> dict[str, str]: @app.post("/search", response_model=SearchAPIResponse) async def search(payload: SearchAPIRequest) -> SearchAPIResponse: - queries = collect_query_variants(payload.question) - if not queries: + query, dense_queries, sparse_queries = build_query_texts(payload.question) + if not query: 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 - all_points: list[Any] = [] - for query_variant in queries: - dense_vector = await embed_dense(client, query_variant) - sparse_query_text = build_sparse_query_text(payload.question, query_variant) - sparse_vector = await embed_sparse(sparse_query_text) - points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question) - if points: - all_points.extend(list(points)) + dense_vectors = [await embed_dense(client, item) for item in dense_queries] + sparse_vectors = [await embed_sparse(item) for item in sparse_queries] + best_points = await qdrant_search(qdrant, dense_vectors, sparse_vectors, payload.question) - 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=[]) - scored_points = await rerank_points(client, query, list(best_points)) - aggregated_scores = aggregate_message_scores(scored_points) - message_ids = select_top_message_ids(aggregated_scores, FINAL_TOP_K) + scored_points = await rerank_points(client, query, deduplicate_points(list(best_points))) + aggregated_scores, first_seen_rank = aggregate_message_scores(scored_points) + message_ids = select_top_message_ids(aggregated_scores, first_seen_rank, FINAL_TOP_K) return SearchAPIResponse( results=[SearchAPIItem(message_ids=message_ids)]