560 lines
16 KiB
Python
560 lines
16 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 = 10
|
||
SPRASE_PREFETCH_K = 30
|
||
RETRIEVE_K = 20
|
||
RERANK_LIMIT = 10
|
||
FINAL_TOP_K = 50
|
||
|
||
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()],
|
||
)
|
||
|
||
# ПЕРЕПИСАТЬ
|
||
|
||
async def qdrant_search(
|
||
client: AsyncQdrantClient,
|
||
dense_vector: list[float],
|
||
sparse_vector: SparseVector,
|
||
question_data: Question
|
||
) -> Any | None:
|
||
must_conditions: []
|
||
|
||
# Фильтр по диапазону дат (поле metadata.start в Qdrant) [cite: 147, 148, 175]
|
||
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_
|
||
)
|
||
)
|
||
)
|
||
|
||
# Фильтр по автору вопроса (поле metadata.participants) [cite: 161, 163]
|
||
if question_data.asker:
|
||
must_conditions.append(
|
||
models.FieldCondition(
|
||
key="metadata.participants",
|
||
match=models.MatchValue(value=question_data.asker)
|
||
)
|
||
)
|
||
|
||
# Создаем итоговый объект фильтра, если есть условия
|
||
search_filter = models.Filter(must=must_conditions) if must_conditions else None
|
||
|
||
response = await client.query_points(
|
||
collection_name=QDRANT_COLLECTION_NAME,
|
||
prefetch=[
|
||
models.Prefetch(
|
||
query=dense_vector,
|
||
using=QDRANT_DENSE_VECTOR_NAME,
|
||
limit=DENSE_PREFETCH_K,
|
||
filter=search_filter,
|
||
),
|
||
models.Prefetch(
|
||
query=models.SparseVector(
|
||
indices=sparse_vector.indices,
|
||
values=sparse_vector.values,
|
||
),
|
||
using=QDRANT_SPARSE_VECTOR_NAME,
|
||
limit=SPRASE_PREFETCH_K,
|
||
filter=search_filter,
|
||
),
|
||
],
|
||
query=models.FusionQuery(fusion=models.Fusion.RRF),
|
||
limit=RETRIEVE_K,
|
||
with_payload=True,
|
||
)
|
||
|
||
if not response.points:
|
||
return None
|
||
|
||
return response.points
|
||
|
||
|
||
async def qdrant_search_dense_only(
|
||
client: AsyncQdrantClient,
|
||
dense_vector: list[float],
|
||
question_data: Question,
|
||
) -> Any | None:
|
||
must_conditions: []
|
||
|
||
if question_data.date_range:
|
||
must_conditions.append(
|
||
models.FieldCondition(
|
||
key="metadata.start",
|
||
range=models.Range(
|
||
gte=question_data.date_range.from_,
|
||
lte=question_data.date_range.to_,
|
||
),
|
||
)
|
||
)
|
||
|
||
if question_data.asker:
|
||
must_conditions.append(
|
||
models.FieldCondition(
|
||
key="metadata.participants",
|
||
match=models.MatchValue(value=question_data.asker),
|
||
)
|
||
)
|
||
|
||
search_filter = models.Filter(must=must_conditions) if must_conditions else None
|
||
|
||
response = await client.query_points(
|
||
collection_name=QDRANT_COLLECTION_NAME,
|
||
prefetch=[
|
||
models.Prefetch(
|
||
query=dense_vector,
|
||
using=QDRANT_DENSE_VECTOR_NAME,
|
||
limit=DENSE_PREFETCH_K,
|
||
filter=search_filter,
|
||
),
|
||
],
|
||
query=models.FusionQuery(fusion=models.Fusion.RRF),
|
||
limit=RETRIEVE_K,
|
||
with_payload=True,
|
||
)
|
||
|
||
if not response.points:
|
||
return None
|
||
|
||
return response.points
|
||
|
||
|
||
def collect_query_variants(question: Question) -> list[str]:
|
||
variants: list[str] = []
|
||
seen: set[str] = set()
|
||
|
||
def add_query(text: str | None) -> None:
|
||
if text is None:
|
||
return
|
||
normalized = text.strip()
|
||
if not normalized:
|
||
return
|
||
if normalized in seen:
|
||
return
|
||
seen.add(normalized)
|
||
variants.append(normalized)
|
||
|
||
add_query(question.search_text)
|
||
add_query(question.text)
|
||
|
||
for variant in question.variants or []:
|
||
add_query(variant)
|
||
|
||
return variants
|
||
|
||
|
||
def collect_hyde_queries(question: Question, base_queries: list[str]) -> list[str]:
|
||
hyde_queries: list[str] = []
|
||
seen: set[str] = set(base_queries)
|
||
|
||
for hyde_query in question.hyde or []:
|
||
normalized = hyde_query.strip()
|
||
if not normalized:
|
||
continue
|
||
if normalized in seen:
|
||
continue
|
||
seen.add(normalized)
|
||
hyde_queries.append(normalized)
|
||
|
||
return hyde_queries
|
||
|
||
|
||
def build_sparse_query_text(question: Question, fallback_query: str) -> str:
|
||
keywords: list[str] = []
|
||
seen: set[str] = set()
|
||
|
||
for keyword in question.keywords or []:
|
||
normalized = keyword.strip()
|
||
if not normalized:
|
||
continue
|
||
if normalized in seen:
|
||
continue
|
||
seen.add(normalized)
|
||
keywords.append(normalized)
|
||
|
||
if keywords:
|
||
return " ".join(keywords)
|
||
|
||
return fallback_query
|
||
|
||
|
||
def deduplicate_points(points: list[Any]) -> list[Any]:
|
||
unique_points: list[Any] = []
|
||
seen_ids: set[str] = set()
|
||
|
||
for point in points:
|
||
point_id = str(getattr(point, "id", ""))
|
||
if not point_id:
|
||
continue
|
||
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 = [point.payload.get("page_content") for point in rerank_candidates]
|
||
scores = await get_rerank_scores(client, query, rerank_targets)
|
||
|
||
reranked_candidates = [
|
||
(point, float(score))
|
||
for score, point in sorted(
|
||
zip(scores, rerank_candidates, strict=True),
|
||
key=lambda item: item[0],
|
||
reverse=True,
|
||
)
|
||
]
|
||
tail_with_scores = [
|
||
(point, extract_point_score(point))
|
||
for _, point in sorted(
|
||
[(extract_point_score(point), point) for point in tail_candidates],
|
||
key=lambda item: item[0],
|
||
reverse=True,
|
||
)
|
||
]
|
||
|
||
return reranked_candidates + tail_with_scores
|
||
|
||
|
||
def aggregate_message_scores(scored_points: list[tuple[Any, float]]) -> dict[str, float]:
|
||
aggregated_scores: dict[str, float] = {}
|
||
|
||
for point, point_score in 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
|
||
|
||
return aggregated_scores
|
||
|
||
|
||
def select_top_message_ids(aggregated_scores: dict[str, float], limit: int) -> list[str]:
|
||
sorted_items = sorted(
|
||
aggregated_scores.items(),
|
||
key=lambda item: (-item[1], 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:
|
||
queries = collect_query_variants(payload.question)
|
||
if not queries:
|
||
raise HTTPException(status_code=400, detail="question.search_text or question.text is required")
|
||
|
||
hyde_queries = collect_hyde_queries(payload.question, queries)
|
||
query = queries[0]
|
||
client: httpx.AsyncClient = app.state.http
|
||
qdrant: AsyncQdrantClient = app.state.qdrant
|
||
|
||
all_points: list[Any] = []
|
||
for query_variant in queries:
|
||
dense_vector = await embed_dense(client, query_variant)
|
||
sparse_query_text = build_sparse_query_text(payload.question, query_variant)
|
||
sparse_vector = await embed_sparse(sparse_query_text)
|
||
points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question)
|
||
if points:
|
||
all_points.extend(list(points))
|
||
|
||
for hyde_query in hyde_queries:
|
||
hyde_dense_vector = await embed_dense(client, hyde_query)
|
||
hyde_points = await qdrant_search_dense_only(qdrant, hyde_dense_vector, payload.question)
|
||
if hyde_points:
|
||
all_points.extend(list(hyde_points))
|
||
|
||
best_points = deduplicate_points(all_points)
|
||
if not best_points:
|
||
return SearchAPIResponse(results=[])
|
||
|
||
scored_points = await rerank_points(client, query, list(best_points))
|
||
aggregated_scores = aggregate_message_scores(scored_points)
|
||
message_ids = select_top_message_ids(aggregated_scores, FINAL_TOP_K)
|
||
|
||
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()
|