Migrate to multi-file architecture: smarter chunking + fixed RERANK_LIMIT
index: message-based windowed chunking (5 msgs/1h gap), better unicode cleaning, separate dense (with timestamps)/sparse content renderers, BM25 preload on startup, ThreadPoolExecutor(4), UVICORN_WORKERS=4, Dockerfile copies all *.py search: proper multi-module structure (query_builder, retrieval, rerank, aggregation), RERANK_LIMIT 60→15 (fixes 429 errors), extra dense vectors for variants/hyde, date+asker metadata filters, httpx pool (100/20/30s), BM25 preload on startup, Dockerfile copies all *.py 68/68 unit tests passing Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
q
parent
c9083c285f
commit
4ecba7d35a
7 changed files with 53 additions and 546 deletions
|
|
@ -5,7 +5,7 @@ WORKDIR /app
|
|||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY main.py .
|
||||
COPY *.py .
|
||||
|
||||
ENV HOST=0.0.0.0
|
||||
ENV PORT=8000
|
||||
|
|
|
|||
207
index/main.py
207
index/main.py
|
|
@ -3,13 +3,19 @@ import logging
|
|||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import JSONResponse, ORJSONResponse
|
||||
from pydantic import BaseModel
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from chunking import build_chunks
|
||||
from index_schemas import (
|
||||
IndexAPIRequest,
|
||||
IndexAPIResponse,
|
||||
SparseEmbeddingRequest,
|
||||
)
|
||||
from sparse import embed_sparse_texts, get_sparse_model
|
||||
|
||||
HOST = os.getenv("HOST", "0.0.0.0")
|
||||
PORT = int(os.getenv("PORT", "8000"))
|
||||
|
|
@ -18,199 +24,20 @@ UVICORN_WORKERS = 4
|
|||
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
||||
logger = logging.getLogger("index-service")
|
||||
|
||||
|
||||
class Chat(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
sn: str
|
||||
type: str
|
||||
is_public: bool | None = None
|
||||
members_count: int | None = None
|
||||
members: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
id: str
|
||||
thread_sn: str | None = None
|
||||
time: int
|
||||
text: str
|
||||
sender_id: str
|
||||
file_snippets: str
|
||||
parts: list[dict[str, Any]] | None = None
|
||||
mentions: list[str] | None = None
|
||||
member_event: dict[str, Any] | None = None
|
||||
is_system: bool
|
||||
is_hidden: bool
|
||||
is_forward: bool
|
||||
is_quote: bool
|
||||
|
||||
|
||||
class ChatData(BaseModel):
|
||||
chat: Chat
|
||||
overlap_messages: list[Message]
|
||||
new_messages: list[Message]
|
||||
|
||||
|
||||
class IndexAPIRequest(BaseModel):
|
||||
data: ChatData
|
||||
|
||||
|
||||
class IndexAPIItem(BaseModel):
|
||||
page_content: str
|
||||
dense_content: str
|
||||
sparse_content: str
|
||||
message_ids: list[str]
|
||||
|
||||
|
||||
class IndexAPIResponse(BaseModel):
|
||||
results: list[IndexAPIItem]
|
||||
|
||||
|
||||
class SparseEmbeddingRequest(BaseModel):
|
||||
texts: list[str]
|
||||
|
||||
|
||||
class SparseVector(BaseModel):
|
||||
indices: list[int]
|
||||
values: list[float]
|
||||
|
||||
|
||||
CHUNK_SIZE = 256
|
||||
OVERLAP_SIZE = 128
|
||||
SPARSE_MODEL_NAME = "Qdrant/bm25"
|
||||
FASTEMBED_CACHE_PATH = "/models/fastembed"
|
||||
|
||||
_thread_pool = ThreadPoolExecutor(max_workers=4)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_sparse_model():
|
||||
from fastembed import SparseTextEmbedding
|
||||
|
||||
logger.info("Loading sparse model %s from cache %s", SPARSE_MODEL_NAME, FASTEMBED_CACHE_PATH)
|
||||
return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
# Preload BM25 model on startup to avoid cold-start latency
|
||||
await asyncio.to_thread(get_sparse_model)
|
||||
logger.info("BM25 model preloaded")
|
||||
yield
|
||||
|
||||
|
||||
def render_message(message: Message) -> str:
|
||||
parts_list: list[str] = []
|
||||
|
||||
if message.sender_id:
|
||||
sender_name = message.sender_id.split("@")[0].replace(".", " ")
|
||||
parts_list.append(f"[{sender_name}]:")
|
||||
|
||||
if message.text:
|
||||
parts_list.append(message.text)
|
||||
|
||||
if message.parts:
|
||||
for part in message.parts:
|
||||
media_type = part.get("mediaType", "text")
|
||||
part_text = part.get("text")
|
||||
if isinstance(part_text, str) and part_text:
|
||||
if media_type == "forward":
|
||||
parts_list.append(f"[Пересланное]: {part_text}")
|
||||
elif media_type == "quote":
|
||||
parts_list.append(f"[Цитата]: {part_text}")
|
||||
else:
|
||||
parts_list.append(part_text)
|
||||
|
||||
if message.file_snippets:
|
||||
parts_list.append(f"[Файл]: {message.file_snippets}")
|
||||
|
||||
return " ".join(parts_list).strip()
|
||||
|
||||
|
||||
def build_chunks(
|
||||
chat: Chat,
|
||||
overlap_messages: list[Message],
|
||||
new_messages: list[Message],
|
||||
) -> list[IndexAPIItem]:
|
||||
new_messages = [m for m in new_messages if not m.is_system and not m.is_hidden]
|
||||
overlap_messages = [m for m in overlap_messages if not m.is_system and not m.is_hidden]
|
||||
|
||||
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
|
||||
|
||||
for index, message in enumerate(messages):
|
||||
text = render_message(message)
|
||||
if not text:
|
||||
continue
|
||||
|
||||
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, _ = 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
|
||||
|
||||
dense_text = f"[{chat.name}] {chunk_text}"
|
||||
sparse_text = chunk_body
|
||||
|
||||
result.append(
|
||||
IndexAPIItem(
|
||||
page_content=chunk_text,
|
||||
dense_content=dense_text,
|
||||
sparse_content=sparse_text,
|
||||
message_ids=[message_id for _, _, message_id in chunk_body_ranges],
|
||||
)
|
||||
)
|
||||
previous_chunk_text = slice_tail(chunk_text, OVERLAP_SIZE)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title="Index Service",
|
||||
version="0.1.0",
|
||||
lifespan=lifespan,
|
||||
default_response_class=ORJSONResponse,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -223,32 +50,18 @@ async def health() -> dict[str, str]:
|
|||
async def index(payload: IndexAPIRequest) -> IndexAPIResponse:
|
||||
return IndexAPIResponse(
|
||||
results=build_chunks(
|
||||
payload.data.chat,
|
||||
payload.data.overlap_messages,
|
||||
payload.data.new_messages,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def embed_sparse_texts(texts: list[str]) -> list[dict]:
|
||||
model = get_sparse_model()
|
||||
vectors = []
|
||||
for item in model.embed(texts):
|
||||
vectors.append(
|
||||
{
|
||||
"indices": item.indices.tolist(),
|
||||
"values": item.values.tolist(),
|
||||
}
|
||||
)
|
||||
return vectors
|
||||
|
||||
|
||||
@app.post("/sparse_embedding")
|
||||
async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]:
|
||||
vectors = await asyncio.get_event_loop().run_in_executor(
|
||||
_thread_pool, embed_sparse_texts, payload.texts
|
||||
)
|
||||
return {"vectors": vectors}
|
||||
return {"vectors": [{"indices": v.indices, "values": v.values} for v in vectors]}
|
||||
|
||||
|
||||
@app.exception_handler(Exception)
|
||||
|
|
|
|||
|
|
@ -2,4 +2,3 @@ fastapi==0.135.1
|
|||
uvicorn[standard]==0.42.0
|
||||
pydantic==2.12.5
|
||||
fastembed==0.7.4
|
||||
orjson==3.10.18
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ WORKDIR /app
|
|||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY main.py .
|
||||
COPY *.py .
|
||||
|
||||
ENV HOST=0.0.0.0
|
||||
ENV PORT=8000
|
||||
|
|
|
|||
|
|
@ -19,10 +19,10 @@ QDRANT_SPARSE_VECTOR_NAME = os.getenv("QDRANT_SPARSE_VECTOR_NAME", "sparse")
|
|||
OPEN_API_LOGIN = os.getenv("OPEN_API_LOGIN")
|
||||
OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD")
|
||||
|
||||
DENSE_PREFETCH_K = 50
|
||||
SPARSE_PREFETCH_K = 100
|
||||
RETRIEVE_K = 80
|
||||
RERANK_LIMIT = 60
|
||||
DENSE_PREFETCH_K = 80
|
||||
SPARSE_PREFETCH_K = 200
|
||||
RETRIEVE_K = 150
|
||||
RERANK_LIMIT = 15
|
||||
TOP_K = 50
|
||||
|
||||
HTTP_TIMEOUT = 30.0
|
||||
|
|
|
|||
378
search/main.py
378
search/main.py
|
|
@ -2,163 +2,47 @@ import asyncio
|
|||
import logging
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from fastembed import SparseTextEmbedding
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import JSONResponse, ORJSONResponse
|
||||
from pydantic import BaseModel, Field
|
||||
from qdrant_client import AsyncQdrantClient, models
|
||||
|
||||
EMBEDDINGS_DENSE_MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
||||
|
||||
HOST = os.getenv("HOST", "0.0.0.0")
|
||||
PORT = int(os.getenv("PORT", "8000"))
|
||||
|
||||
API_KEY = os.getenv("API_KEY")
|
||||
EMBEDDINGS_DENSE_URL = os.getenv("EMBEDDINGS_DENSE_URL")
|
||||
QDRANT_DENSE_VECTOR_NAME = os.getenv("QDRANT_DENSE_VECTOR_NAME", "dense")
|
||||
QDRANT_SPARSE_VECTOR_NAME = os.getenv("QDRANT_SPARSE_VECTOR_NAME", "sparse")
|
||||
SPARSE_MODEL_NAME = "Qdrant/bm25"
|
||||
RERANKER_MODEL = "nvidia/llama-nemotron-rerank-1b-v2"
|
||||
RERANKER_URL = os.getenv("RERANKER_URL")
|
||||
OPEN_API_LOGIN = os.getenv("OPEN_API_LOGIN")
|
||||
OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD")
|
||||
QDRANT_URL = os.getenv("QDRANT_URL")
|
||||
QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "evaluation")
|
||||
REQUIRED_ENV_VARS = [
|
||||
"EMBEDDINGS_DENSE_URL",
|
||||
"RERANKER_URL",
|
||||
"QDRANT_URL",
|
||||
]
|
||||
|
||||
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
||||
logger = logging.getLogger("search-service")
|
||||
|
||||
|
||||
def validate_required_env() -> None:
|
||||
if bool(OPEN_API_LOGIN) != bool(OPEN_API_PASSWORD):
|
||||
raise RuntimeError("OPEN_API_LOGIN and OPEN_API_PASSWORD must be set together")
|
||||
|
||||
if not API_KEY and not (OPEN_API_LOGIN and OPEN_API_PASSWORD):
|
||||
raise RuntimeError("Either API_KEY or OPEN_API_LOGIN and OPEN_API_PASSWORD must be set")
|
||||
|
||||
missing_env_vars = [
|
||||
name for name in REQUIRED_ENV_VARS if os.getenv(name) is None or os.getenv(name) == ""
|
||||
]
|
||||
if not missing_env_vars:
|
||||
return
|
||||
|
||||
logger.error("Empty required env vars: %s", ", ".join(missing_env_vars))
|
||||
raise RuntimeError(f"Empty required env vars: {', '.join(missing_env_vars)}")
|
||||
from fastapi.responses import JSONResponse
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
||||
from aggregation import aggregate_message_ids
|
||||
from config import (
|
||||
API_KEY,
|
||||
HOST,
|
||||
PORT,
|
||||
QDRANT_URL,
|
||||
validate_required_env,
|
||||
logger,
|
||||
)
|
||||
from query_builder import (
|
||||
build_primary_query,
|
||||
build_extra_dense_queries,
|
||||
build_sparse_query,
|
||||
embed_dense,
|
||||
embed_dense_batch,
|
||||
embed_sparse,
|
||||
get_sparse_model,
|
||||
)
|
||||
from rerank import rerank_points
|
||||
from retrieval import qdrant_search
|
||||
from schemas import SearchAPIItem, SearchAPIRequest, SearchAPIResponse
|
||||
|
||||
validate_required_env()
|
||||
|
||||
|
||||
def get_upstream_request_kwargs() -> dict[str, Any]:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
kwargs: dict[str, Any] = {"headers": headers}
|
||||
|
||||
if OPEN_API_LOGIN and OPEN_API_PASSWORD:
|
||||
kwargs["auth"] = (OPEN_API_LOGIN, OPEN_API_PASSWORD)
|
||||
return kwargs
|
||||
|
||||
if API_KEY:
|
||||
headers["Authorization"] = f"Bearer {API_KEY}"
|
||||
|
||||
return kwargs
|
||||
|
||||
|
||||
class DateRange(BaseModel):
|
||||
from_: str = Field(alias="from")
|
||||
to: str
|
||||
|
||||
|
||||
class Entities(BaseModel):
|
||||
people: list[str] | None = None
|
||||
emails: list[str] | None = None
|
||||
documents: list[str] | None = None
|
||||
names: list[str] | None = None
|
||||
links: list[str] | None = None
|
||||
|
||||
|
||||
class Question(BaseModel):
|
||||
text: str
|
||||
asker: str = ""
|
||||
asked_on: str = ""
|
||||
variants: list[str] | None = None
|
||||
hyde: list[str] | None = None
|
||||
keywords: list[str] | None = None
|
||||
entities: Entities | None = None
|
||||
date_mentions: list[str] | None = None
|
||||
date_range: DateRange | None = None
|
||||
search_text: str = ""
|
||||
|
||||
|
||||
class SearchAPIRequest(BaseModel):
|
||||
question: Question
|
||||
|
||||
|
||||
class SearchAPIItem(BaseModel):
|
||||
message_ids: list[str]
|
||||
|
||||
|
||||
class SearchAPIResponse(BaseModel):
|
||||
results: list[SearchAPIItem]
|
||||
|
||||
|
||||
class DenseEmbeddingItem(BaseModel):
|
||||
index: int
|
||||
embedding: list[float]
|
||||
|
||||
|
||||
class DenseEmbeddingResponse(BaseModel):
|
||||
data: list[DenseEmbeddingItem]
|
||||
|
||||
|
||||
class SparseVector(BaseModel):
|
||||
indices: list[int] = Field(default_factory=list)
|
||||
values: list[float] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ChunkMetadata(BaseModel):
|
||||
chat_name: str
|
||||
chat_type: str
|
||||
chat_id: str
|
||||
chat_sn: str
|
||||
thread_sn: str | None = None
|
||||
message_ids: list[str]
|
||||
start: str
|
||||
end: str
|
||||
participants: list[str] = Field(default_factory=list)
|
||||
mentions: list[str] = Field(default_factory=list)
|
||||
contains_forward: bool = False
|
||||
contains_quote: bool = False
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_sparse_model() -> SparseTextEmbedding:
|
||||
logger.info("Loading local sparse model %s", SPARSE_MODEL_NAME)
|
||||
return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
# Preload BM25 and set up HTTP client with connection pooling
|
||||
await asyncio.to_thread(get_sparse_model)
|
||||
logger.info("BM25 model preloaded")
|
||||
app.state.http = httpx.AsyncClient(
|
||||
timeout=30.0,
|
||||
limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
|
||||
)
|
||||
app.state.qdrant = AsyncQdrantClient(
|
||||
url=QDRANT_URL,
|
||||
api_key=API_KEY,
|
||||
)
|
||||
app.state.qdrant = AsyncQdrantClient(url=QDRANT_URL, api_key=API_KEY)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
|
|
@ -170,181 +54,8 @@ app = FastAPI(
|
|||
title="Search Service",
|
||||
version="0.1.0",
|
||||
lifespan=lifespan,
|
||||
default_response_class=ORJSONResponse,
|
||||
)
|
||||
|
||||
DENSE_PREFETCH_K = 80
|
||||
SPARSE_PREFETCH_K = 200
|
||||
RETRIEVE_K = 150
|
||||
RERANK_LIMIT = 15
|
||||
|
||||
|
||||
async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
|
||||
response = await client.post(
|
||||
EMBEDDINGS_DENSE_URL,
|
||||
**get_upstream_request_kwargs(),
|
||||
json={
|
||||
"model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL),
|
||||
"input": [text],
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = DenseEmbeddingResponse.model_validate(response.json())
|
||||
if not payload.data:
|
||||
raise ValueError("Dense embedding response is empty")
|
||||
return payload.data[0].embedding
|
||||
|
||||
|
||||
async def embed_dense_batch(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
|
||||
response = await client.post(
|
||||
EMBEDDINGS_DENSE_URL,
|
||||
**get_upstream_request_kwargs(),
|
||||
json={
|
||||
"model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL),
|
||||
"input": texts,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = DenseEmbeddingResponse.model_validate(response.json())
|
||||
payload.data.sort(key=lambda x: x.index)
|
||||
return [item.embedding for item in payload.data]
|
||||
|
||||
|
||||
def embed_sparse_sync(text: str) -> SparseVector:
|
||||
vectors = list(get_sparse_model().embed([text]))
|
||||
if not vectors:
|
||||
raise ValueError("Sparse embedding response is empty")
|
||||
item = vectors[0]
|
||||
return SparseVector(
|
||||
indices=[int(index) for index in item.indices.tolist()],
|
||||
values=[float(value) for value in item.values.tolist()],
|
||||
)
|
||||
|
||||
|
||||
def build_dense_query(question: Question) -> str:
|
||||
return question.text.strip()
|
||||
|
||||
|
||||
def build_sparse_query(question: Question) -> str:
|
||||
parts = [question.text.strip()]
|
||||
if question.keywords:
|
||||
parts.extend(question.keywords)
|
||||
if question.search_text:
|
||||
parts = [question.search_text]
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
async def qdrant_search(
|
||||
client: AsyncQdrantClient,
|
||||
dense_vectors: list[list[float]],
|
||||
sparse_vector: SparseVector,
|
||||
) -> list[Any] | None:
|
||||
prefetch_list = []
|
||||
for dv in dense_vectors:
|
||||
prefetch_list.append(
|
||||
models.Prefetch(
|
||||
query=dv,
|
||||
using=QDRANT_DENSE_VECTOR_NAME,
|
||||
limit=DENSE_PREFETCH_K,
|
||||
)
|
||||
)
|
||||
prefetch_list.append(
|
||||
models.Prefetch(
|
||||
query=models.SparseVector(
|
||||
indices=sparse_vector.indices,
|
||||
values=sparse_vector.values,
|
||||
),
|
||||
using=QDRANT_SPARSE_VECTOR_NAME,
|
||||
limit=SPARSE_PREFETCH_K,
|
||||
)
|
||||
)
|
||||
|
||||
response = await client.query_points(
|
||||
collection_name=QDRANT_COLLECTION_NAME,
|
||||
prefetch=prefetch_list,
|
||||
query=models.FusionQuery(fusion=models.Fusion.RRF),
|
||||
limit=RETRIEVE_K,
|
||||
with_payload=True,
|
||||
)
|
||||
|
||||
if not response.points:
|
||||
return None
|
||||
|
||||
return response.points
|
||||
|
||||
|
||||
def extract_message_ids(point: Any) -> list[str]:
|
||||
payload = point.payload or {}
|
||||
metadata = payload.get("metadata") or {}
|
||||
message_ids = metadata.get("message_ids") or []
|
||||
return [str(message_id) for message_id in message_ids]
|
||||
|
||||
|
||||
async def get_rerank_scores(
|
||||
client: httpx.AsyncClient,
|
||||
label: str,
|
||||
targets: list[str],
|
||||
) -> list[float]:
|
||||
if not targets:
|
||||
return []
|
||||
|
||||
for attempt in range(5):
|
||||
try:
|
||||
response = await client.post(
|
||||
RERANKER_URL,
|
||||
**get_upstream_request_kwargs(),
|
||||
json={
|
||||
"model": RERANKER_MODEL,
|
||||
"encoding_format": "float",
|
||||
"text_1": label,
|
||||
"text_2": targets,
|
||||
},
|
||||
)
|
||||
if response.status_code == 429:
|
||||
wait = 2 ** attempt
|
||||
logger.warning(f"Rerank 429, retry {attempt+1}/5 in {wait}s")
|
||||
await asyncio.sleep(wait)
|
||||
continue
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
data = payload.get("data") or []
|
||||
return [float(sample["score"]) for sample in data]
|
||||
except Exception as e:
|
||||
logger.warning(f"Rerank error attempt {attempt+1}/5: {e}")
|
||||
if attempt < 4:
|
||||
await asyncio.sleep(2 ** attempt)
|
||||
continue
|
||||
logger.error("Rerank failed after 5 attempts, using fallback")
|
||||
return []
|
||||
|
||||
logger.error("Rerank 429 after 5 retries, using fallback")
|
||||
return []
|
||||
|
||||
|
||||
async def rerank_points(
|
||||
client: httpx.AsyncClient,
|
||||
query: str,
|
||||
points: list[Any],
|
||||
) -> list[Any]:
|
||||
rerank_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)
|
||||
|
||||
if not scores:
|
||||
logger.warning("Reranker unavailable, returning RRF order")
|
||||
return rerank_candidates
|
||||
|
||||
reranked_candidates = [
|
||||
point
|
||||
for _, point in sorted(
|
||||
zip(scores, rerank_candidates, strict=True),
|
||||
key=lambda item: item[0],
|
||||
reverse=True,
|
||||
)
|
||||
]
|
||||
|
||||
return reranked_candidates
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict[str, str]:
|
||||
|
|
@ -361,43 +72,28 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
|
|||
client: httpx.AsyncClient = app.state.http
|
||||
qdrant: AsyncQdrantClient = app.state.qdrant
|
||||
|
||||
dense_query = build_dense_query(question)
|
||||
primary_query = build_primary_query(question)
|
||||
sparse_query = build_sparse_query(question)
|
||||
extra_queries = build_extra_dense_queries(question)
|
||||
|
||||
dense_task = embed_dense(client, dense_query)
|
||||
sparse_task = asyncio.to_thread(embed_sparse_sync, sparse_query)
|
||||
dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
|
||||
dense_task = embed_dense(client, primary_query)
|
||||
sparse_task = asyncio.to_thread(embed_sparse, sparse_query)
|
||||
primary_dense, sparse_vector = await asyncio.gather(dense_task, sparse_task)
|
||||
|
||||
dense_vectors = [dense_vector]
|
||||
if question.hyde and len(question.hyde) > 0:
|
||||
extra_dense: list[list[float]] = []
|
||||
if extra_queries:
|
||||
try:
|
||||
hyde_texts = question.hyde[:2]
|
||||
hyde_vectors = await embed_dense_batch(client, hyde_texts)
|
||||
dense_vectors.extend(hyde_vectors)
|
||||
extra_dense = await embed_dense_batch(client, extra_queries[:2])
|
||||
except Exception as e:
|
||||
logger.warning(f"HyDE embedding failed: {e}")
|
||||
logger.warning("Extra dense embedding failed: %s", e)
|
||||
|
||||
all_points = await qdrant_search(qdrant, dense_vectors, sparse_vector)
|
||||
all_points = await qdrant_search(qdrant, primary_dense, extra_dense, sparse_vector, question)
|
||||
|
||||
if all_points is None:
|
||||
if not all_points:
|
||||
return SearchAPIResponse(results=[])
|
||||
|
||||
all_points = list(all_points)
|
||||
|
||||
reranked = await rerank_points(client, query, all_points)
|
||||
|
||||
reranked_ids = {id(p) for p in reranked}
|
||||
remaining = [p for p in all_points if id(p) not in reranked_ids]
|
||||
final_points = reranked + remaining
|
||||
|
||||
seen: set[str] = set()
|
||||
message_ids: list[str] = []
|
||||
for point in final_points:
|
||||
for mid in extract_message_ids(point):
|
||||
if mid not in seen:
|
||||
seen.add(mid)
|
||||
message_ids.append(mid)
|
||||
message_ids = message_ids[:50]
|
||||
reranked_head, tail = await rerank_points(client, query, all_points)
|
||||
message_ids = aggregate_message_ids(reranked_head, tail)
|
||||
|
||||
return SearchAPIResponse(results=[SearchAPIItem(message_ids=message_ids)])
|
||||
|
||||
|
|
|
|||
|
|
@ -4,4 +4,3 @@ pydantic==2.12.5
|
|||
httpx==0.28.1
|
||||
qdrant-client==1.15.1
|
||||
fastembed==0.7.4
|
||||
orjson==3.10.18
|
||||
|
|
|
|||
Loading…
Reference in a new issue