feat(ebook): migrate to async DB/HTTP and parallelize phrase pipeline
Convert the ebook-search web app to async end to end and add concurrency to the protected-phrase extraction and judging pipeline so large books no longer block the event loop or the UI. ORM / infra: - Add get_async_postgres_engine and factor shared URL/connect_args building into build_postgres_url (reused by the sync and async engine builders) - Add async FastAPI session helpers (get_async_db, AsyncDbSession) with expire_on_commit=False to avoid implicit IO under asyncio App: - Use AsyncEngine/AsyncSession throughout routes, search, ingest, embeddings, answer, rerank and LLM calls; convert handlers to async - Share a single httpx.AsyncClient in app state for LLM requests; size the connection pool for concurrent phrase-judging workers - Add judge_tasks: run per-book judging as tracked background tasks so a book already being judged isn't double-queued Protected phrases: - Add a process pool (pool.py) and worker-count config (extraction/judge book/phrase workers) to parallelize candidate generation and judging - Split admin actions into all/missing variants for generation and judging Config: - Add protected_phrase_extraction_workers, phrase_judge_book_workers, phrase_judge_phrase_workers
This commit is contained in:
@@ -22,10 +22,26 @@ def auth_headers(api_key: str) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
|
||||
def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
|
||||
"""Request embeddings from the configured OpenAI-compatible endpoint."""
|
||||
async def request_embeddings(
|
||||
client: httpx.AsyncClient,
|
||||
texts: Sequence[str],
|
||||
config: EbookSearchConfig,
|
||||
) -> list[list[float]]:
|
||||
"""Request embeddings from the configured OpenAI-compatible endpoint.
|
||||
|
||||
Args:
|
||||
client (httpx.AsyncClient): Shared async client for LLM calls.
|
||||
texts (Sequence[str]): Texts to embed.
|
||||
config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth.
|
||||
|
||||
Returns:
|
||||
list[list[float]]: One embedding vector per input text.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the request fails or the response cannot be parsed.
|
||||
"""
|
||||
try:
|
||||
response = httpx.post(
|
||||
response = await client.post(
|
||||
f"{config.embedding_base_url.rstrip('/')}/embeddings",
|
||||
headers=auth_headers(config.embedding_api_key),
|
||||
json={"model": config.embedding_model, "input": list(texts)},
|
||||
@@ -44,10 +60,15 @@ def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[
|
||||
raise RuntimeError(msg) from error
|
||||
|
||||
|
||||
def check_embedding_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool:
|
||||
async def check_embedding_endpoint(
|
||||
client: httpx.AsyncClient,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> bool:
|
||||
"""Return whether the configured embedding endpoint answers a model listing."""
|
||||
try:
|
||||
response = httpx.get(
|
||||
response = await client.get(
|
||||
f"{config.embedding_base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(config.embedding_api_key),
|
||||
timeout=timeout_seconds,
|
||||
@@ -59,10 +80,15 @@ def check_embedding_endpoint(config: EbookSearchConfig, *, timeout_seconds: floa
|
||||
return True
|
||||
|
||||
|
||||
def check_chat_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool:
|
||||
async def check_chat_endpoint(
|
||||
client: httpx.AsyncClient,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> bool:
|
||||
"""Return whether the configured chat (answering) endpoint answers a model listing."""
|
||||
try:
|
||||
response = httpx.get(
|
||||
response = await client.get(
|
||||
f"{config.vllm_base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(config.vllm_api_key),
|
||||
timeout=timeout_seconds,
|
||||
@@ -98,18 +124,29 @@ def embedding_vectors_from_response(body: object) -> list[list[float]]:
|
||||
return vectors
|
||||
|
||||
|
||||
def request_rerank(
|
||||
async def request_rerank(
|
||||
client: httpx.AsyncClient,
|
||||
query: str,
|
||||
documents: Sequence[str],
|
||||
config: RerankConfig,
|
||||
) -> object | None:
|
||||
"""Request rerank scores from the configured vLLM endpoint."""
|
||||
"""Request rerank scores from the configured vLLM endpoint.
|
||||
|
||||
Args:
|
||||
client (httpx.AsyncClient): Shared async client for LLM calls.
|
||||
query (str): Query the documents are scored against.
|
||||
documents (Sequence[str]): Candidate documents to score.
|
||||
config (RerankConfig): Rerank endpoint settings.
|
||||
|
||||
Returns:
|
||||
object | None: The decoded response body, or ``None`` when it is not valid JSON.
|
||||
"""
|
||||
payload = {
|
||||
"model": config.model,
|
||||
"query": query,
|
||||
"documents": list(documents),
|
||||
}
|
||||
response = httpx.post(
|
||||
response = await client.post(
|
||||
f"{config.base_url.rstrip('/')}/rerank",
|
||||
json=payload,
|
||||
timeout=config.timeout_seconds,
|
||||
@@ -122,13 +159,26 @@ def request_rerank(
|
||||
return None
|
||||
|
||||
|
||||
def request_chat_completion(
|
||||
async def request_chat_completion(
|
||||
client: httpx.AsyncClient,
|
||||
config: EbookSearchConfig,
|
||||
messages: Sequence[dict[str, str]],
|
||||
) -> str:
|
||||
"""Request a chat completion from the configured OpenAI-compatible endpoint."""
|
||||
"""Request a chat completion over a shared async client.
|
||||
|
||||
Args:
|
||||
client (httpx.AsyncClient): Shared async client whose connection pool bounds concurrency.
|
||||
config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth.
|
||||
messages (Sequence[dict[str, str]]): OpenAI-style chat messages.
|
||||
|
||||
Returns:
|
||||
str: The assistant message text.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the request fails or the response cannot be parsed.
|
||||
"""
|
||||
try:
|
||||
response = httpx.post(
|
||||
response = await client.post(
|
||||
f"{config.vllm_base_url.rstrip('/')}/chat/completions",
|
||||
headers=auth_headers(config.vllm_api_key),
|
||||
json={
|
||||
|
||||
Reference in New Issue
Block a user