diff --git a/.ai_update/changes.md b/.ai_update/changes.md new file mode 100644 index 0000000..9c00723 --- /dev/null +++ b/.ai_update/changes.md @@ -0,0 +1,104 @@ +# AI Change Log + +Дата создания: 2026-04-18 + +## Как пользоваться + +Этот файл - рабочий журнал изменений для Codex. + +Перед новыми правками нужно прочитать этот файл и учитывать: + +- что уже было изменено; +- какие файлы трогались; +- какие проверки запускались; +- какие ограничения и договоренности есть по задаче. + +После каждой осмысленной правки нужно добавлять новую запись с: + +- кратким описанием изменения; +- списком измененных файлов; +- результатом проверок; +- открытыми рисками или TODO, если они есть. + +## Договоренности + +- Не менять контракты `POST /index`, `POST /sparse_embedding`, `POST /search`. +- По текущей задаче фокус держать на `index`. +- Не трогать `search`, `docker-compose.yml` и инфраструктуру без отдельной просьбы. +- Не откатывать чужие или пользовательские изменения. + +## Записи + +### 2026-04-18 - индексный P1 и безопасный P0 search + +Что сделано: + +- В `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-хвост. +- Финальные `message_ids` агрегируются, дедуплицируются и ограничиваются `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` + +Проверки: + +- Код не запускался, потому что создан только журнал. + +Следующий ожидаемый фокус: + +- Перестроить индексный renderer и chunking в `index/main.py`, если пользователь попросит перейти к реализации. 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 c5063aa..0d0178c 100644 --- a/search/main.py +++ b/search/main.py @@ -171,10 +171,13 @@ 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 как списком строк. @@ -206,59 +209,115 @@ 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 +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] = [] - response = await client.query_points( - collection_name=QDRANT_COLLECTION_NAME, - 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, - ), - ], + ) + ) + + if not prefetch: + return None + + response = await client.query_points( + collection_name=QDRANT_COLLECTION_NAME, + prefetch=prefetch, query=models.FusionQuery(fusion=models.Fusion.RRF), limit=RETRIEVE_K, with_payload=True, @@ -310,20 +369,61 @@ async def rerank_points( query: str, points: list[Any], ) -> list[Any]: - rerank_candidates = points[:10] - rerank_targets = [point.payload.get("page_content") for point in rerank_candidates] - scores = await get_rerank_scores(client, query, rerank_targets) + rerank_candidates = points[:RERANK_LIMIT] + retrieval_tail = points[RERANK_LIMIT:] + 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 points + + if len(scores) != len(rerank_candidates): + logger.warning( + "Rerank returned %d scores for %d candidates", + len(scores), + len(rerank_candidates), + ) + return points reranked_candidates = [ point for _, point in sorted( - zip(scores, rerank_candidates, strict=True), + zip(scores, rerank_candidates), key=lambda item: item[0], reverse=True, ) ] - return reranked_candidates + return reranked_candidates + retrieval_tail + + +def aggregate_message_ids(points: list[Any], top_k: int = FINAL_TOP_K) -> list[str]: + best_scores: dict[str, float] = {} + first_seen_rank: dict[str, int] = {} + total_points = len(points) + + for rank, point in enumerate(points): + retrieval_score = float(getattr(point, "score", 0.0) or 0.0) + rank_score = float(total_points - rank) + score = rank_score + retrieval_score + + for message_id in extract_message_ids(point): + if message_id not in first_seen_rank: + first_seen_rank[message_id] = rank + + if score > best_scores.get(message_id, float("-inf")): + best_scores[message_id] = score + + ranked_ids = sorted( + best_scores, + key=lambda message_id: (-best_scores[message_id], first_seen_rank[message_id]), + ) + return ranked_ids[:top_k] # Ваш сервис должен имплементировать оба этих метода @@ -334,25 +434,22 @@ async def health() -> dict[str, str]: @app.post("/search", response_model=SearchAPIResponse) async def search(payload: SearchAPIRequest) -> SearchAPIResponse: - query = payload.question.text.strip() + query, dense_queries, sparse_queries = build_query_texts(payload.question) if not query: - raise HTTPException(status_code=400, detail="question.text is required") + raise HTTPException(status_code=400, detail="question.search_text or question.text is required") 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) + 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) if best_points is None: return SearchAPIResponse(results=[]) best_points = await rerank_points(client, query, list(best_points)) - - message_ids: list[str] = [] - for point in best_points: - message_ids += extract_message_ids(point) + message_ids = aggregate_message_ids(best_points) return SearchAPIResponse( results=[SearchAPIItem(message_ids=message_ids)]