Compare commits
5 commits
8d9b0d5f63
...
4800e25dd6
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4800e25dd6 | ||
|
|
c0f2d52f70 | ||
|
|
92cde65e42 | ||
|
|
19fec2361e | ||
|
|
4ecba7d35a |
6 changed files with 125 additions and 93 deletions
|
|
@ -1,19 +1,17 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
|
||||||
from contextlib import asynccontextmanager
|
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from fastapi import FastAPI, Request
|
from fastapi import FastAPI, Request
|
||||||
from fastapi.exceptions import RequestValidationError
|
from fastapi.exceptions import RequestValidationError
|
||||||
from fastapi.responses import JSONResponse, ORJSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
HOST = os.getenv("HOST", "0.0.0.0")
|
HOST = os.getenv("HOST", "0.0.0.0")
|
||||||
PORT = int(os.getenv("PORT", "8000"))
|
PORT = int(os.getenv("PORT", "8000"))
|
||||||
UVICORN_WORKERS = 4
|
UVICORN_WORKERS = 8
|
||||||
|
|
||||||
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
|
||||||
logger = logging.getLogger("index-service")
|
logger = logging.getLogger("index-service")
|
||||||
|
|
@ -80,24 +78,6 @@ OVERLAP_SIZE = 128
|
||||||
SPARSE_MODEL_NAME = "Qdrant/bm25"
|
SPARSE_MODEL_NAME = "Qdrant/bm25"
|
||||||
FASTEMBED_CACHE_PATH = "/models/fastembed"
|
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:
|
def render_message(message: Message) -> str:
|
||||||
parts_list: list[str] = []
|
parts_list: list[str] = []
|
||||||
|
|
@ -206,12 +186,7 @@ def build_chunks(
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
app = FastAPI(
|
app = FastAPI(title="Index Service", version="0.1.0")
|
||||||
title="Index Service",
|
|
||||||
version="0.1.0",
|
|
||||||
lifespan=lifespan,
|
|
||||||
default_response_class=ORJSONResponse,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@app.get("/health")
|
@app.get("/health")
|
||||||
|
|
@ -230,6 +205,14 @@ async def index(payload: IndexAPIRequest) -> IndexAPIResponse:
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@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]:
|
def embed_sparse_texts(texts: list[str]) -> list[dict]:
|
||||||
model = get_sparse_model()
|
model = get_sparse_model()
|
||||||
vectors = []
|
vectors = []
|
||||||
|
|
@ -245,9 +228,7 @@ def embed_sparse_texts(texts: list[str]) -> list[dict]:
|
||||||
|
|
||||||
@app.post("/sparse_embedding")
|
@app.post("/sparse_embedding")
|
||||||
async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]:
|
async def sparse_embedding(payload: SparseEmbeddingRequest) -> dict[str, Any]:
|
||||||
vectors = await asyncio.get_event_loop().run_in_executor(
|
vectors = await asyncio.to_thread(embed_sparse_texts, payload.texts)
|
||||||
_thread_pool, embed_sparse_texts, payload.texts
|
|
||||||
)
|
|
||||||
return {"vectors": vectors}
|
return {"vectors": vectors}
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -2,4 +2,3 @@ fastapi==0.135.1
|
||||||
uvicorn[standard]==0.42.0
|
uvicorn[standard]==0.42.0
|
||||||
pydantic==2.12.5
|
pydantic==2.12.5
|
||||||
fastembed==0.7.4
|
fastembed==0.7.4
|
||||||
orjson==3.10.18
|
|
||||||
|
|
|
||||||
|
|
@ -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_LOGIN = os.getenv("OPEN_API_LOGIN")
|
||||||
OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD")
|
OPEN_API_PASSWORD = os.getenv("OPEN_API_PASSWORD")
|
||||||
|
|
||||||
DENSE_PREFETCH_K = 50
|
DENSE_PREFETCH_K = 80
|
||||||
SPARSE_PREFETCH_K = 100
|
SPARSE_PREFETCH_K = 200
|
||||||
RETRIEVE_K = 80
|
RETRIEVE_K = 150
|
||||||
RERANK_LIMIT = 60
|
RERANK_LIMIT = 15
|
||||||
TOP_K = 50
|
TOP_K = 50
|
||||||
|
|
||||||
HTTP_TIMEOUT = 30.0
|
HTTP_TIMEOUT = 30.0
|
||||||
|
|
|
||||||
143
search/main.py
143
search/main.py
|
|
@ -9,7 +9,7 @@ import httpx
|
||||||
from fastembed import SparseTextEmbedding
|
from fastembed import SparseTextEmbedding
|
||||||
from fastapi import FastAPI, HTTPException, Request
|
from fastapi import FastAPI, HTTPException, Request
|
||||||
from fastapi.exceptions import RequestValidationError
|
from fastapi.exceptions import RequestValidationError
|
||||||
from fastapi.responses import JSONResponse, ORJSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
from qdrant_client import AsyncQdrantClient, models
|
from qdrant_client import AsyncQdrantClient, models
|
||||||
|
|
||||||
|
|
@ -148,13 +148,7 @@ def get_sparse_model() -> SparseTextEmbedding:
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(app: FastAPI):
|
async def lifespan(app: FastAPI):
|
||||||
# Preload BM25 and set up HTTP client with connection pooling
|
app.state.http = httpx.AsyncClient()
|
||||||
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(
|
app.state.qdrant = AsyncQdrantClient(
|
||||||
url=QDRANT_URL,
|
url=QDRANT_URL,
|
||||||
api_key=API_KEY,
|
api_key=API_KEY,
|
||||||
|
|
@ -166,17 +160,13 @@ async def lifespan(app: FastAPI):
|
||||||
await app.state.qdrant.close()
|
await app.state.qdrant.close()
|
||||||
|
|
||||||
|
|
||||||
app = FastAPI(
|
app = FastAPI(title="Search Service", version="0.1.0", lifespan=lifespan)
|
||||||
title="Search Service",
|
|
||||||
version="0.1.0",
|
|
||||||
lifespan=lifespan,
|
|
||||||
default_response_class=ORJSONResponse,
|
|
||||||
)
|
|
||||||
|
|
||||||
DENSE_PREFETCH_K = 80
|
DENSE_PREFETCH_K = 120
|
||||||
SPARSE_PREFETCH_K = 200
|
SPARSE_PREFETCH_K = 200
|
||||||
RETRIEVE_K = 150
|
RETRIEVE_K = 150
|
||||||
RERANK_LIMIT = 15
|
RERANK_LIMIT = 35
|
||||||
|
KEYWORD_BOOST_EXTRA = 10
|
||||||
|
|
||||||
|
|
||||||
async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
|
async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
|
||||||
|
|
@ -222,18 +212,64 @@ def embed_sparse_sync(text: str) -> SparseVector:
|
||||||
|
|
||||||
|
|
||||||
def build_dense_query(question: Question) -> str:
|
def build_dense_query(question: Question) -> str:
|
||||||
return question.text.strip()
|
q = question.search_text.strip() if question.search_text else question.text.strip()
|
||||||
|
return q
|
||||||
|
|
||||||
|
|
||||||
def build_sparse_query(question: Question) -> str:
|
def build_sparse_query(question: Question) -> str:
|
||||||
parts = [question.text.strip()]
|
base = question.search_text.strip() if question.search_text else question.text.strip()
|
||||||
|
parts = [base]
|
||||||
if question.keywords:
|
if question.keywords:
|
||||||
parts.extend(question.keywords)
|
parts.extend(question.keywords)
|
||||||
if question.search_text:
|
|
||||||
parts = [question.search_text]
|
|
||||||
return " ".join(parts)
|
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(
|
async def qdrant_search(
|
||||||
client: AsyncQdrantClient,
|
client: AsyncQdrantClient,
|
||||||
dense_vectors: list[list[float]],
|
dense_vectors: list[list[float]],
|
||||||
|
|
@ -326,25 +362,24 @@ async def rerank_points(
|
||||||
query: str,
|
query: str,
|
||||||
points: list[Any],
|
points: list[Any],
|
||||||
) -> list[Any]:
|
) -> list[Any]:
|
||||||
rerank_candidates = points[:RERANK_LIMIT]
|
if not points:
|
||||||
rerank_targets = [point.payload.get("page_content") for point in rerank_candidates]
|
return []
|
||||||
scores = await get_rerank_scores(client, query, rerank_targets)
|
targets = [point.payload.get("page_content") for point in points]
|
||||||
|
scores = await get_rerank_scores(client, query, targets)
|
||||||
|
|
||||||
if not scores:
|
if not scores or len(scores) != len(points):
|
||||||
logger.warning("Reranker unavailable, returning RRF order")
|
logger.warning("Reranker unavailable or score mismatch, returning RRF order")
|
||||||
return rerank_candidates
|
return points
|
||||||
|
|
||||||
reranked_candidates = [
|
return [
|
||||||
point
|
point
|
||||||
for _, point in sorted(
|
for _, point in sorted(
|
||||||
zip(scores, rerank_candidates, strict=True),
|
zip(scores, points),
|
||||||
key=lambda item: item[0],
|
key=lambda item: item[0],
|
||||||
reverse=True,
|
reverse=True,
|
||||||
)
|
)
|
||||||
]
|
]
|
||||||
|
|
||||||
return reranked_candidates
|
|
||||||
|
|
||||||
|
|
||||||
@app.get("/health")
|
@app.get("/health")
|
||||||
async def health() -> dict[str, str]:
|
async def health() -> dict[str, str]:
|
||||||
|
|
@ -365,17 +400,29 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
|
||||||
sparse_query = build_sparse_query(question)
|
sparse_query = build_sparse_query(question)
|
||||||
|
|
||||||
dense_task = embed_dense(client, dense_query)
|
dense_task = embed_dense(client, dense_query)
|
||||||
sparse_task = asyncio.to_thread(embed_sparse_sync, sparse_query)
|
sparse_task = asyncio.to_thread(lambda: embed_sparse_sync(sparse_query))
|
||||||
dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
|
dense_vector, sparse_vector = await asyncio.gather(dense_task, sparse_task)
|
||||||
|
|
||||||
dense_vectors = [dense_vector]
|
dense_vectors = [dense_vector]
|
||||||
if question.hyde and len(question.hyde) > 0:
|
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:
|
try:
|
||||||
hyde_texts = question.hyde[:2]
|
extra_vecs = await embed_dense_batch(client, extra_texts)
|
||||||
hyde_vectors = await embed_dense_batch(client, hyde_texts)
|
dense_vectors.extend(extra_vecs)
|
||||||
dense_vectors.extend(hyde_vectors)
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"HyDE embedding failed: {e}")
|
logger.warning(f"Extra dense embedding failed: {e}")
|
||||||
|
|
||||||
all_points = await qdrant_search(qdrant, dense_vectors, sparse_vector)
|
all_points = await qdrant_search(qdrant, dense_vectors, sparse_vector)
|
||||||
|
|
||||||
|
|
@ -384,20 +431,20 @@ async def search(payload: SearchAPIRequest) -> SearchAPIResponse:
|
||||||
|
|
||||||
all_points = list(all_points)
|
all_points = list(all_points)
|
||||||
|
|
||||||
reranked = await rerank_points(client, query, 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
|
||||||
|
|
||||||
reranked_ids = {id(p) for p in reranked}
|
msg_score: dict[str, float] = {}
|
||||||
remaining = [p for p in all_points if id(p) not in reranked_ids]
|
for rank, point in enumerate(reranked):
|
||||||
final_points = reranked + remaining
|
score = 1.0 / (rank + 1)
|
||||||
|
|
||||||
seen: set[str] = set()
|
|
||||||
message_ids: list[str] = []
|
|
||||||
for point in final_points:
|
|
||||||
for mid in extract_message_ids(point):
|
for mid in extract_message_ids(point):
|
||||||
if mid not in seen:
|
msg_score[mid] = msg_score.get(mid, 0.0) + score
|
||||||
seen.add(mid)
|
for rank, point in enumerate(rerank_tail):
|
||||||
message_ids.append(mid)
|
score = 1.0 / (60 + rank + 1)
|
||||||
message_ids = message_ids[:50]
|
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)])
|
return SearchAPIResponse(results=[SearchAPIItem(message_ids=message_ids)])
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -4,4 +4,3 @@ pydantic==2.12.5
|
||||||
httpx==0.28.1
|
httpx==0.28.1
|
||||||
qdrant-client==1.15.1
|
qdrant-client==1.15.1
|
||||||
fastembed==0.7.4
|
fastembed==0.7.4
|
||||||
orjson==3.10.18
|
|
||||||
|
|
|
||||||
|
|
@ -18,15 +18,21 @@ def _build_filter(question: Question) -> models.Filter | None:
|
||||||
must_conditions: list[models.Condition] = []
|
must_conditions: list[models.Condition] = []
|
||||||
|
|
||||||
if question.date_range:
|
if question.date_range:
|
||||||
must_conditions.append(
|
try:
|
||||||
models.FieldCondition(
|
must_conditions.append(
|
||||||
key="metadata.start",
|
models.FieldCondition(
|
||||||
datetime_range=models.DatetimeRange(
|
key="metadata.end",
|
||||||
gte=question.date_range.from_,
|
range=models.Range(gte=question.date_range.from_),
|
||||||
lte=question.date_range.to,
|
)
|
||||||
),
|
|
||||||
)
|
)
|
||||||
)
|
must_conditions.append(
|
||||||
|
models.FieldCondition(
|
||||||
|
key="metadata.start",
|
||||||
|
range=models.Range(lte=question.date_range.to),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning("Date filter failed: %s", e)
|
||||||
|
|
||||||
if question.asker:
|
if question.asker:
|
||||||
must_conditions.append(
|
must_conditions.append(
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue