From 418fab29c32fc41e465b39f79f3c1ad26a2f865b Mon Sep 17 00:00:00 2001 From: kert Date: Tue, 8 Sep 2026 23:36:01 -0400 Subject: [PATCH] fix(llm): per-call DuckDB cursor for the cached replica handle; lock the cache; concurrency test (refs P48) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A /chat stream is a sync generator, so Starlette runs it in a threadpool worker: concurrent turns shared the one cached DuckDBPyConnection and read each other's result sets ("not enough values to unpack"), which the broad except swallowed as "valuation evidence skipped" — the table just vanished. Each call now runs on its own cursor off the shared handle and closes it in a finally, and the cache check-and-open is under a lock so a race on the first open cannot leak a connection. --- src/llm/evidence.py | 36 +++++++----- tests/llm/test_evidence.py | 115 +++++++++++++++++++++++++++++++++++-- 2 files changed, 131 insertions(+), 20 deletions(-) diff --git a/src/llm/evidence.py b/src/llm/evidence.py index 125f21f..80e1cc3 100644 --- a/src/llm/evidence.py +++ b/src/llm/evidence.py @@ -12,6 +12,7 @@ from __future__ import annotations import logging import os +import threading from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Sequence @@ -98,6 +99,7 @@ class ValuationEvidence: #: so the handle stays valid on the old inode until we notice the new #: mtime and reopen. _REPLICA: tuple[str, int, Any] | None = None +_REPLICA_LOCK = threading.Lock() def _mtime(path: str) -> int: @@ -112,18 +114,19 @@ def _connect(path: str) -> Any: replica file has been republished since we opened it.""" global _REPLICA mtime = _mtime(path) - if _REPLICA is not None: - cached_path, cached_mtime, con = _REPLICA - if (cached_path, cached_mtime) == (path, mtime): - return con - _REPLICA = None - try: - con.close() - except Exception as e: # noqa: BLE001 — a stale handle must not break the chat - log.warning("stale replica handle not closed: %s", e) - con = duckdb.connect(path, read_only=True) - _REPLICA = (path, mtime, con) - return con + with _REPLICA_LOCK: # two chat threads racing the first open would leak one + if _REPLICA is not None: + cached_path, cached_mtime, con = _REPLICA + if (cached_path, cached_mtime) == (path, mtime): + return con + _REPLICA = None + try: + con.close() + except Exception as e: # noqa: BLE001 — a stale handle must not break the chat + log.warning("stale replica handle not closed: %s", e) + con = duckdb.connect(path, read_only=True) + _REPLICA = (path, mtime, con) + return con def _replica_path(cfg: LlmConfig) -> str: @@ -141,7 +144,10 @@ def valuation_evidence(question: str, cfg: LlmConfig) -> ValuationEvidence | Non return None path = _replica_path(cfg) try: - con = _connect(path) + # A chat turn runs in a threadpool worker, and DuckDB gives each + # thread its own cursor off the shared connection — without one, + # concurrent turns read each other's result sets. + cur = _connect(path).cursor() except Exception as e: # noqa: BLE001 — duckdb raises several types log.warning("valuation evidence skipped (%s): %s", path, e) return None @@ -149,10 +155,12 @@ def valuation_evidence(question: str, cfg: LlmConfig) -> ValuationEvidence | Non # crowd the excerpts out of the prompt. years = cfg.valuation_years if len(det.codes) <= 3 else 2 try: - rows, unpriced = valuation(con, list(det.codes), years=years) + rows, unpriced = valuation(cur, list(det.codes), years=years) except Exception as e: # noqa: BLE001 log.warning("valuation evidence skipped (query): %s", e) return None + finally: + cur.close() # the cached parent connection stays open return ValuationEvidence(det.codes, det.families, tuple(rows), tuple(unpriced)) diff --git a/tests/llm/test_evidence.py b/tests/llm/test_evidence.py index 1b6ffee..882dd90 100644 --- a/tests/llm/test_evidence.py +++ b/tests/llm/test_evidence.py @@ -2,8 +2,11 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace from unittest.mock import MagicMock, patch +import duckdb import pytest from llm import evidence @@ -32,6 +35,14 @@ CFG = LlmConfig( ) +@pytest.fixture(autouse=True) +def _clear_replica_cache(): + """The replica handle is process-global; no test may inherit another's.""" + evidence._REPLICA = None + yield + evidence._REPLICA = None + + def _row(**kw) -> ValuationRow: base = dict( code="G0556", @@ -161,12 +172,6 @@ class TestEvidenceText: class TestValuationEvidence: - @pytest.fixture(autouse=True) - def _clear_replica_cache(self): - evidence._REPLICA = None - yield - evidence._REPLICA = None - def test_none_when_no_codes(self): assert ( valuation_evidence("what did commenters say about telehealth?", CFG) is None @@ -223,6 +228,104 @@ class TestValuationEvidence: assert mock_val.call_args.kwargs == {"years": 2} +RVU_COLS = ( + "hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, " + "work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, " + "non_fac_total DOUBLE, fac_total DOUBLE, conv_factor DOUBLE, year INTEGER" +) +PROPOSED_COLS = ( + "hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, " + "work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, " + "cms_rule_id VARCHAR" +) + + +class TestConcurrentChats: + """/chat streams run in Starlette's threadpool, so several turns share + the cached replica handle. Without a per-call cursor their result sets + cross-talk and the broad except swallows it as "no valuation".""" + + @pytest.fixture + def replica(self, tmp_path): + db = tmp_path / "aco.ro.duckdb" + con = duckdb.connect(str(db)) + con.execute("CREATE SCHEMA pfs") + con.execute(f"CREATE TABLE pfs.rvu ({RVU_COLS})") + con.execute(f"CREATE TABLE pfs.rvu_proposed ({PROPOSED_COLS})") + # two unfamilied codes with different row counts, so a swapped + # result set is visible as a wrong count + con.executemany( + "INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?,?,?,?,?,?,?)", + [ + ( + "99213", + None, + "Office visit", + "A", + 1.0, + 1.0, + 0.5, + 0.05, + 2.05, + 1.55, + 30.0, + 2026, + ), + *[ + ( + "99214", + None, + "Office visit", + "A", + 1.5, + 1.0, + 0.5, + 0.05, + 2.55, + 2.05, + 30.0, + y, + ) + for y in (2024, 2025, 2026) + ], + ], + ) + con.close() + return db + + def test_parallel_turns_share_one_handle_without_cross_talk(self, replica): + cfg = replace(CFG, duckdb_replica=str(replica)) + real_connect, opens = duckdb.connect, [] + + def counting_connect(*a, **kw): + opens.append(a) + return real_connect(*a, **kw) + + expected = {"99213": 1, "99214": 3} + questions = ["What does 99213 pay?", "What does 99214 pay?"] + + def one(i): + q = questions[i % 2] + ev = valuation_evidence(q, cfg) + assert ev is not None, "valuation dropped under concurrency" + (code,) = ev.codes + return code, len(ev.rows) + + with patch("llm.evidence.duckdb.connect", counting_connect): + with ThreadPoolExecutor(max_workers=8) as pool: + results = list(pool.map(one, range(160))) + + assert all(n == expected[code] for code, n in results), sorted(set(results)) + assert len(opens) == 1 # one handle, shared + + def test_cursor_is_closed_but_the_handle_is_kept(self, replica): + cfg = replace(CFG, duckdb_replica=str(replica)) + assert valuation_evidence("99213?", cfg) is not None + assert evidence._REPLICA is not None + # the cached parent still answers — it was never closed + assert valuation_evidence("99214?", cfg) is not None + + def _engine(rows): engine = MagicMock() conn = engine.begin.return_value.__enter__.return_value