mempalace.kg_triple_worker

Source: mempalace/kg_triple_worker.py

Async worker that drains mempalace_kg_extraction_queue.

The worker pulls drawer_ids from the postgres queue, calls the LLM extractor on each drawer's text, and writes the resulting triples into the AGE knowledge graph via a per-coroutine connection from an AsyncConnectionPool.

Concurrency model:

• One httpx.AsyncClient shared across the loop.
• asyncio.Semaphore(max_concurrency) caps in-flight LLM calls so we never overrun llama-server --parallel N.
• Postgres I/O uses psycopg3 AsyncConnectionPool. Each coroutine acquires its own connection for queue ops and AGE triple writes, so there is no global write lock. Pool size matches the LLM concurrency cap so the worker can overlap N drawer extractions with N writes.

Queue claim uses FOR UPDATE SKIP LOCKED so multiple worker processes (or threads) can drain the queue without colliding.

Migration note: this module previously serialized triple writes through a single psycopg2 connection wrapped in an asyncio.Lock. After llm_refine ReDoS fix #206 raised the LLM throughput ~10×, that lock became the binding constraint. The pool-per-coroutine design replaces both the lock and the per-call asyncio.to_thread hop.

The CLI entry point is exposed as mempalace-kg-extract via pyproject.toml; see cli_main at the bottom of this module.

Classes

class WorkerStats

Lightweight in-process counters that wrap the queue table snapshot.

snapshot

def snapshot(self) -> dict

Functions

run_worker

async def run_worker(dsn: str, llm_endpoint: str = DEFAULT_ENDPOINT, model: str = DEFAULT_MODEL, *, batch_size: int = DEFAULT_BATCH_SIZE, poll_interval: int = DEFAULT_POLL_INTERVAL, max_concurrency: int = DEFAULT_CONCURRENCY, db_pool_size: Optional[int] = None, worker_id: Optional[str] = None, backfill: bool = False, backfill_limit: Optional[int] = None, once: bool = False, pool_factory: Optional[Callable[[str, int, int], Any]] = None, kg_factory: Optional[Callable[[Any], _KGHandle]] = None, http_client_factory: Optional[Callable[[], Awaitable[Any]]] = None, stats: Optional[WorkerStats] = None, stop_event: Optional[asyncio.Event] = None) -> WorkerStats

Drain the extraction queue until cancelled (or once=True).

Concurrency model: a producer task tops up an asyncio.Queue of claimed drawers from the postgres queue whenever it dips below a low-water mark, while max_concurrency long-lived consumer tasks pull from that internal queue and call _process_one. Replaces the old claim-batch → gather → claim-batch barrier so the slowest drawer in a batch can no longer stall the next claim cycle.

Args: dsn: Postgres DSN for both the queue and mempalace_drawers. llm_endpoint: Base URL for the OpenAI-compatible inference server. model: Model alias. batch_size: How many rows to claim per refill cycle. Also drives the internal queue's high-water mark (= batch_size) and low-water refill threshold (= batch_size // 2). poll_interval: Seconds to sleep when the queue is empty. max_concurrency: Cap on in-flight LLM calls (matches llama-server --parallel N). Equals the number of persistent consumer tasks. db_pool_size: psycopg pool max_size. Defaults to max_concurrency + 8 so every in-flight LLM call has a write conn plus slack for the producer's claim/refill ops. Caller invariant: db_pool_size >= max_concurrency. worker_id: Identifier written to worker_id on each claim. Defaults to hostname:pid:short-uuid. backfill: If true, bulk-enqueue every uncompleted drawer before entering the normal claim loop. backfill_limit: Cap the number of rows seeded in backfill mode. once: If true, drain the current claimable set then exit. The producer claims once (no refill loop) and consumers exit after the internal queue is empty. pool_factory / kg_factory / http_client_factory: Test seams. stats: Pre-existing WorkerStats to mutate; one is created if None. stop_event: Async event that causes the loop to exit cleanly when set.

Returns the final WorkerStats.

get_status

def get_status(dsn: str) -> dict

One-shot status query used by the CLI's --status flag.

Kept synchronous: callers run it from non-async contexts (CLI flag, monitoring scripts) where spinning up an async pool would be overkill. Uses a single short-lived psycopg3 connection.

cli_main

def cli_main(argv: Optional[list[str]] = None) -> int