Speed up knowledge graph builds by extracting pages concurrently

pipeline/kg.py

@@ -44,6 +44,7 @@ import time
 import urllib.error
 import urllib.request
 from collections import deque
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from typing import Optional
 
@@ -98,6 +99,19 @@ _RETRY_STATUSES = frozenset({429, 503})
 _RETRY_ATTEMPTS = 4
 _RETRY_BASE_DELAY = 1.5  # seconds; doubled each attempt
 
+# Build scalability: per-page triple extraction is one (I/O-bound) Gemini call
+# each, so they're fanned across this SHARED, bounded pool instead of running
+# strictly sequentially — a multi-page build is ~Nx faster on a key with real
+# throughput. The pool is module-level on purpose: being shared, it also caps the
+# TOTAL number of concurrent Gemini calls across simultaneous builds/users, so one
+# big document can't blow the rate limit for everyone. Overshoot of the provider's
+# rate limit is absorbed by the 429/503 backoff in _post_json. Graph mutation stays
+# single-threaded (results are collected, then applied in deterministic page order).
+_BUILD_CONCURRENCY = 5
+_EXTRACT_POOL = ThreadPoolExecutor(
+    max_workers=_BUILD_CONCURRENCY, thread_name_prefix="kg-extract"
+)
+
 # Instruction for triple extraction. We ask for typed entities AND triples and
 # force a JSON response via responseSchema for deterministic parsing.
 _TRIPLE_PROMPT = (
@@ -906,13 +920,21 @@ def build_graph(
     api_key : str
         A Gemini API key (required; this is the opt-in online feature).
     triple_model, embed_model : str, optional
-        Model overrides; default to Flash + text-embedding-004.
+        Model overrides; default to Flash + a discovered embedding model
+        (``_pick_embedding_model``, prefers gemini-embedding-001).
     max_pages : int, optional
         Cap the number of non-empty pages processed (cost control). None = all.
     embed : bool
         If True (default) also embed the nodes so semantic search works; if
         False the graph supports lexical search only (cheaper).
 
+    Per-page triple extraction is fanned across a shared bounded pool
+    (``_EXTRACT_POOL``), so a multi-page build is much faster on a key with real
+    throughput; the shared pool also caps total concurrent Gemini calls. Each
+    page is isolated — a single failed page is skipped and counted
+    (``pages_failed``); only an all-pages failure raises. Embedding failure is
+    non-fatal (``embed_error`` set, lexical search retained).
+
     Returns a graph (possibly empty if the document yields no facts). Raises a
     single-line RuntimeError only for a missing key or a hard Gemini failure.
     """
@@ -928,15 +950,28 @@ def build_graph(
 
     graph = KnowledgeGraph()
     last_error: Optional[Exception] = None
-    for page_no, text in usable:
-        # Isolate each page: a single rate-limited / safety-blocked page (common
-        # on the free tier) must NOT discard every page already processed (and
-        # its API spend). Skip it and keep going.
+
+    # Fan the per-page Gemini calls across the shared pool (concurrency), then
+    # build the graph single-threaded below in PAGE ORDER so node ids stay
+    # deterministic and the adjacency dicts are never mutated from two threads.
+    def _extract_one(text):
+        return extract_triples(text, api_key=api_key, model=triple_model, timeout=timeout)
+
+    futures = [(page_no, _EXTRACT_POOL.submit(_extract_one, text)) for page_no, text in usable]
+    results: dict = {}
+    for page_no, fut in futures:
+        # Isolate each page: one rate-limited / safety-blocked page (common on the
+        # free tier) must NOT discard every other page (and its API spend).
         try:
-            data = extract_triples(text, api_key=api_key, model=triple_model, timeout=timeout)
+            results[page_no] = fut.result()
         except RuntimeError as exc:
-            graph.pages_failed += 1
+            results[page_no] = exc
             last_error = exc
+
+    for page_no, _text in usable:
+        data = results.get(page_no)
+        if not isinstance(data, dict):  # a skipped page (exception) or missing
+            graph.pages_failed += 1
             continue
         graph.pages_built += 1
         # Register typed entities first so types are known, then triples.