forked from zovos/vk_hackathon
521 lines
15 KiB
Python
521 lines
15 KiB
Python
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
|
||
from pydantic import BaseModel, Field
|
||
from qdrant_client import AsyncQdrantClient, models
|
||
|
||
EMBEDDINGS_DENSE_MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
||
|
||
# Ваш сервис должен считывать эти переменные из окружения (env), так как проверяющая система управляет ими
|
||
HOST = os.getenv("HOST", "0.0.0.0")
|
||
PORT = int(os.getenv("PORT", "8003"))
|
||
|
||
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)}")
|
||
|
||
|
||
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 SparseEmbeddingResponse(BaseModel):
|
||
vectors: list[SparseVector]
|
||
|
||
# Метадата чанков в Qdrant'e, по которой вы можете фильтровать
|
||
class ChunkMetadata(BaseModel):
|
||
chat_name: str
|
||
chat_type: str # channel, group, private, thread
|
||
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):
|
||
app.state.http = httpx.AsyncClient()
|
||
app.state.qdrant = AsyncQdrantClient(
|
||
url=QDRANT_URL,
|
||
api_key=API_KEY,
|
||
)
|
||
try:
|
||
yield
|
||
finally:
|
||
await app.state.http.aclose()
|
||
await app.state.qdrant.close()
|
||
|
||
|
||
app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan)
|
||
|
||
|
||
# Внутри шаблона dense и rerank берутся из внешних HTTP endpoint'ов,
|
||
# которые предоставляет проверяющая система.
|
||
# Текущий код ниже — минимальный пример search pipeline.
|
||
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 как списком строк.
|
||
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_sparse(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 unique_non_empty(values: list[str | None]) -> list[str]:
|
||
result: list[str] = []
|
||
seen: set[str] = set()
|
||
|
||
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.DatetimeRange(
|
||
gte=question.date_range.from_,
|
||
lte=question.date_range.to,
|
||
),
|
||
)
|
||
)
|
||
|
||
return 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] = []
|
||
|
||
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=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,
|
||
)
|
||
|
||
if not response.points:
|
||
return None
|
||
|
||
return response.points
|
||
|
||
|
||
def deduplicate_points(points: list[Any]) -> list[Any]:
|
||
unique_points: list[Any] = []
|
||
seen_ids: set[str] = set()
|
||
|
||
for point in points:
|
||
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)
|
||
unique_points.append(point)
|
||
|
||
return unique_points
|
||
|
||
|
||
def extract_point_score(point: Any) -> float:
|
||
score = getattr(point, "score", 0.0)
|
||
if score is None:
|
||
return 0.0
|
||
return float(score)
|
||
|
||
|
||
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 []
|
||
|
||
# Rerank endpoint возвращает score для пары query -> candidate text.
|
||
response = await client.post(
|
||
RERANKER_URL,
|
||
**get_upstream_request_kwargs(),
|
||
json={
|
||
"model": RERANKER_MODEL,
|
||
"encoding_format": "float",
|
||
"text_1": label,
|
||
"text_2": targets,
|
||
},
|
||
)
|
||
response.raise_for_status()
|
||
|
||
payload = response.json()
|
||
data = payload.get("data") or []
|
||
|
||
return [float(sample["score"]) for sample in data]
|
||
|
||
|
||
async def rerank_points(
|
||
client: httpx.AsyncClient,
|
||
query: str,
|
||
points: list[Any],
|
||
) -> list[tuple[Any, float]]:
|
||
rerank_candidates = points[:RERANK_LIMIT]
|
||
tail_candidates = 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 [(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),
|
||
key=lambda item: item[0],
|
||
reverse=True,
|
||
)
|
||
]
|
||
tail_with_scores = [
|
||
(point, extract_point_score(point))
|
||
for point in tail_candidates
|
||
]
|
||
|
||
return reranked_candidates + tail_with_scores
|
||
|
||
|
||
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 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, first_seen_rank
|
||
|
||
|
||
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], first_seen_rank[item[0]]),
|
||
)
|
||
return [message_id for message_id, _ in sorted_items[:limit]]
|
||
|
||
|
||
# Ваш сервис должен имплементировать оба этих метода
|
||
@app.get("/health")
|
||
async def health() -> dict[str, str]:
|
||
return {"status": "ok"}
|
||
|
||
|
||
@app.post("/search", response_model=SearchAPIResponse)
|
||
async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
|
||
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")
|
||
|
||
client: httpx.AsyncClient = app.state.http
|
||
qdrant: AsyncQdrantClient = app.state.qdrant
|
||
|
||
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 not best_points:
|
||
return SearchAPIResponse(results=[])
|
||
|
||
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)]
|
||
)
|
||
|
||
|
||
@app.exception_handler(Exception)
|
||
async def exception_handler(request: Request, exc: Exception) -> JSONResponse:
|
||
logger.exception(exc)
|
||
detail = str(exc) or repr(exc)
|
||
|
||
if isinstance(exc, RequestValidationError):
|
||
return JSONResponse(status_code=422, content={"detail": exc.errors()})
|
||
|
||
if isinstance(exc, HTTPException):
|
||
return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})
|
||
|
||
return JSONResponse(status_code=500, content={"detail": detail})
|
||
|
||
|
||
def main() -> None:
|
||
import uvicorn
|
||
|
||
uvicorn.run(
|
||
"main:app",
|
||
host=HOST,
|
||
port=PORT,
|
||
reload=False,
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|