diff --git a/src/agent/graph.py b/src/agent/graph.py index f51666b..7b50cf7 100644 --- a/src/agent/graph.py +++ b/src/agent/graph.py @@ -50,8 +50,38 @@ def __init__( self.pool: AsyncConnectionPool[AsyncConnection[dict[str, Any]]] | None = None def __del__(self) -> None: - if self.pool: - asyncio.run(self.close_pool()) + """Close the connection pool if nothing else did. + + This used to call asyncio.run() unconditionally, which raises + RuntimeError when a loop is already running -- and __del__ can fire at + any point, including inside the running server. Exceptions in __del__ are + swallowed and printed, so it surfaced as noise in production logs with the + pool still open. + + Scheduling the close with loop.create_task() is not a fix either: the + task is not guaranteed to run if the loop is shutting down, which is + exactly when a graph is usually collected. + + Nothing calls close_pool() explicitly today, so this is the only cleanup + there is. The real fix is an explicit lifecycle -- close the pool from the + application's shutdown hook -- which belongs with the agent-API work. + """ + if self.pool is None: + return + try: + asyncio.get_running_loop() + except RuntimeError: + # No loop running, so we can drive the close to completion. + try: + asyncio.run(self.close_pool()) + except Exception as e: + logging.warning(f"Could not close the connection pool: {e}") + return + logging.warning( + "AgentGraph was garbage-collected while an event loop is running; " + "its Postgres pool is still open. Close it explicitly from the " + "application shutdown hook." + ) async def initialize(self) -> dict[str, CompiledStateGraph]: checkpointer: BaseCheckpointSaver[str] = await self.create_checkpointer() diff --git a/src/agent/profiles/base.py b/src/agent/profiles/base.py index b11d653..3045c7e 100644 --- a/src/agent/profiles/base.py +++ b/src/agent/profiles/base.py @@ -1,3 +1,4 @@ +import asyncio from typing import Annotated, Literal, TypedDict from langchain_core.embeddings import Embeddings @@ -57,11 +58,15 @@ async def preprocess(self, state: BaseState, config: RunnableConfig) -> BaseStat }, config, ) - safety_check: SafetyCheck = await self.safety_checker.ainvoke( - {"rephrased_input": rephrased_input}, config - ) - detected_language: str = await self.language_detector.ainvoke( - {"user_input": state["user_input"]}, config + # The safety check needs the rephrased text, so it has to follow. Language + # detection reads the raw user input and does not, so the two overlap + # instead of running back to back -- one LLM round trip saved on every + # message, on the path every profile uses. + safety_check: SafetyCheck + detected_language: str + safety_check, detected_language = await asyncio.gather( + self.safety_checker.ainvoke({"rephrased_input": rephrased_input}, config), + self.language_detector.ainvoke({"user_input": state["user_input"]}, config), ) return BaseState( rephrased_input=rephrased_input, diff --git a/src/data_generation/alliance/__init__.py b/src/data_generation/alliance/__init__.py index 2c79c9b..badbc3e 100644 --- a/src/data_generation/alliance/__init__.py +++ b/src/data_generation/alliance/__init__.py @@ -113,7 +113,8 @@ def upload_to_chromadb( "Alt. ID(s) interactor B", "Alias(es) interactor A", "Alias(es) interactor B", - "Interaction detection method(s) Publication 1st author(s)", + "Interaction detection method(s)", + "Publication 1st author(s)", "Publication Identifier(s)", "Taxid interactor A", "Taxid interactor B", @@ -134,7 +135,8 @@ def upload_to_chromadb( "Annotation(s) interactor A", "Annotation(s) interactor B", "Interaction annotation(s)", - "Host organism(s)Interaction parameter(s)", + "Host organism(s)", + "Interaction parameter(s)", "Creation date", "Update date", "Checksum(s) interactor A", @@ -173,11 +175,11 @@ def upload_to_chromadb( "Xref(s) interactor A", "Xref(s) interactor B", "Interaction Xref(s)", - "Annotation(s) interactor A" - "Annotation(s) interactor B" - "Interaction annotation(s)" - "Host organism(s)" - "Interaction parameter(s)" + "Annotation(s) interactor A", + "Annotation(s) interactor B", + "Interaction annotation(s)", + "Host organism(s)", + "Interaction parameter(s)", "Creation date", "Update date", "Checksum(s) interactor A", diff --git a/tests/agent/test_graph_lifecycle.py b/tests/agent/test_graph_lifecycle.py new file mode 100644 index 0000000..97dfb8a --- /dev/null +++ b/tests/agent/test_graph_lifecycle.py @@ -0,0 +1,70 @@ +"""AgentGraph.__del__ must never raise, whatever the loop state. + +It used to call asyncio.run() unconditionally. That raises RuntimeError when a +loop is already running, and __del__ can fire at any moment -- including inside +the running server. Reported independently in PRs #147 and #156. +""" + +import asyncio + +from agent.graph import AgentGraph + + +class _FakePool: + """Stands in for AsyncConnectionPool; records whether it was closed.""" + + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + +def _graph_with_pool() -> tuple[AgentGraph, _FakePool]: + """Build an AgentGraph without running __init__, which would need an LLM.""" + graph = AgentGraph.__new__(AgentGraph) + pool = _FakePool() + graph.pool = pool # type: ignore[assignment] + graph.graph = None + return graph, pool + + +def test_del_closes_the_pool_when_no_loop_is_running() -> None: + graph, pool = _graph_with_pool() + graph.__del__() + assert pool.closed is True + + +def test_del_does_not_raise_inside_a_running_loop() -> None: + """The regression: asyncio.run() here raised RuntimeError.""" + + async def collect() -> None: + graph, pool = _graph_with_pool() + graph.__del__() # must not raise + # The pool is deliberately left open rather than closed unreliably; + # scheduling from __del__ is not guaranteed to run. + assert pool.closed is False + + asyncio.run(collect()) + + +def test_del_is_a_no_op_without_a_pool() -> None: + graph = AgentGraph.__new__(AgentGraph) + graph.pool = None + graph.__del__() + + +def test_del_survives_a_failing_close() -> None: + """__del__ must swallow errors; exceptions raised there are printed, not raised.""" + + class _Boom(_FakePool): + async def close(self) -> None: + raise RuntimeError("pool already gone") + + graph = AgentGraph.__new__(AgentGraph) + graph.pool = _Boom() # type: ignore[assignment] + graph.__del__() + # Detach the pool so the object does not retry at interpreter shutdown, when + # logging's handlers are closed and the warning itself fails. That failure + # mode is the reason __del__ is the wrong place for this work at all. + graph.pool = None diff --git a/tests/agent/test_preprocess_concurrency.py b/tests/agent/test_preprocess_concurrency.py new file mode 100644 index 0000000..90fbb38 --- /dev/null +++ b/tests/agent/test_preprocess_concurrency.py @@ -0,0 +1,78 @@ +"""The safety check and language detection must overlap. + +They used to run back to back, costing an extra LLM round trip on every message +on the path every profile shares. The safety check needs the rephrased text so it +has to follow the rephrase; language detection reads the raw input and does not. + +Idea from @bleedblack1 in PR #111. +""" + +import asyncio +import time +from typing import Any + +from langchain_core.runnables import RunnableConfig, RunnableLambda + +from agent.profiles.base import BaseGraphBuilder, BaseState +from agent.tasks.safety_checker import SafetyCheck + +DELAY = 0.2 + + +def _slow( + result: Any, log: list[tuple[str, float, float]], name: str +) -> RunnableLambda: + async def run(_: Any) -> Any: + start = time.perf_counter() + await asyncio.sleep(DELAY) + log.append((name, start, time.perf_counter())) + return result + + return RunnableLambda(run) + + +def _builder(log: list[tuple[str, float, float]]) -> BaseGraphBuilder: + """Build without __init__, which would construct real LLM chains.""" + b = BaseGraphBuilder.__new__(BaseGraphBuilder) + b.rephrase_chain = _slow("rephrased", log, "rephrase") # type: ignore[assignment] + b.safety_checker = _slow( # type: ignore[assignment] + SafetyCheck(safety="true", reason_unsafe=""), log, "safety" + ) + b.language_detector = _slow("en", log, "language") # type: ignore[assignment] + return b + + +def test_safety_and_language_overlap() -> None: + log: list[tuple[str, float, float]] = [] + state = BaseState(user_input="what is TP53?") # type: ignore[typeddict-item] + + elapsed = time.perf_counter() + result = asyncio.run(_builder(log).preprocess(state, RunnableConfig())) + elapsed = time.perf_counter() - elapsed + + assert result["rephrased_input"] == "rephrased" + assert result["safety"] == "true" + assert result["detected_language"] == "en" + + spans = {name: (start, end) for name, start, end in log} + safety, language = spans["safety"], spans["language"] + # Two intervals overlap when each starts before the other ends. + assert safety[0] < language[1], "safety started after language detection finished" + assert language[0] < safety[1], "language detection started after safety finished" + + # Three sequential calls would take 3 * DELAY; overlapping two takes about 2. + assert ( + elapsed < DELAY * 2.8 + ), f"took {elapsed:.2f}s, expected roughly {DELAY * 2:.2f}s" + + +def test_rephrase_still_precedes_the_safety_check() -> None: + """Ordering that must not be lost: the safety check reads the rephrased text.""" + log: list[tuple[str, float, float]] = [] + state = BaseState(user_input="what is TP53?") # type: ignore[typeddict-item] + asyncio.run(_builder(log).preprocess(state, RunnableConfig())) + + spans = {name: (start, end) for name, start, end in log} + assert ( + spans["rephrase"][1] <= spans["safety"][0] + ), "the safety check started before the rephrase finished" diff --git a/tests/data_generation/test_alliance_columns.py b/tests/data_generation/test_alliance_columns.py new file mode 100644 index 0000000..fea1cdb --- /dev/null +++ b/tests/data_generation/test_alliance_columns.py @@ -0,0 +1,64 @@ +"""The Alliance metadata column lists are MITAB schemas, declared twice. + +A missing comma between two adjacent string literals is not a syntax error in +Python -- it silently concatenates them. That had happened three times in +`molecular_interaction`, collapsing seven columns into three, so those metadata +fields were never populated. `genetic_interaction` is the same schema written +correctly, which makes it a usable oracle. + +Note the linter cannot catch this: ruff's ISC001 flags implicit concatenation on +one line, but it is disabled because it conflicts with the formatter -- and the +formatter *joins* such literals, destroying the evidence. Hence a test. +""" + +import ast +from pathlib import Path + +SOURCE = ( + Path(__file__).resolve().parents[2] / "src/data_generation/alliance/__init__.py" +) + + +def _column_lists() -> dict[str, list[str]]: + tree = ast.parse(SOURCE.read_text()) + found: dict[str, list[str]] = {} + for node in ast.walk(tree): + if not isinstance(node, ast.Dict): + continue + for key, value in zip(node.keys, node.values, strict=False): + if isinstance(key, ast.Constant) and isinstance(value, ast.List): + items = [e.value for e in value.elts if isinstance(e, ast.Constant)] + if items and "interactor" in " ".join(items): + found.setdefault(str(key.value), items) + return found + + +def test_the_two_mitab_schemas_match() -> None: + lists = _column_lists() + molecular = lists["molecular_interaction"] + genetic = lists["genetic_interaction"] + assert set(molecular) == set(genetic), ( + "the two MITAB column lists have diverged; a missing comma silently " + "concatenates adjacent entries" + ) + assert len(molecular) == len(genetic) + + +def test_no_column_name_starts_with_another_column_name() -> None: + """Catches the concatenation directly, for lists with no second copy. + + Tests the prefix rather than a substring: "Alt. ID(s) interactor A" legitimately + *contains* "ID(s) interactor A", but a concatenation always begins with the + column that swallowed the comma. + """ + for name, columns in _column_lists().items(): + for column in columns: + prefixes = [ + other + for other in columns + if other != column and column.startswith(other) and len(other) > 8 + ] + assert not prefixes, ( + f"{name}: {column!r} begins with {prefixes} -- " + "likely a missing comma between two adjacent string literals" + )