main #5

Closed
jeffrey228 wants to merge 3 commits from jeffrey228/vk_hackathon:main into main
5 changed files with 758 additions and 274 deletions

View file

@ -1,27 +1,125 @@
# AI Update Log # AI Change Log
## Scope Дата создания: 2026-04-18
- File: `search/main.py`
- Purpose: fixed and extended retrieval/rerank pipeline according to TODO items.
## 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 Этот файл - рабочий журнал изменений для Codex.
- Earlier step introduced direct `message_ids` deduplication before response.
- Current logic supersedes this by ranking and selecting unique `message_id` values from aggregated scores.
## 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`
Проверки:
- Код не запускался, потому что создан только журнал.

1
.gitignore vendored
View file

@ -214,3 +214,4 @@ __marimo__/
# Streamlit # Streamlit
.streamlit/secrets.toml .streamlit/secrets.toml
.DS_Store

View file

@ -9,7 +9,11 @@ COPY main.py .
ENV HOST=0.0.0.0 ENV HOST=0.0.0.0
ENV PORT=8000 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 FASTEMBED_CACHE_PATH=/models/fastembed
ENV HF_HOME=/models/huggingface ENV HF_HOME=/models/huggingface

View file

@ -1,9 +1,12 @@
import logging import logging
import os import os
import json
import re
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache from functools import lru_cache
from typing import Any from typing import Any
import asyncio import asyncio
import hashlib
from fastapi import FastAPI, Request from fastapi import FastAPI, Request
from fastapi.exceptions import RequestValidationError from fastapi.exceptions import RequestValidationError
@ -88,31 +91,441 @@ app = FastAPI(title="Index Service", version="0.1.0")
# Ваша внутренняя логика построения чанков. Можете делать всё, что посчитаете нужным. # Ваша внутренняя логика построения чанков. Можете делать всё, что посчитаете нужным.
# Текущий код минимальный пример # Текущий код минимальный пример
CHUNK_SIZE = 512 CHUNK_MAX_CHARS = int(os.getenv("INDEX_CHUNK_MAX_CHARS", "2200"))
OVERLAP_SIZE = 256 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" SPARSE_MODEL_NAME = "Qdrant/bm25"
FASTEMBED_CACHE_PATH = "/models/fastembed" FASTEMBED_CACHE_PATH = "/models/fastembed"
# Важная переманная, которая позволяет вычислять sparse вектор в несколько ядер. Не рекомендуется изменять. # Важная переманная, которая позволяет вычислять sparse вектор в несколько ядер. Не рекомендуется изменять.
UVICORN_WORKERS=8 UVICORN_WORKERS = 8
def render_message(message: Message) -> str:
text = ""
if message.text: @dataclass(frozen=True)
text += message.text 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( def build_chunks(
@ -121,71 +534,78 @@ def build_chunks(
) -> list[IndexAPIItem]: ) -> list[IndexAPIItem]:
result: list[IndexAPIItem] = [] result: list[IndexAPIItem] = []
def build_text_and_ranges(messages: list[Message]) -> tuple[str, list[tuple[int, int, str]]]: overlap_units = [
text_parts: list[str] = [] unit
message_ranges: list[tuple[int, int, str]] = [] for message in overlap_messages
position = 0 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): def flush_current() -> None:
text = render_message(message) nonlocal current
if not text: if not current:
continue return
if index > 0 and text_parts: message_ids = unique_preserve_order(
text_parts.append("\n") [item.unit.message_id for item in current if item.is_new]
position += 1 )
if not message_ids:
start = position current = []
text_parts.append(text) return
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
text = chunk_text(current)
result.append( result.append(
IndexAPIItem( IndexAPIItem(
page_content=chunk_text, page_content=text.page,
dense_content=chunk_text, dense_content=text.dense,
sparse_content=chunk_text, sparse_content=text.sparse,
message_ids=[message_id for _, _, message_id in chunk_body_ranges], 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 return result

View file

@ -171,11 +171,14 @@ app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan)
# Внутри шаблона dense и rerank берутся из внешних HTTP endpoint'ов, # Внутри шаблона dense и rerank берутся из внешних HTTP endpoint'ов,
# которые предоставляет проверяющая система. # которые предоставляет проверяющая система.
# Текущий код ниже — минимальный пример search pipeline. # Текущий код ниже — минимальный пример search pipeline.
DENSE_PREFETCH_K = 10 DENSE_PREFETCH_K = 30
SPRASE_PREFETCH_K = 30 SPARSE_PREFETCH_K = 40
RETRIEVE_K = 20 RETRIEVE_K = 80
RERANK_LIMIT = 10 RERANK_LIMIT = 20
FINAL_TOP_K = 50 FINAL_TOP_K = 50
MAX_DENSE_QUERIES = 4
MAX_SPARSE_QUERIES = 3
async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]: async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
# Dense endpoint ожидает OpenAI-compatible body с input как списком строк. # 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()], values=[float(value) for value in item.values.tolist()],
) )
# ПЕРЕПИСАТЬ
async def qdrant_search( def unique_non_empty(values: list[str | None]) -> list[str]:
client: AsyncQdrantClient, result: list[str] = []
dense_vector: list[float], seen: set[str] = set()
sparse_vector: SparseVector,
question_data: Question
) -> Any | None:
must_conditions: []
# Фильтр по диапазону дат (поле metadata.start в Qdrant) [cite: 147, 148, 175] for value in values:
if question_data.date_range: 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( must_conditions.append(
models.FieldCondition( models.FieldCondition(
key="metadata.start", key="metadata.start",
range=models.Range( range=models.DatetimeRange(
gte=question_data.date_range.from_, gte=question.date_range.from_,
lte=question_data.date_range.to_ lte=question.date_range.to,
) ),
) )
) )
# Фильтр по автору вопроса (поле metadata.participants) [cite: 161, 163] return models.Filter(must=must_conditions) if must_conditions else None
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( async def qdrant_search(
collection_name=QDRANT_COLLECTION_NAME, client: AsyncQdrantClient,
prefetch=[ 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( models.Prefetch(
query=dense_vector, query=dense_vector,
using=QDRANT_DENSE_VECTOR_NAME, using=QDRANT_DENSE_VECTOR_NAME,
limit=DENSE_PREFETCH_K, limit=DENSE_PREFETCH_K,
filter=search_filter, filter=search_filter,
), )
)
for sparse_vector in sparse_vectors:
if not sparse_vector.indices:
continue
prefetch.append(
models.Prefetch( models.Prefetch(
query=models.SparseVector( query=models.SparseVector(
indices=sparse_vector.indices, indices=sparse_vector.indices,
values=sparse_vector.values, values=sparse_vector.values,
), ),
using=QDRANT_SPARSE_VECTOR_NAME, using=QDRANT_SPARSE_VECTOR_NAME,
limit=SPRASE_PREFETCH_K, limit=SPARSE_PREFETCH_K,
filter=search_filter, 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 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( response = await client.query_points(
collection_name=QDRANT_COLLECTION_NAME, collection_name=QDRANT_COLLECTION_NAME,
prefetch=[ prefetch=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), query=models.FusionQuery(fusion=models.Fusion.RRF),
limit=RETRIEVE_K, limit=RETRIEVE_K,
with_payload=True, with_payload=True,
@ -320,73 +331,17 @@ async def qdrant_search_dense_only(
return response.points 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]: def deduplicate_points(points: list[Any]) -> list[Any]:
unique_points: list[Any] = [] unique_points: list[Any] = []
seen_ids: set[str] = set() seen_ids: set[str] = set()
for point in points: for point in points:
point_id = str(getattr(point, "id", "")) point_id = getattr(point, "id", None)
if not point_id: if point_id is None:
unique_points.append(point)
continue continue
point_id = str(point_id)
if point_id in seen_ids: if point_id in seen_ids:
continue continue
seen_ids.add(point_id) seen_ids.add(point_id)
@ -444,44 +399,64 @@ async def rerank_points(
) -> list[tuple[Any, float]]: ) -> list[tuple[Any, float]]:
rerank_candidates = points[:RERANK_LIMIT] rerank_candidates = points[:RERANK_LIMIT]
tail_candidates = points[RERANK_LIMIT:] tail_candidates = points[RERANK_LIMIT:]
rerank_targets = [point.payload.get("page_content") for point in rerank_candidates] rerank_targets = [
scores = await get_rerank_scores(client, query, 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 = [ reranked_candidates = [
(point, float(score)) (point, float(score))
for score, point in sorted( for score, point in sorted(
zip(scores, rerank_candidates, strict=True), zip(scores, rerank_candidates),
key=lambda item: item[0], key=lambda item: item[0],
reverse=True, reverse=True,
) )
] ]
tail_with_scores = [ tail_with_scores = [
(point, extract_point_score(point)) (point, extract_point_score(point))
for _, point in sorted( for point in tail_candidates
[(extract_point_score(point), point) for point in tail_candidates],
key=lambda item: item[0],
reverse=True,
)
] ]
return reranked_candidates + tail_with_scores 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] = {} 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)) point_message_ids = set(extract_message_ids(point))
for message_id in point_message_ids: for message_id in point_message_ids:
aggregated_scores[message_id] = aggregated_scores.get(message_id, 0.0) + point_score 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( sorted_items = sorted(
aggregated_scores.items(), 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]] 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) @app.post("/search", response_model=SearchAPIResponse)
async def search(payload: SearchAPIRequest) -> SearchAPIResponse: async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
queries = collect_query_variants(payload.question) query, dense_queries, sparse_queries = build_query_texts(payload.question)
if not queries: if not query:
raise HTTPException(status_code=400, detail="question.search_text or question.text is required") 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 client: httpx.AsyncClient = app.state.http
qdrant: AsyncQdrantClient = app.state.qdrant qdrant: AsyncQdrantClient = app.state.qdrant
all_points: list[Any] = [] dense_vectors = [await embed_dense(client, item) for item in dense_queries]
for query_variant in queries: sparse_vectors = [await embed_sparse(item) for item in sparse_queries]
dense_vector = await embed_dense(client, query_variant) best_points = await qdrant_search(qdrant, dense_vectors, sparse_vectors, payload.question)
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))
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: if not best_points:
return SearchAPIResponse(results=[]) return SearchAPIResponse(results=[])
scored_points = await rerank_points(client, query, list(best_points)) scored_points = await rerank_points(client, query, deduplicate_points(list(best_points)))
aggregated_scores = aggregate_message_scores(scored_points) aggregated_scores, first_seen_rank = aggregate_message_scores(scored_points)
message_ids = select_top_message_ids(aggregated_scores, FINAL_TOP_K) message_ids = select_top_message_ids(aggregated_scores, first_seen_rank, FINAL_TOP_K)
return SearchAPIResponse( return SearchAPIResponse(
results=[SearchAPIItem(message_ids=message_ids)] results=[SearchAPIItem(message_ids=message_ids)]