Richie
2706c4417d
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
125 lines
4.4 KiB
Python
125 lines
4.4 KiB
Python
"""Shared ORM definitions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from os import getenv
|
|
from typing import cast
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.engine import URL, Engine
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
|
|
|
NAMING_CONVENTION = {
|
|
"ix": "ix_%(table_name)s_%(column_0_name)s",
|
|
"uq": "uq_%(table_name)s_%(column_0_name)s",
|
|
"ck": "ck_%(table_name)s_%(constraint_name)s",
|
|
"fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
|
|
"pk": "pk_%(table_name)s",
|
|
}
|
|
|
|
|
|
def get_connection_info(name: str) -> tuple[str, str, str, str, str | None]:
|
|
"""Get connection info from environment variables."""
|
|
database = getenv(f"{name}_DB")
|
|
host = getenv(f"{name}_HOST")
|
|
port = getenv(f"{name}_PORT")
|
|
username = getenv(f"{name}_USER")
|
|
password = getenv(f"{name}_PASSWORD")
|
|
|
|
if None in (database, host, port, username):
|
|
error = f"Missing environment variables for Postgres connection.\n{database=}\n{host=}\n{port=}\n{username=}\n"
|
|
raise ValueError(error)
|
|
return cast("tuple[str, str, str, str, str | None]", (database, host, port, username, password))
|
|
|
|
|
|
def build_postgres_url(name: str, *, vector_engine: bool = False) -> tuple[URL, dict[str, str]]:
|
|
"""Build the Postgres connection URL and connect_args from environment variables.
|
|
|
|
Args:
|
|
name (str): The name of the environment variable prefix.
|
|
vector_engine (bool, optional): Whether to use the vector search schema. Defaults to False.
|
|
This updates the search path to include the vector types and operators.
|
|
|
|
Returns:
|
|
tuple[URL, dict[str, str]]: The SQLAlchemy URL and connect_args for create_engine.
|
|
"""
|
|
database, host, port, username, password = get_connection_info(name)
|
|
|
|
url = URL.create(
|
|
drivername="postgresql+psycopg",
|
|
username=username,
|
|
password=password,
|
|
host=host,
|
|
port=int(port),
|
|
database=database,
|
|
)
|
|
|
|
connect_args = {}
|
|
# There more better way to do this is with separate PG account and a dedicated vector schema for the vector types
|
|
if vector_engine:
|
|
connect_args["options"] = "-csearch_path=main,public"
|
|
|
|
return url, connect_args
|
|
|
|
|
|
def get_postgres_engine(
|
|
*,
|
|
name: str = "POSTGRES",
|
|
pool_pre_ping: bool = True,
|
|
vector_engine: bool = False,
|
|
pool_size: int = 8,
|
|
) -> Engine:
|
|
"""Create a SQLAlchemy engine from environment variables.
|
|
|
|
Args:
|
|
name (str, optional): The name of the environment variable prefix. Defaults to "POSTGRES".
|
|
pool_pre_ping (bool, optional): Whether to ping the database before each connection. Defaults to True.
|
|
This fixes the issue of trying to use a conection that has timed out on the database side.
|
|
vector_engine (bool, optional): Whether to use the vector search schema. Defaults to False.
|
|
This updates the search path the incldued the vecore types and operators.
|
|
pool_size (int, optional): Number of connections to keep in the pool. Defaults to 8.
|
|
|
|
Returns:
|
|
Engine: The SQLAlchemy engine.
|
|
"""
|
|
url, connect_args = build_postgres_url(name, vector_engine=vector_engine)
|
|
|
|
return create_engine(
|
|
url=url,
|
|
pool_pre_ping=pool_pre_ping,
|
|
pool_recycle=1800,
|
|
connect_args=connect_args,
|
|
pool_size=pool_size,
|
|
)
|
|
|
|
|
|
def get_async_postgres_engine(
|
|
*,
|
|
name: str = "POSTGRES",
|
|
pool_pre_ping: bool = True,
|
|
vector_engine: bool = False,
|
|
pool_size: int = 8,
|
|
) -> AsyncEngine:
|
|
"""Create an async SQLAlchemy engine from environment variables.
|
|
|
|
Args:
|
|
name (str, optional): The name of the environment variable prefix. Defaults to "POSTGRES".
|
|
pool_pre_ping (bool, optional): Whether to ping the database before each connection. Defaults to True.
|
|
This fixes the issue of trying to use a conection that has timed out on the database side.
|
|
vector_engine (bool, optional): Whether to use the vector search schema. Defaults to False.
|
|
This updates the search path the incldued the vecore types and operators.
|
|
pool_size (int, optional): Number of connections to keep in the pool. Defaults to 8.
|
|
|
|
Returns:
|
|
AsyncEngine: The async SQLAlchemy engine.
|
|
"""
|
|
url, connect_args = build_postgres_url(name, vector_engine=vector_engine)
|
|
|
|
return create_async_engine(
|
|
url=url,
|
|
pool_pre_ping=pool_pre_ping,
|
|
pool_recycle=1800,
|
|
connect_args=connect_args,
|
|
pool_size=pool_size,
|
|
)
|