сделал первые 3 задачи из todo_people.md
This commit is contained in:
parent
3ea1cd3066
commit
97c2e1710c
1 changed files with 125 additions and 7 deletions
132
search/main.py
132
search/main.py
|
|
@ -270,6 +270,111 @@ async def qdrant_search(
|
||||||
return response.points
|
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 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_message_ids(point: Any) -> list[str]:
|
def extract_message_ids(point: Any) -> list[str]:
|
||||||
payload = point.payload or {}
|
payload = point.payload or {}
|
||||||
metadata = payload.get("metadata") or {}
|
metadata = payload.get("metadata") or {}
|
||||||
|
|
@ -334,18 +439,31 @@ 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:
|
||||||
query = payload.question.text.strip()
|
queries = collect_query_variants(payload.question)
|
||||||
if not query:
|
if not queries:
|
||||||
raise HTTPException(status_code=400, detail="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
|
||||||
|
|
||||||
dense_vector = await embed_dense(client, query)
|
all_points: list[Any] = []
|
||||||
sparse_vector = await embed_sparse(query)
|
for query_variant in queries:
|
||||||
best_points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question)
|
dense_vector = await embed_dense(client, query_variant)
|
||||||
|
sparse_vector = await embed_sparse(query_variant)
|
||||||
|
points = await qdrant_search(qdrant, dense_vector, sparse_vector, payload.question)
|
||||||
|
if points:
|
||||||
|
all_points.extend(list(points))
|
||||||
|
|
||||||
if best_points is None:
|
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=[])
|
return SearchAPIResponse(results=[])
|
||||||
|
|
||||||
best_points = await rerank_points(client, query, list(best_points))
|
best_points = await rerank_points(client, query, list(best_points))
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue