Compare commits
27 commits
19af82ba11
...
4800e25dd6
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4800e25dd6 | ||
|
|
c0f2d52f70 | ||
|
|
92cde65e42 | ||
|
|
19fec2361e | ||
|
|
4ecba7d35a | ||
|
|
c9083c285f | ||
|
|
f6e53758f7 | ||
|
|
f6d66854b9 | ||
|
|
2bb595e452 | ||
|
|
57a5229c95 | ||
|
|
775c874399 | ||
|
|
c0506c49aa | ||
|
|
3689d4f3ec | ||
|
|
6831a5d149 | ||
|
|
4822bbb24b | ||
|
|
84eec2321e | ||
|
|
fec71a98b9 | ||
|
|
5d50a219bf | ||
|
|
68e2dc22c4 | ||
|
|
878971bb57 | ||
|
|
7e40bb6e17 | ||
|
|
6a25927813 | ||
|
|
f16df83601 | ||
|
|
72991ff71a | ||
|
|
4bff9e5ea2 | ||
|
|
1104ed936c | ||
|
|
1f976cf297 |
23 changed files with 1109 additions and 164 deletions
0
.codex
Normal file
0
.codex
Normal file
366
doc/curl_api_test.md
Normal file
366
doc/curl_api_test.md
Normal file
|
|
@ -0,0 +1,366 @@
|
|||
# Curl API Test
|
||||
|
||||
## Sources
|
||||
|
||||
- Canonical contracts: `doc/ТЗ_на_хакатон_Индексация_и_поиск_по_сообщениям.pdf`
|
||||
- Runnable examples and local launch notes: `README.md`
|
||||
- Actual local wiring: `docker-compose.yml`
|
||||
|
||||
PDF gives the strict request/response schemas for `POST /index`, `POST /sparse_embedding`, and `POST /search`.
|
||||
`README.md` adds ready curl examples for the minimal requests.
|
||||
This file normalizes both into checks against the current local compose stack.
|
||||
|
||||
## Compose Wiring
|
||||
|
||||
- `index`: `http://localhost:8001`
|
||||
- `search`: `http://localhost:8002`
|
||||
- `qdrant`: `http://localhost:6334`
|
||||
- Inside compose, services use `QDRANT_URL=http://qdrant:6333`
|
||||
- Collection name from `.env`: `evaluation`
|
||||
- Vector names from `.env`: `dense` and `sparse`
|
||||
|
||||
Note: current `docker-compose.yml` publishes Qdrant as `6334:6333`, while `README.md` still says `localhost:6333`. For local checks in this repo state, use `localhost:6334`.
|
||||
|
||||
## Extracted API Requests
|
||||
|
||||
### `GET /health`
|
||||
|
||||
Both services must answer `200 OK`.
|
||||
|
||||
```bash
|
||||
curl -sS http://localhost:8001/health
|
||||
curl -sS http://localhost:8002/health
|
||||
```
|
||||
|
||||
Expected shape:
|
||||
|
||||
```json
|
||||
{"status":"ok"}
|
||||
```
|
||||
|
||||
### `POST /index`
|
||||
|
||||
Schema from the PDF:
|
||||
|
||||
- body root: `data`
|
||||
- `data.chat`
|
||||
- `data.overlap_messages[]`
|
||||
- `data.new_messages[]`
|
||||
|
||||
Runnable request:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8001/index \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"data": {
|
||||
"chat": {
|
||||
"id": "chat-1",
|
||||
"name": "Go Nova",
|
||||
"sn": "chat-1@chat.agent",
|
||||
"type": "channel",
|
||||
"is_public": true
|
||||
},
|
||||
"overlap_messages": [
|
||||
{
|
||||
"id": "1",
|
||||
"time": 1710000000,
|
||||
"text": "Обсуждаем релиз Go",
|
||||
"sender_id": "u1",
|
||||
"file_snippets": "",
|
||||
"parts": [],
|
||||
"mentions": [],
|
||||
"member_event": null,
|
||||
"is_system": false,
|
||||
"is_hidden": false,
|
||||
"is_forward": false,
|
||||
"is_quote": false
|
||||
}
|
||||
],
|
||||
"new_messages": [
|
||||
{
|
||||
"id": "2",
|
||||
"time": 1710000060,
|
||||
"text": "Релиз Go перенесли на следующую неделю",
|
||||
"sender_id": "u2",
|
||||
"file_snippets": "",
|
||||
"parts": [],
|
||||
"mentions": [],
|
||||
"member_event": null,
|
||||
"is_system": false,
|
||||
"is_hidden": false,
|
||||
"is_forward": false,
|
||||
"is_quote": false
|
||||
}
|
||||
]
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
Observed response:
|
||||
|
||||
```json
|
||||
{
|
||||
"results": [
|
||||
{
|
||||
"page_content": "u1: Обсуждаем релиз Go\nu2: Релиз Go перенесли на следующую неделю",
|
||||
"dense_content": "[2024-03-09 16:00] sender:u1\nОбсуждаем релиз Go\n[2024-03-09 16:01] sender:u2\nРелиз Go перенесли на следующую неделю",
|
||||
"sparse_content": "u1 Обсуждаем релиз Go u2 Релиз Go перенесли на следующую неделю",
|
||||
"message_ids": ["2"]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Note: overlap messages are used as context, but are not included in returned `message_ids`.
|
||||
|
||||
### `POST /sparse_embedding`
|
||||
|
||||
Schema from the PDF:
|
||||
|
||||
- body root: `texts: string[]`
|
||||
|
||||
Runnable request:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8001/sparse_embedding \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"texts": [
|
||||
"Релиз Go перенесли на следующую неделю",
|
||||
"VK GPT обсуждали в отдельном чате"
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Observed response:
|
||||
|
||||
```json
|
||||
{
|
||||
"vectors": [
|
||||
{
|
||||
"indices": [275068001, 108710752, 842257583, 1159207840, 2129888840, 703082301],
|
||||
"values": [1.6652868125369606, 1.6652868125369606, 1.6652868125369606, 1.6652868125369606, 1.6652868125369606, 1.6652868125369606]
|
||||
},
|
||||
{
|
||||
"indices": [73209461, 751565418, 59863655, 1856729543, 2036701913, 1943620510],
|
||||
"values": [1.6652868125369606, 1.6652868125369606, 1.6652868125369606, 1.6652868125369606, 1.6652868125369606, 1.6652868125369606]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### `POST /search`
|
||||
|
||||
Minimal request from `README.md`:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8002/search \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"question": {
|
||||
"text": "Что писали про релиз Go?"
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
Full schema from the PDF:
|
||||
|
||||
```json
|
||||
{
|
||||
"question": {
|
||||
"text": "Что писали про релиз Go?",
|
||||
"asker": "u2",
|
||||
"asked_on": "2024-03-09",
|
||||
"variants": ["релиз go перенесли?", "обсуждение релиза go"],
|
||||
"hyde": ["В чате пишут, что релиз Go перенесли на следующую неделю."],
|
||||
"keywords": ["релиз", "Go", "перенесли"],
|
||||
"entities": {
|
||||
"people": ["u2"],
|
||||
"emails": [],
|
||||
"documents": [],
|
||||
"names": ["Go"],
|
||||
"links": []
|
||||
},
|
||||
"date_mentions": ["следующая неделя", "2024-03-09"],
|
||||
"date_range": {
|
||||
"from": "2024-03-09T00:00:00Z",
|
||||
"to": "2024-03-10T00:00:00Z"
|
||||
},
|
||||
"search_text": "релиз Go перенесли на следующую неделю"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Checks Run
|
||||
|
||||
### 1. Health checks
|
||||
|
||||
Commands:
|
||||
|
||||
```bash
|
||||
curl -sS http://localhost:8001/health
|
||||
curl -sS http://localhost:8002/health
|
||||
```
|
||||
|
||||
Observed:
|
||||
|
||||
```json
|
||||
{"status":"ok"}
|
||||
{"status":"ok"}
|
||||
```
|
||||
|
||||
### 2. Qdrant collection exists, but starts empty
|
||||
|
||||
Command:
|
||||
|
||||
```bash
|
||||
curl -sS http://localhost:6334/collections/evaluation
|
||||
```
|
||||
|
||||
Observed before manual insert:
|
||||
|
||||
- `points_count: 0`
|
||||
- `indexed_vectors_count: 0`
|
||||
|
||||
This matches the README note that local compose creates the collection, but the template flow does not automatically upsert `/index` output into Qdrant.
|
||||
|
||||
### 3. `/index` works
|
||||
|
||||
Observed:
|
||||
|
||||
- HTTP request completed successfully
|
||||
- service returned one chunk
|
||||
- returned fields match the contract: `page_content`, `dense_content`, `sparse_content`, `message_ids`
|
||||
|
||||
### 4. `/sparse_embedding` works
|
||||
|
||||
Observed:
|
||||
|
||||
- HTTP request completed successfully
|
||||
- response returned `vectors[]`
|
||||
- each vector contains `indices[]` and `values[]`
|
||||
|
||||
### 5. `/search` on an empty collection returns an empty result
|
||||
|
||||
Command:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8002/search \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"question":{"text":"Что писали про релиз Go?"}}'
|
||||
```
|
||||
|
||||
Observed:
|
||||
|
||||
```json
|
||||
{"results":[]}
|
||||
```
|
||||
|
||||
This is expected while `evaluation` has no points.
|
||||
|
||||
### 6. Manual Qdrant upsert for end-to-end smoke test
|
||||
|
||||
To verify `/search` end-to-end, I inserted one synthetic point into local Qdrant with:
|
||||
|
||||
- point id `1001`
|
||||
- dummy dense vector of size `1024`
|
||||
- sparse vector under field `sparse`
|
||||
- payload containing `page_content` and `metadata.message_ids=["2"]`
|
||||
|
||||
Command:
|
||||
|
||||
```bash
|
||||
vec=$(awk 'BEGIN{for(i=0;i<1024;i++) printf "%s%d", (i?",":""), (i==0)}')
|
||||
curl -sS -X PUT 'http://localhost:6334/collections/evaluation/points?wait=true' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"points\":[{\"id\":1001,\"vector\":{\"dense\":[${vec}],\"sparse\":{\"indices\":[1],\"values\":[1.0]}},\"payload\":{\"page_content\":\"u1: Обсуждаем релиз Go\\nu2: Релиз Go перенесли на следующую неделю\",\"metadata\":{\"message_ids\":[\"2\"],\"participants\":[\"u1\",\"u2\"],\"start\":\"2024-03-09T16:00:00Z\",\"end\":\"2024-03-09T16:01:00Z\",\"chat_id\":\"chat-1\",\"chat_name\":\"Go Nova\",\"chat_type\":\"channel\",\"chat_sn\":\"chat-1@chat.agent\"}}}]}"
|
||||
```
|
||||
|
||||
Observed:
|
||||
|
||||
```json
|
||||
{"result":{"operation_id":0,"status":"completed"},"status":"ok","time":0.008234969}
|
||||
```
|
||||
|
||||
Collection state after insert:
|
||||
|
||||
- `points_count: 1`
|
||||
- `indexed_vectors_count: 1`
|
||||
|
||||
### 7. `/search` works after one point is present
|
||||
|
||||
Minimal request:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8002/search \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"question":{"text":"Что писали про релиз Go?"}}'
|
||||
```
|
||||
|
||||
Observed:
|
||||
|
||||
```json
|
||||
{"results":[{"message_ids":["2"]}]}
|
||||
```
|
||||
|
||||
Enriched request without `date_range`:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8002/search \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"question": {
|
||||
"text": "Что писали про релиз Go?",
|
||||
"asker": "u2",
|
||||
"asked_on": "2024-03-09",
|
||||
"variants": ["релиз go перенесли?", "обсуждение релиза go"],
|
||||
"hyde": ["В чате пишут, что релиз Go перенесли на следующую неделю."],
|
||||
"keywords": ["релиз", "Go", "перенесли"],
|
||||
"entities": {
|
||||
"people": ["u2"],
|
||||
"emails": [],
|
||||
"documents": [],
|
||||
"names": ["Go"],
|
||||
"links": []
|
||||
},
|
||||
"date_mentions": ["следующая неделя", "2024-03-09"],
|
||||
"search_text": "релиз Go перенесли на следующую неделю"
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
Observed:
|
||||
|
||||
```json
|
||||
{"results":[{"message_ids":["2"]}]}
|
||||
```
|
||||
|
||||
### 8. Defect: `date_range` request currently fails
|
||||
|
||||
The full PDF-shaped request with ISO timestamps in `question.date_range` does not work in the current implementation.
|
||||
|
||||
Observed:
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": "2 validation errors for Range\ngte\n Input should be a valid number, unable to parse string as a number [type=float_parsing, input_value='2024-03-09T00:00:00Z', input_type=str]\n For further information visit https://errors.pydantic.dev/2.12/v/float_parsing\nlte\n Input should be a valid number, unable to parse string as a number [type=float_parsing, input_value='2024-03-10T00:00:00Z', input_type=str]\n For further information visit https://errors.pydantic.dev/2.12/v/float_parsing"
|
||||
}
|
||||
```
|
||||
|
||||
Interpretation:
|
||||
|
||||
- the public request schema accepts ISO date strings
|
||||
- current `search` code tries to pass them into a numeric `qdrant_client.models.Range`
|
||||
- so `date_range` is a real runtime bug in the current local build
|
||||
|
||||
## Bottom Line
|
||||
|
||||
- `index /health`: OK
|
||||
- `search /health`: OK
|
||||
- `POST /index`: OK
|
||||
- `POST /sparse_embedding`: OK
|
||||
- `POST /search` on empty collection: OK, returns empty list
|
||||
- `POST /search` after one test point is inserted: OK
|
||||
- `POST /search` with enriched request excluding `date_range`: OK
|
||||
- `POST /search` with `date_range` from the PDF schema: FAILS in current implementation
|
||||
|
|
@ -2,7 +2,7 @@ services:
|
|||
qdrant:
|
||||
image: qdrant/qdrant:v1.14.1
|
||||
ports:
|
||||
- "6333:6333"
|
||||
- "6334:6333"
|
||||
|
||||
qdrant-init:
|
||||
image: curlimages/curl:8.12.1
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ COPY main.py .
|
|||
|
||||
ENV HOST=0.0.0.0
|
||||
ENV PORT=8000
|
||||
ENV CHUNK_SIZE=10
|
||||
ENV FASTEMBED_CACHE_PATH=/models/fastembed
|
||||
ENV HF_HOME=/models/huggingface
|
||||
|
||||
|
|
|
|||
|
|
@ -1,24 +1,26 @@
|
|||
LOGIN ?=
|
||||
PASSWORD ?=
|
||||
TEAM_ID ?=
|
||||
TEAM_ID ?= 35230
|
||||
DOCKER_REGISTRY_URL ?= 83.166.249.64:5000
|
||||
PORT ?= 8000
|
||||
|
||||
IMAGE = $(DOCKER_REGISTRY_URL)/$(TEAM_ID)/index-service:latest
|
||||
|
||||
.PHONY: login build run push
|
||||
.PHONY: login build run push release
|
||||
|
||||
login:
|
||||
@: $(if $(LOGIN),,$(error LOGIN is required for make login))
|
||||
@: $(if $(PASSWORD),,$(error PASSWORD is required for make login))
|
||||
@: $(if $(LOGIN),,$(error LOGIN is required))
|
||||
@: $(if $(PASSWORD),,$(error PASSWORD is required))
|
||||
docker login $(DOCKER_REGISTRY_URL) -u $(LOGIN) -p $(PASSWORD)
|
||||
|
||||
build:
|
||||
@: $(if $(TEAM_ID),,$(error TEAM_ID is required for make build))
|
||||
docker build -t $(IMAGE) ./
|
||||
docker build --platform linux/amd64 -t $(IMAGE) ./
|
||||
|
||||
run: build
|
||||
docker run --rm -p $(PORT):8000 $(IMAGE)
|
||||
|
||||
push: login build
|
||||
push:
|
||||
docker push $(IMAGE)
|
||||
|
||||
release: build push
|
||||
@echo "index-service pushed → $(IMAGE)"
|
||||
|
|
|
|||
|
|
@ -1,17 +1,17 @@
|
|||
"""Message-based chunking with window by count, length, and time gap."""
|
||||
|
||||
from .cleaning import CleanedMessage, clean_message
|
||||
from .rendering import render_dense_content, render_page_content, render_sparse_content
|
||||
from .schemas import IndexAPIItem, Message
|
||||
from cleaning import CleanedMessage, clean_message
|
||||
from rendering import render_dense_content, render_page_content, render_sparse_content
|
||||
from index_schemas import IndexAPIItem, Message
|
||||
|
||||
WINDOW_MAX_MESSAGES = 10
|
||||
WINDOW_MAX_CHARS = 2048
|
||||
WINDOW_MAX_MESSAGES = 5
|
||||
WINDOW_MAX_CHARS = 512
|
||||
TIME_GAP_SECONDS = 3600
|
||||
OVERLAP_MESSAGES = 3
|
||||
OVERLAP_MESSAGES = 2
|
||||
|
||||
|
||||
def _clean_all(messages: list[Message]) -> list[CleanedMessage]:
|
||||
cleaned = [clean_message(m) for m in messages]
|
||||
cleaned = [clean_message(m) for m in messages if not m.is_system and not m.is_hidden]
|
||||
return [c for c in cleaned if not c.is_empty]
|
||||
|
||||
|
||||
|
|
|
|||
204
index/main.py
204
index/main.py
|
|
@ -1,24 +1,192 @@
|
|||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from .chunking import build_chunks
|
||||
from .schemas import IndexAPIRequest, IndexAPIResponse, SparseEmbeddingRequest
|
||||
from .sparse import embed_sparse_texts
|
||||
from pydantic import BaseModel
|
||||
|
||||
HOST = os.getenv("HOST", "0.0.0.0")
|
||||
PORT = int(os.getenv("PORT", "8004"))
|
||||
PORT = int(os.getenv("PORT", "8000"))
|
||||
UVICORN_WORKERS = 8
|
||||
|
||||
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
||||
logger = logging.getLogger("index-service")
|
||||
|
||||
app = FastAPI(title="Index Service", version="0.2.0")
|
||||
|
||||
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"
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
|
|
@ -30,16 +198,38 @@ 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,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
|
||||
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.to_thread(embed_sparse_texts, payload.texts)
|
||||
return {"vectors": [{"indices": v.indices, "values": v.values} for v in vectors]}
|
||||
return {"vectors": vectors}
|
||||
|
||||
|
||||
@app.exception_handler(Exception)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
import datetime
|
||||
|
||||
from .cleaning import CleanedMessage
|
||||
from cleaning import CleanedMessage
|
||||
|
||||
|
||||
def _format_time(ts: int) -> str:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import logging
|
|||
import os
|
||||
from functools import lru_cache
|
||||
|
||||
from .schemas import SparseVector
|
||||
from index_schemas import SparseVector
|
||||
|
||||
SPARSE_MODEL_NAME = "Qdrant/bm25"
|
||||
FASTEMBED_CACHE_PATH = "/models/fastembed"
|
||||
|
|
|
|||
3
kredit.md
Normal file
3
kredit.md
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
team_id: 35230
|
||||
vk login: 56aa86799bb9edc4
|
||||
vk password: edd89cea9ed0734d00ba6904cf7475d7
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
LOGIN ?=
|
||||
PASSWORD ?=
|
||||
TEAM_ID ?=
|
||||
TEAM_ID ?= 35230
|
||||
DOCKER_REGISTRY_URL ?= 83.166.249.64:5000
|
||||
PORT ?= 8000
|
||||
QDRANT_URL ?=
|
||||
|
|
@ -16,16 +16,15 @@ REQUIRED_RUN_VARS := QDRANT_URL EMBEDDINGS_DENSE_URL API_KEY RERANKER_URL
|
|||
|
||||
IMAGE = $(DOCKER_REGISTRY_URL)/$(TEAM_ID)/search-service:latest
|
||||
|
||||
.PHONY: login build run push check-run-env
|
||||
.PHONY: login build run push release
|
||||
|
||||
login:
|
||||
@: $(if $(LOGIN),,$(error LOGIN is required for make login))
|
||||
@: $(if $(PASSWORD),,$(error PASSWORD is required for make login))
|
||||
@: $(if $(LOGIN),,$(error LOGIN is required))
|
||||
@: $(if $(PASSWORD),,$(error PASSWORD is required))
|
||||
docker login $(DOCKER_REGISTRY_URL) -u $(LOGIN) -p $(PASSWORD)
|
||||
|
||||
build:
|
||||
@: $(if $(TEAM_ID),,$(error TEAM_ID is required for make build))
|
||||
docker build -t $(IMAGE) ./
|
||||
docker build --platform linux/amd64 -t $(IMAGE) ./
|
||||
|
||||
run: build
|
||||
@: $(foreach var,$(REQUIRED_RUN_VARS),$(if $($(var)),,$(error $(var) is required for make run)))
|
||||
|
|
@ -41,5 +40,8 @@ run: build
|
|||
-e QDRANT_SPARSE_VECTOR_NAME=$(QDRANT_SPARSE_VECTOR_NAME) \
|
||||
$(IMAGE)
|
||||
|
||||
push: login build
|
||||
push:
|
||||
docker push $(IMAGE)
|
||||
|
||||
release: build push
|
||||
@echo "search-service pushed → $(IMAGE)"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from typing import Any
|
||||
|
||||
from .config import TOP_K
|
||||
from .retrieval import extract_message_ids
|
||||
from config import TOP_K
|
||||
from retrieval import extract_message_ids
|
||||
|
||||
|
||||
def aggregate_message_ids(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
464
search/main.py
464
search/main.py
|
|
@ -2,56 +2,157 @@ 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
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
from pydantic import BaseModel, Field
|
||||
from qdrant_client import AsyncQdrantClient, models
|
||||
|
||||
from .aggregation import aggregate_message_ids
|
||||
from .config import (
|
||||
API_KEY,
|
||||
HOST,
|
||||
HTTP_MAX_RETRIES,
|
||||
HTTP_TIMEOUT,
|
||||
PORT,
|
||||
QDRANT_URL,
|
||||
logger,
|
||||
validate_required_env,
|
||||
)
|
||||
from .query_builder import (
|
||||
build_extra_dense_queries,
|
||||
build_primary_query,
|
||||
build_sparse_query,
|
||||
embed_dense,
|
||||
embed_dense_multi,
|
||||
embed_sparse,
|
||||
)
|
||||
from .rerank import rerank_points
|
||||
from .retrieval import qdrant_search
|
||||
from .schemas import SearchAPIItem, SearchAPIRequest, SearchAPIResponse, SparseVector
|
||||
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")
|
||||
|
||||
|
||||
async def _embed_dense_with_retry(client: httpx.AsyncClient, text: str) -> list[float]:
|
||||
last_exc: Exception | None = None
|
||||
for attempt in range(HTTP_MAX_RETRIES + 1):
|
||||
try:
|
||||
return await embed_dense(client, text)
|
||||
except (httpx.TransportError, httpx.HTTPStatusError) as exc:
|
||||
if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code < 500:
|
||||
raise
|
||||
last_exc = exc
|
||||
if attempt < HTTP_MAX_RETRIES:
|
||||
await asyncio.sleep(0.5 * (attempt + 1))
|
||||
raise RuntimeError(f"Dense embedding failed after retries: {last_exc}")
|
||||
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 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):
|
||||
validate_required_env()
|
||||
app.state.http = httpx.AsyncClient(timeout=HTTP_TIMEOUT)
|
||||
app.state.qdrant = AsyncQdrantClient(url=QDRANT_URL, api_key=API_KEY)
|
||||
app.state.http = httpx.AsyncClient()
|
||||
app.state.qdrant = AsyncQdrantClient(
|
||||
url=QDRANT_URL,
|
||||
api_key=API_KEY,
|
||||
)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
|
|
@ -59,7 +160,225 @@ async def lifespan(app: FastAPI):
|
|||
await app.state.qdrant.close()
|
||||
|
||||
|
||||
app = FastAPI(title="Search Service", version="0.2.0", lifespan=lifespan)
|
||||
app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan)
|
||||
|
||||
DENSE_PREFETCH_K = 120
|
||||
SPARSE_PREFETCH_K = 200
|
||||
RETRIEVE_K = 150
|
||||
RERANK_LIMIT = 35
|
||||
KEYWORD_BOOST_EXTRA = 10
|
||||
|
||||
|
||||
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:
|
||||
q = question.search_text.strip() if question.search_text else question.text.strip()
|
||||
return q
|
||||
|
||||
|
||||
def build_sparse_query(question: Question) -> str:
|
||||
base = question.search_text.strip() if question.search_text else question.text.strip()
|
||||
parts = [base]
|
||||
if question.keywords:
|
||||
parts.extend(question.keywords)
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def _build_keyword_set(question: Question) -> list[str]:
|
||||
tokens: list[str] = []
|
||||
if question.keywords:
|
||||
tokens.extend(kw.lower() for kw in question.keywords if kw)
|
||||
if question.entities:
|
||||
for field in (
|
||||
question.entities.people,
|
||||
question.entities.emails,
|
||||
question.entities.documents,
|
||||
question.entities.names,
|
||||
question.entities.links,
|
||||
):
|
||||
tokens.extend(e.lower() for e in (field or []) if e)
|
||||
return tokens
|
||||
|
||||
|
||||
def prefilter_for_rerank(
|
||||
points: list[Any],
|
||||
question: Question,
|
||||
) -> tuple[list[Any], list[Any]]:
|
||||
"""Select candidates for reranking: top by RRF + keyword-boosted stragglers."""
|
||||
if not points:
|
||||
return [], []
|
||||
|
||||
head = points[:RERANK_LIMIT]
|
||||
tail = points[RERANK_LIMIT:]
|
||||
|
||||
keywords = _build_keyword_set(question)
|
||||
if not keywords or not tail:
|
||||
return head, tail
|
||||
|
||||
extra: list[Any] = []
|
||||
remaining_tail: list[Any] = []
|
||||
for p in tail:
|
||||
if len(extra) >= KEYWORD_BOOST_EXTRA:
|
||||
remaining_tail.append(p)
|
||||
continue
|
||||
content = ((p.payload or {}).get("page_content") or "").lower()
|
||||
if any(kw in content for kw in keywords):
|
||||
extra.append(p)
|
||||
else:
|
||||
remaining_tail.append(p)
|
||||
|
||||
return head + extra, remaining_tail
|
||||
|
||||
|
||||
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]:
|
||||
if not points:
|
||||
return []
|
||||
targets = [point.payload.get("page_content") for point in points]
|
||||
scores = await get_rerank_scores(client, query, targets)
|
||||
|
||||
if not scores or len(scores) != len(points):
|
||||
logger.warning("Reranker unavailable or score mismatch, returning RRF order")
|
||||
return points
|
||||
|
||||
return [
|
||||
point
|
||||
for _, point in sorted(
|
||||
zip(scores, points),
|
||||
key=lambda item: item[0],
|
||||
reverse=True,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
|
|
@ -70,39 +389,62 @@ async def health() -> dict[str, str]:
|
|||
@app.post("/search", response_model=SearchAPIResponse)
|
||||
async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
|
||||
question = payload.question
|
||||
primary_query = build_primary_query(question)
|
||||
if not primary_query:
|
||||
query = question.text.strip()
|
||||
if not query:
|
||||
raise HTTPException(status_code=400, detail="question.text is required")
|
||||
|
||||
client: httpx.AsyncClient = app.state.http
|
||||
qdrant: AsyncQdrantClient = app.state.qdrant
|
||||
|
||||
extra_texts = build_extra_dense_queries(question)
|
||||
sparse_text = build_sparse_query(question)
|
||||
dense_query = build_dense_query(question)
|
||||
sparse_query = build_sparse_query(question)
|
||||
|
||||
async def _no_extra() -> list:
|
||||
return []
|
||||
dense_task = embed_dense(client, dense_query)
|
||||
sparse_task = asyncio.to_thread(lambda: embed_sparse_sync(sparse_query))
|
||||
dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
|
||||
|
||||
extra_task = embed_dense_multi(client, extra_texts) if extra_texts else _no_extra()
|
||||
primary_dense, extra_dense_vecs, sparse_vec = await asyncio.gather(
|
||||
_embed_dense_with_retry(client, primary_query),
|
||||
extra_task,
|
||||
asyncio.to_thread(embed_sparse, sparse_text),
|
||||
)
|
||||
dense_vectors = [dense_vector]
|
||||
extra_texts: list[str] = []
|
||||
raw_text = question.text.strip()
|
||||
if raw_text and raw_text != dense_query:
|
||||
extra_texts.append(raw_text)
|
||||
for v in (question.variants or []):
|
||||
q_v = v.strip()
|
||||
if q_v and q_v != dense_query and q_v not in extra_texts:
|
||||
extra_texts.append(q_v)
|
||||
for h in (question.hyde or []):
|
||||
q_h = h.strip()
|
||||
if q_h and q_h != dense_query and q_h not in extra_texts:
|
||||
extra_texts.append(q_h)
|
||||
extra_texts = extra_texts[:3]
|
||||
if extra_texts:
|
||||
try:
|
||||
extra_vecs = await embed_dense_batch(client, extra_texts)
|
||||
dense_vectors.extend(extra_vecs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Extra dense embedding failed: {e}")
|
||||
|
||||
points = await qdrant_search(
|
||||
qdrant,
|
||||
primary_dense,
|
||||
extra_dense_vecs,
|
||||
sparse_vec,
|
||||
question,
|
||||
)
|
||||
all_points = await qdrant_search(qdrant, dense_vectors, sparse_vector)
|
||||
|
||||
if not points:
|
||||
if all_points is None:
|
||||
return SearchAPIResponse(results=[])
|
||||
|
||||
reranked_head, retrieval_tail = await rerank_points(client, primary_query, points)
|
||||
message_ids = aggregate_message_ids(reranked_head, retrieval_tail)
|
||||
all_points = list(all_points)
|
||||
|
||||
rerank_pool, rerank_tail = prefilter_for_rerank(all_points, question)
|
||||
reranked = await rerank_points(client, query, rerank_pool)
|
||||
final_points = reranked + rerank_tail
|
||||
|
||||
msg_score: dict[str, float] = {}
|
||||
for rank, point in enumerate(reranked):
|
||||
score = 1.0 / (rank + 1)
|
||||
for mid in extract_message_ids(point):
|
||||
msg_score[mid] = msg_score.get(mid, 0.0) + score
|
||||
for rank, point in enumerate(rerank_tail):
|
||||
score = 1.0 / (60 + rank + 1)
|
||||
for mid in extract_message_ids(point):
|
||||
msg_score[mid] = msg_score.get(mid, 0.0) + score
|
||||
message_ids = sorted(msg_score, key=lambda m: msg_score[m], reverse=True)[:50]
|
||||
|
||||
return SearchAPIResponse(results=[SearchAPIItem(message_ids=message_ids)])
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import asyncio
|
||||
import os
|
||||
import re
|
||||
from functools import lru_cache
|
||||
|
|
@ -6,14 +5,14 @@ from functools import lru_cache
|
|||
import httpx
|
||||
from fastembed import SparseTextEmbedding
|
||||
|
||||
from .config import (
|
||||
from config import (
|
||||
EMBEDDINGS_DENSE_MODEL,
|
||||
EMBEDDINGS_DENSE_URL,
|
||||
SPARSE_MODEL_NAME,
|
||||
get_upstream_kwargs,
|
||||
logger,
|
||||
)
|
||||
from .schemas import DenseEmbeddingResponse, Question, SparseVector
|
||||
from schemas import DenseEmbeddingResponse, Question, SparseVector
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
|
|
@ -38,9 +37,20 @@ async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
|
|||
return payload.data[0].embedding
|
||||
|
||||
|
||||
async def embed_dense_multi(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
|
||||
tasks = [embed_dense(client, t) for t in texts]
|
||||
return list(await asyncio.gather(*tasks))
|
||||
async def embed_dense_batch(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
|
||||
"""Single request for multiple texts — avoids N parallel calls and rate limiting."""
|
||||
response = await client.post(
|
||||
str(EMBEDDINGS_DENSE_URL),
|
||||
**get_upstream_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(text: str) -> SparseVector:
|
||||
|
|
|
|||
|
|
@ -1,12 +1,13 @@
|
|||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from .config import RERANK_LIMIT, RERANKER_MODEL, RERANKER_URL, get_upstream_kwargs, logger
|
||||
from .retrieval import extract_page_content
|
||||
from config import RERANK_LIMIT, RERANKER_MODEL, RERANKER_URL, get_upstream_kwargs, logger
|
||||
from retrieval import extract_page_content
|
||||
|
||||
|
||||
async def get_rerank_scores(
|
||||
async def _get_rerank_scores(
|
||||
client: httpx.AsyncClient,
|
||||
query: str,
|
||||
targets: list[str],
|
||||
|
|
@ -14,6 +15,8 @@ async def get_rerank_scores(
|
|||
if not targets:
|
||||
return []
|
||||
|
||||
for attempt in range(5):
|
||||
try:
|
||||
response = await client.post(
|
||||
str(RERANKER_URL),
|
||||
**get_upstream_kwargs(),
|
||||
|
|
@ -24,35 +27,47 @@ async def get_rerank_scores(
|
|||
"text_2": targets,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
except Exception as exc:
|
||||
if attempt < 4:
|
||||
await asyncio.sleep(2 ** attempt)
|
||||
continue
|
||||
raise exc
|
||||
|
||||
if response.status_code == 429:
|
||||
wait = 2 ** attempt
|
||||
logger.warning("Rerank 429, retry %d/5 in %ds", attempt + 1, wait)
|
||||
await asyncio.sleep(wait)
|
||||
continue
|
||||
|
||||
response.raise_for_status()
|
||||
data = response.json().get("data") or []
|
||||
return [float(sample["score"]) for sample in data]
|
||||
|
||||
logger.error("Rerank 429 after all retries, falling back")
|
||||
return []
|
||||
|
||||
|
||||
async def rerank_points(
|
||||
client: httpx.AsyncClient,
|
||||
query: str,
|
||||
points: list[Any],
|
||||
) -> tuple[list[Any], list[Any]]:
|
||||
"""Return (reranked_head, retrieval_tail) so we don't lose candidates."""
|
||||
if not points:
|
||||
return [], []
|
||||
|
||||
rerank_candidates = points[:RERANK_LIMIT]
|
||||
head = points[:RERANK_LIMIT]
|
||||
tail = points[RERANK_LIMIT:]
|
||||
targets = [extract_page_content(p) for p in head]
|
||||
|
||||
targets = [extract_page_content(p) for p in rerank_candidates]
|
||||
try:
|
||||
scores = await get_rerank_scores(client, query, targets)
|
||||
scores = await _get_rerank_scores(client, query, targets)
|
||||
except Exception as exc:
|
||||
logger.warning("Rerank failed, using retrieval order: %s", exc)
|
||||
return rerank_candidates, tail
|
||||
return head, tail
|
||||
|
||||
if len(scores) != len(rerank_candidates):
|
||||
if len(scores) != len(head):
|
||||
logger.warning("Rerank score count mismatch, using retrieval order")
|
||||
return rerank_candidates, tail
|
||||
return head, tail
|
||||
|
||||
paired = sorted(zip(scores, rerank_candidates), key=lambda x: x[0], reverse=True)
|
||||
reranked = [p for _, p in paired]
|
||||
reranked = [p for _, p in sorted(zip(scores, head), key=lambda x: x[0], reverse=True)]
|
||||
return reranked, tail
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from typing import Any
|
|||
|
||||
from qdrant_client import AsyncQdrantClient, models
|
||||
|
||||
from .config import (
|
||||
from config import (
|
||||
DENSE_PREFETCH_K,
|
||||
QDRANT_COLLECTION_NAME,
|
||||
QDRANT_DENSE_VECTOR_NAME,
|
||||
|
|
@ -11,22 +11,28 @@ from .config import (
|
|||
SPARSE_PREFETCH_K,
|
||||
logger,
|
||||
)
|
||||
from .schemas import Question, SparseVector
|
||||
from schemas import Question, SparseVector
|
||||
|
||||
|
||||
def _build_filter(question: Question) -> models.Filter | None:
|
||||
must_conditions: list[models.Condition] = []
|
||||
|
||||
if question.date_range:
|
||||
try:
|
||||
must_conditions.append(
|
||||
models.FieldCondition(
|
||||
key="metadata.end",
|
||||
range=models.Range(gte=question.date_range.from_),
|
||||
)
|
||||
)
|
||||
must_conditions.append(
|
||||
models.FieldCondition(
|
||||
key="metadata.start",
|
||||
range=models.Range(
|
||||
gte=question.date_range.from_,
|
||||
lte=question.date_range.to,
|
||||
),
|
||||
range=models.Range(lte=question.date_range.to),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Date filter failed: %s", e)
|
||||
|
||||
if question.asker:
|
||||
must_conditions.append(
|
||||
|
|
|
|||
|
|
@ -1,15 +1,17 @@
|
|||
"""Unit tests for search/aggregation.py"""
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
_SEARCH_DIR = os.path.join(os.path.dirname(__file__), "..", "search")
|
||||
sys.path.insert(0, _SEARCH_DIR)
|
||||
|
||||
os.environ.setdefault("EMBEDDINGS_DENSE_URL", "http://localhost/embed")
|
||||
os.environ.setdefault("RERANKER_URL", "http://localhost/rerank")
|
||||
os.environ.setdefault("QDRANT_URL", "http://localhost:6333")
|
||||
os.environ.setdefault("API_KEY", "test-key")
|
||||
|
||||
from search.aggregation import aggregate_message_ids
|
||||
from search.config import TOP_K
|
||||
from aggregation import aggregate_message_ids
|
||||
from config import TOP_K
|
||||
|
||||
|
||||
def _point(message_ids: list[str]):
|
||||
|
|
|
|||
|
|
@ -1,11 +1,13 @@
|
|||
"""Unit tests for index/chunking.py"""
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
from index.chunking import build_chunks, _split_windows, WINDOW_MAX_MESSAGES, TIME_GAP_SECONDS
|
||||
from index.cleaning import CleanedMessage
|
||||
from index.schemas import Message
|
||||
_INDEX_DIR = os.path.join(os.path.dirname(__file__), "..", "index")
|
||||
sys.path.insert(0, _INDEX_DIR)
|
||||
|
||||
from chunking import build_chunks, _split_windows, WINDOW_MAX_MESSAGES, TIME_GAP_SECONDS
|
||||
from cleaning import CleanedMessage
|
||||
from index_schemas import Message
|
||||
|
||||
|
||||
def _make_message(id: str, time: int, text: str = "hello", **kwargs) -> Message:
|
||||
|
|
|
|||
|
|
@ -1,17 +1,19 @@
|
|||
"""Unit tests for index/cleaning.py"""
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
_INDEX_DIR = os.path.join(os.path.dirname(__file__), "..", "index")
|
||||
sys.path.insert(0, _INDEX_DIR)
|
||||
|
||||
import pytest
|
||||
from index.cleaning import (
|
||||
from cleaning import (
|
||||
normalize_unicode,
|
||||
parse_file_snippets,
|
||||
normalize_member_event,
|
||||
normalize_part,
|
||||
clean_message,
|
||||
)
|
||||
from index.schemas import Message
|
||||
from index_schemas import Message
|
||||
|
||||
|
||||
def _make_message(**kwargs) -> Message:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
"""Unit tests for search/query_builder.py (pure logic only, no HTTP)"""
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
_SEARCH_DIR = os.path.join(os.path.dirname(__file__), "..", "search")
|
||||
sys.path.insert(0, _SEARCH_DIR)
|
||||
|
||||
# Stub env vars before importing search modules
|
||||
os.environ.setdefault("EMBEDDINGS_DENSE_URL", "http://localhost/embed")
|
||||
|
|
@ -9,8 +11,8 @@ os.environ.setdefault("RERANKER_URL", "http://localhost/rerank")
|
|||
os.environ.setdefault("QDRANT_URL", "http://localhost:6333")
|
||||
os.environ.setdefault("API_KEY", "test-key")
|
||||
|
||||
from search.schemas import Entities, Question
|
||||
from search.query_builder import (
|
||||
from schemas import Entities, Question
|
||||
from query_builder import (
|
||||
build_primary_query,
|
||||
build_extra_dense_queries,
|
||||
build_sparse_query,
|
||||
|
|
|
|||
|
|
@ -1,10 +1,12 @@
|
|||
"""Unit tests for index/rendering.py"""
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
from index.cleaning import CleanedMessage
|
||||
from index.rendering import render_page_content, render_dense_content, render_sparse_content
|
||||
_INDEX_DIR = os.path.join(os.path.dirname(__file__), "..", "index")
|
||||
sys.path.insert(0, _INDEX_DIR)
|
||||
|
||||
from cleaning import CleanedMessage
|
||||
from rendering import render_page_content, render_dense_content, render_sparse_content
|
||||
|
||||
|
||||
def _make_cleaned(**kwargs) -> CleanedMessage:
|
||||
|
|
|
|||
Loading…
Reference in a new issue