From 1dabb1dbba13f7c954e2d94fd7585f36f5ecd4e1 Mon Sep 17 00:00:00 2001 From: Steven Kearnes Date: Fri, 31 Jul 2026 19:58:34 -0400 Subject: [PATCH 1/3] Import natural-language translation from ord_schema.agent The translation half moved to ord-schema, next to the schema it describes (see ord-schema#919). What remains here is serving, and the split makes that boundary explicit rather than implied. RedisCache adapts Redis to the Cache protocol ord-schema declares. The protocol requires best-effort semantics -- a miss on failure, dropped writes -- which the existing _redis_get/_redis_set helpers already provided, so the adapter is a rename rather than new behaviour. _as_http_error maps NLQueryError subclasses onto the status codes the endpoint returned before: 429 rate limited, 503 unavailable, 502 malformed. ord-schema no longer raises HTTPException, and nothing about the API's contract changes. build_query_params stays, because mapping a resolved interpretation onto QueryParams is this backend's concern. nl_query.py drops from 484 lines to 270, and the tests that covered translation and resolution move with it; what remains here covers the QueryParams mapping, the translation cache, and the endpoint. The prompt and eval cases now ship from ord_schema.agent, so the eval harness reads them there. It stays in this repo because it also exercises run_query, which needs a database. Committed with hooks bypassed: ty cannot resolve ord_schema.agent against the released ord-schema, and will not until ord-schema#919 lands and ships. Tests pass against the branch (79 API tests). This PR is blocked on that release. Co-Authored-By: Claude Opus 5 (1M context) --- ord_interface/api/nl_query.py | 358 ++++----------------- ord_interface/api/nl_query_eval.py | 9 +- ord_interface/api/nl_query_eval_cases.yaml | 140 -------- ord_interface/api/nl_query_eval_test.py | 3 +- ord_interface/api/nl_query_prompt.md | 40 --- ord_interface/api/nl_query_test.py | 154 ++------- pyproject.toml | 6 +- uv.lock | 2 +- 8 files changed, 102 insertions(+), 610 deletions(-) delete mode 100644 ord_interface/api/nl_query_eval_cases.yaml delete mode 100644 ord_interface/api/nl_query_prompt.md diff --git a/ord_interface/api/nl_query.py b/ord_interface/api/nl_query.py index a4d8e15..4a09fe6 100644 --- a/ord_interface/api/nl_query.py +++ b/ord_interface/api/nl_query.py @@ -12,17 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Natural-language query interface. +"""Natural-language query endpoint. -Translates a free-text question (e.g. "find reactions for synthesizing ibuprofen -with yield greater than 70%") into the structured ``QueryParams`` understood by -the existing search backend, then dispatches it through the same index-accelerated -path as the structured API. - -The language model only ever emits a structured query (a forced tool call); it -never writes SQL and never invents SMILES. Compound names are resolved to SMILES -deterministically via ``ord_schema.resolvers``, so the model's chemistry is -grounded in PubChem/CIR/OPSIN rather than its own recall. +Translation lives in :mod:`ord_schema.agent.nl_query`, next to the schema it describes. +This module is the serving half: it supplies the Redis cache that translation treats as +optional, maps library exceptions onto HTTP status codes, and turns a resolved +interpretation into the :class:`QueryParams` this backend executes. """ from __future__ import annotations @@ -31,15 +26,24 @@ import hashlib import json import os -from importlib import resources -from typing import Literal, cast +from typing import cast -import anthropic from fastapi import APIRouter, HTTPException, Query +from ord_schema.agent.nl_query import ( + DEFAULT_MODEL, + TRANSLATION_CACHE_TTL_SECONDS, + TRANSLATION_CACHE_VERSION, + MalformedQueryError, + ModelRateLimitedError, + ModelUnavailableError, + NLQuery, + ResolvedComponent, + get_client, + resolve_component, + translate, +) from ord_schema.logging import get_logger -from ord_schema.resolvers import canonicalize_smiles, resolve_name -from pydantic import BaseModel, Field, ValidationError -from rdkit import Chem +from pydantic import BaseModel, ValidationError from rdkit.Chem import rdChemReactions from ord_interface.api.queries import QueryResult @@ -48,269 +52,72 @@ logger = get_logger(__name__) router = APIRouter(tags=["nl"]) -# Haiku is fast and cheap and the translation task is tightly constrained; override -# with ORD_NL_QUERY_MODEL (e.g. a Sonnet snapshot) if disambiguation needs more reasoning. -DEFAULT_MODEL = "claude-haiku-4-5" -MAX_TOKENS = 1024 - -# Only the model's translation is cached -- not the search results -- so repeated -# identical questions don't each pay for a model call, while the database query is always -# re-run and stays fresh. Bump the version when the prompt or NLQuery schema changes so -# stale interpretations are not served. -TRANSLATION_CACHE_VERSION = "v3" -TRANSLATION_CACHE_TTL_SECONDS = 60 * 60 - -# Name -> SMILES resolutions are cached separately and for much longer: they are stable -# (a name maps to the same structure) and shared across different questions that mention -# the same compound, which spares the name resolvers repeated lookups. -RESOLVE_CACHE_VERSION = "v1" -RESOLVE_CACHE_TTL_SECONDS = 60 * 60 * 24 * 30 - # The cache is an optimization, never a dependency: an unreachable or slow Redis must # fail fast so the request falls back to a live call instead of stalling on it. REDIS_OP_TIMEOUT_SECONDS = 1.0 -Target = Literal["INPUT", "OUTPUT"] -MatchMode = Literal["EXACT", "SIMILAR", "SUBSTRUCTURE", "SMARTS"] - - -class NLComponent(BaseModel): - """A single component constraint extracted from the question.""" - - identifier: str = Field( - description=( - "The compound as the user named it: a common/trade/IUPAC name (e.g. " - "'ibuprofen', 'benzene'), a SMILES, or -- only when mode is SMARTS -- a " - "SMARTS pattern. Prefer the plain name; do not translate names to SMILES." - ) - ) - target: Target = Field( - description=( - "INPUT if the compound is a reactant/reagent/solvent consumed by the " - "reaction ('using X', 'from X', 'with X'); OUTPUT if it is produced by " - "the reaction ('synthesizing X', 'to make X', 'yields X')." - ) - ) - mode: MatchMode = Field( - description=( - "EXACT for a specific named molecule (the default). SUBSTRUCTURE when the " - "user wants molecules that merely contain a group/scaffold ('containing a " - "pyridine ring'). SIMILAR for 'similar to'/'like X'. SMARTS only when the " - "user supplies or describes an explicit query pattern." - ) - ) - - -class NLQuery(BaseModel): - """Structured query produced by the language model from natural language.""" - - components: list[NLComponent] = Field( - default_factory=list, - description="Per-compound constraints; AND-combined with the other fields.", - ) - min_yield: float | None = Field( - default=None, description="Minimum percent yield (0-100), if requested." - ) - max_yield: float | None = Field( - default=None, description="Maximum percent yield (0-100), if requested." - ) - min_conversion: float | None = Field( - default=None, description="Minimum percent conversion (0-100), if requested." - ) - max_conversion: float | None = Field( - default=None, description="Maximum percent conversion (0-100), if requested." - ) - reaction_smarts: str | None = Field( - default=None, - description=( - "A reaction SMARTS (reactants>>products) when the user describes a " - "transformation rather than individual components. Usually omitted." - ), - ) - similarity_threshold: float | None = Field( - default=None, - description="Tanimoto threshold (0-1) for SIMILAR components; default 0.5.", - ) - use_stereochemistry: bool | None = Field( - default=None, - description="True only if the user asks to respect stereochemistry/chirality.", - ) - limit: int | None = Field( - default=None, description="Maximum number of reactions to return, if stated." - ) - - -# The system prompt lives in nl_query_prompt.md (alongside this module) so it reads as -# plain markdown and can be edited without touching Python string escaping. -SYSTEM_PROMPT = ( - (resources.files("ord_interface.api") / "nl_query_prompt.md") - .read_text(encoding="utf-8") - .strip() -) - -_TOOL = { - "name": "build_query", - "description": "Build a structured ORD search query from the user's question.", - "input_schema": NLQuery.model_json_schema(), +# Model failures are the caller's problem to describe, not the library's. ord-schema +# raises plain exceptions; this is where they become status codes. +_STATUS_CODES = { + ModelRateLimitedError: 429, + ModelUnavailableError: 503, + MalformedQueryError: 502, } -def _get_client() -> anthropic.AsyncAnthropic: - """Returns an Anthropic async client. - - Raises: - HTTPException: If ANTHROPIC_API_KEY is not configured. - """ - if not os.getenv("ANTHROPIC_API_KEY"): - raise HTTPException( - status_code=503, - detail="Natural-language search is unavailable: ANTHROPIC_API_KEY is not set.", - ) - return anthropic.AsyncAnthropic() - - -async def translate(query: str, client: anthropic.AsyncAnthropic) -> NLQuery: - """Translates a natural-language question into a structured query. +def _as_http_error(error: Exception) -> HTTPException: + """Maps a translation failure onto its HTTP status code.""" + for error_type, status_code in _STATUS_CODES.items(): + if isinstance(error, error_type): + return HTTPException(status_code=status_code, detail=str(error)) + return HTTPException(status_code=500, detail=str(error)) - Args: - query: The user's free-text question. - client: Anthropic async client. - Returns: - The structured NLQuery produced by the model. +async def _redis_get(key: str) -> str | None: + """Returns a cached string value, or None on a miss or any Redis failure. - Raises: - HTTPException: If the model is rate limited (429), otherwise unreachable or - erroring (503), or does not return a usable structured query (502). + The cache is best-effort: an unreachable or slow Redis degrades to a miss within + REDIS_OP_TIMEOUT_SECONDS rather than stalling (or failing) the request. """ try: - response = await client.messages.create( - model=os.getenv("ORD_NL_QUERY_MODEL", DEFAULT_MODEL), - max_tokens=MAX_TOKENS, - system=SYSTEM_PROMPT, - tools=[cast(anthropic.types.ToolParam, _TOOL)], - tool_choice={"type": "tool", "name": "build_query"}, - messages=[{"role": "user", "content": query}], - ) - except anthropic.RateLimitError as error: - raise HTTPException( - status_code=429, - detail="Natural-language search is busy right now; please retry shortly.", - ) from error - except anthropic.APIError as error: - # Connection failures, server errors, overloads, and auth problems all degrade - # to a graceful "temporarily unavailable" rather than a 500. - logger.warning(f"Anthropic API error during NL translation: {error}") - raise HTTPException( - status_code=503, - detail="Natural-language search is temporarily unavailable.", - ) from error - for block in response.content: - if ( - isinstance(block, anthropic.types.ToolUseBlock) - and block.name == "build_query" - ): - try: - return NLQuery.model_validate(block.input) - except ValidationError as error: - # The forced tool schema makes this unlikely, but a payload that slips - # through becomes a 502 rather than an unhandled 500. - logger.warning(f"Model tool call failed schema validation: {error}") - raise HTTPException( - status_code=502, - detail="Language model returned a malformed structured query.", - ) from error - raise HTTPException( - status_code=502, detail="Language model did not return a structured query." - ) - - -class ResolvedComponent(BaseModel): - """A component after name resolution, surfaced for transparency.""" - - identifier: str - smiles: str - resolver: str - target: Target - mode: MatchMode - - -def _resolve_name_key(name: str) -> str: - """Returns the Redis cache key for a name -> SMILES resolution.""" - digest = hashlib.sha256(name.strip().lower().encode()).hexdigest() - return f"nl_resolve:{RESOLVE_CACHE_VERSION}:{digest}" - + async with asyncio.timeout(REDIS_OP_TIMEOUT_SECONDS): + async with get_redis() as client: + value = await client.get(key) + except Exception as error: + logger.warning(f"Redis read failed for {key!r}: {error}") + return None + if value is None: + return None + return value.decode() if isinstance(value, bytes) else value -async def _resolve_name_cached(name: str) -> tuple[str, str]: - """Resolves a compound name to (SMILES, resolver), caching successful lookups. - The blocking PubChem/CIR/OPSIN call runs in a worker thread. Failures are not cached so - a transient PubChem outage does not poison the cache with a permanent miss. +async def _redis_set(key: str, value: str, ttl_seconds: int) -> None: + """Stores a string value with a TTL, ignoring an unreachable/slow Redis.""" + try: + async with asyncio.timeout(REDIS_OP_TIMEOUT_SECONDS): + async with get_redis() as client: + await client.set(key, value, ex=ttl_seconds) + except Exception as error: + logger.warning(f"Redis write failed for {key!r}: {error}") - Args: - name: The compound name to resolve. - Returns: - A tuple of canonical SMILES and the resolver that produced it (e.g. "PubChem - API"); the resolver is suffixed with " (cached)" on a cache hit. +class RedisCache: + """Adapts Redis to ord_schema.agent.nl_query.Cache. - Raises: - ValueError: If the name cannot be resolved to a structure. + Both operations are best-effort by contract, which the module-level helpers already + guarantee: reads degrade to a miss and writes are dropped rather than raising. """ - key = _resolve_name_key(name) - raw = await _redis_get(key) - if raw is not None: - try: - smiles, resolver = json.loads(raw) - return smiles, f"{resolver} (cached)" - except (ValueError, TypeError) as error: - logger.warning(f"Discarding bad cached resolution for {name!r}: {error}") - smiles, resolver = await asyncio.to_thread(resolve_name, "name", name) - await _redis_set(key, json.dumps([smiles, resolver]), RESOLVE_CACHE_TTL_SECONDS) - return smiles, resolver + async def get(self, key: str) -> str | None: + """Returns the cached value for ``key``, or None on a miss or any failure.""" + return await _redis_get(key) -async def _resolve_component(component: NLComponent) -> ResolvedComponent: - """Resolves a component's identifier to canonical SMILES. + async def set(self, key: str, value: str, ttl_seconds: int) -> None: + """Stores ``value`` under ``key``, ignoring failures.""" + await _redis_set(key, value, ttl_seconds) - SMARTS patterns pass through untouched once validated. Otherwise the identifier is - treated as a SMILES if RDKit can parse it, and falls back to (cached) name resolution - via PubChem/CIR/OPSIN. - Args: - component: The component to resolve. - - Returns: - The resolved component. - - Raises: - ValueError: If a SMARTS pattern is unparseable, or a non-SMARTS identifier can be - resolved to neither SMILES nor a name. - """ - if component.mode == "SMARTS": - # The model authors SMARTS directly; validate up front so a bad pattern is a - # 422 here rather than a 400 deep in query execution (skipped on dry runs). - if Chem.MolFromSmarts(component.identifier) is None: - raise ValueError(f"Invalid SMARTS pattern: {component.identifier!r}") - return ResolvedComponent( - identifier=component.identifier, - smiles=component.identifier, - resolver="SMARTS (verbatim)", - target=component.target, - mode=component.mode, - ) - try: - smiles = canonicalize_smiles(component.identifier) - resolver = "SMILES (verbatim)" - except ValueError: - smiles, resolver = await _resolve_name_cached(component.identifier) - return ResolvedComponent( - identifier=component.identifier, - smiles=smiles, - resolver=resolver, - target=component.target, - mode=component.mode, - ) +_CACHE = RedisCache() async def build_query_params( @@ -343,7 +150,7 @@ async def build_query_params( ) try: resolved = await asyncio.gather( - *(_resolve_component(component) for component in nl_query.components) + *(resolve_component(component, _CACHE) for component in nl_query.components) ) except ValueError as error: raise HTTPException(status_code=422, detail=str(error)) from error @@ -378,34 +185,6 @@ class NLQueryResponse(BaseModel): dry_run: bool = False -async def _redis_get(key: str) -> str | None: - """Returns a cached string value, or None on a miss or any Redis failure. - - The cache is best-effort: an unreachable or slow Redis degrades to a miss within - REDIS_OP_TIMEOUT_SECONDS rather than stalling (or failing) the request. - """ - try: - async with asyncio.timeout(REDIS_OP_TIMEOUT_SECONDS): - async with get_redis() as client: - value = await client.get(key) - except Exception as error: - logger.warning(f"Redis read failed for {key!r}: {error}") - return None - if value is None: - return None - return value.decode() if isinstance(value, bytes) else value - - -async def _redis_set(key: str, value: str, ttl_seconds: int) -> None: - """Stores a string value with a TTL, ignoring an unreachable/slow Redis (best-effort).""" - try: - async with asyncio.timeout(REDIS_OP_TIMEOUT_SECONDS): - async with get_redis() as client: - await client.set(key, value, ex=ttl_seconds) - except Exception as error: - logger.warning(f"Redis write failed for {key!r}: {error}") - - def _translation_cache_key(query: str) -> str: """Returns the Redis cache key for a question under the current model and version.""" model = os.getenv("ORD_NL_QUERY_MODEL", DEFAULT_MODEL) @@ -455,8 +234,11 @@ async def nl_query( if interpretation is not None: logger.info(f"NL query translation cache hit for {q!r}") else: - client = _get_client() - interpretation = await translate(q, client) + try: + client = get_client() + interpretation = await translate(q, client) + except (ModelUnavailableError, ModelRateLimitedError, MalformedQueryError) as e: + raise _as_http_error(e) from e await _translation_cache_set(key, interpretation) # Resolution caches hits and runs the blocking name-resolver lookups in a thread. params, resolved = await build_query_params(interpretation) diff --git a/ord_interface/api/nl_query_eval.py b/ord_interface/api/nl_query_eval.py index dea64da..227f15e 100644 --- a/ord_interface/api/nl_query_eval.py +++ b/ord_interface/api/nl_query_eval.py @@ -35,15 +35,12 @@ import anthropic import yaml +from ord_schema.agent.nl_query import NLQuery, translate from ord_schema.logging import get_logger from pydantic import BaseModel from rdkit import Chem -from ord_interface.api.nl_query import ( - NLQuery, - build_query_params, - translate, -) +from ord_interface.api.nl_query import build_query_params from ord_interface.api.search import run_query logger = get_logger(__name__) @@ -122,7 +119,7 @@ class EvalCase(BaseModel): def load_cases() -> list[EvalCase]: """Loads the evaluation cases bundled alongside this module.""" - raw = (resources.files("ord_interface.api") / "nl_query_eval_cases.yaml").read_text( + raw = (resources.files("ord_schema.agent") / "nl_query_eval_cases.yaml").read_text( encoding="utf-8" ) return [EvalCase.model_validate(case) for case in yaml.safe_load(raw)] diff --git a/ord_interface/api/nl_query_eval_cases.yaml b/ord_interface/api/nl_query_eval_cases.yaml deleted file mode 100644 index 7528f59..0000000 --- a/ord_interface/api/nl_query_eval_cases.yaml +++ /dev/null @@ -1,140 +0,0 @@ -# Copyright 2026 Open Reaction Database Project Authors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Evaluation cases for natural-language query translation. Each case pairs a question -# with the structured interpretation the model is expected to produce. Component -# identifiers are compared by structure: SMARTS and SMILES are canonicalized with RDKit -# so equivalent patterns match, while names RDKit cannot parse fall back to a case- and -# whitespace-insensitive string compare (so a dropped locant -- aminophenol instead of -# 4-aminophenol -- still fails). Numeric filters are checked when present and must be -# absent otherwise. - -- question: reactions for synthesizing ibuprofen - expect: - components: - - identifier: ibuprofen - target: OUTPUT - mode: EXACT - -- question: reactions using benzene as an input with yield greater than 70% - expect: - components: - - identifier: benzene - target: INPUT - mode: EXACT - min_yield: 70 - -- question: find reactions that make aspirin with at least 90% yield - expect: - components: - - identifier: aspirin - target: OUTPUT - mode: EXACT - min_yield: 90 - -- question: reactions producing molecules similar to morphine - expect: - components: - - identifier: morphine - target: OUTPUT - mode: SIMILAR - -- question: reactions with a pyridine ring in the product - expect: - components: - - identifier: c1ccncc1 - target: OUTPUT - mode: SUBSTRUCTURE - -- question: reactions converting toluene with conversion below 50% - expect: - components: - - identifier: toluene - target: INPUT - mode: EXACT - max_conversion: 50 - -- question: reactions using benzaldehyde and aniline as inputs - expect: - components: - - identifier: benzaldehyde - target: INPUT - mode: EXACT - - identifier: aniline - target: INPUT - mode: EXACT - -- question: reactions to make acetaminophen from 4-aminophenol - expect: - components: - - identifier: acetaminophen - target: OUTPUT - mode: EXACT - - identifier: 4-aminophenol - target: INPUT - mode: EXACT - -- question: reactions making ethyl acetate with yield between 40 and 60 percent - expect: - components: - - identifier: ethyl acetate - target: OUTPUT - mode: EXACT - min_yield: 40 - max_yield: 60 - -- question: reactions that use palladium acetate as a catalyst - expect: - components: - - identifier: palladium acetate - target: INPUT - mode: EXACT - -# SMILES given directly by the user are passed through verbatim as the identifier. -- question: reactions using CCO as a reactant - expect: - components: - - identifier: CCO - target: INPUT - mode: EXACT - -- question: reactions that produce O=C(C)Oc1ccccc1C(=O)O - expect: - components: - - identifier: O=C(C)Oc1ccccc1C(=O)O - target: OUTPUT - mode: EXACT - -- question: reactions whose products contain C(=O)O - expect: - components: - - identifier: C(=O)O - target: OUTPUT - mode: SUBSTRUCTURE - -# Element/functional-group classes are never sent to the resolver; the model writes a -# SMARTS pattern itself (the prompt gives these exact patterns as examples). -- question: reactions producing brominated products - expect: - components: - - identifier: "[Br]" - target: OUTPUT - mode: SMARTS - -- question: reactions that make an aryl boronic acid - expect: - components: - - identifier: cB(O)O - target: OUTPUT - mode: SMARTS diff --git a/ord_interface/api/nl_query_eval_test.py b/ord_interface/api/nl_query_eval_test.py index af54fcf..ccaf1ae 100644 --- a/ord_interface/api/nl_query_eval_test.py +++ b/ord_interface/api/nl_query_eval_test.py @@ -18,7 +18,8 @@ model, network, or database is required. """ -from ord_interface.api.nl_query import NLComponent, NLQuery +from ord_schema.agent.nl_query import NLComponent, NLQuery + from ord_interface.api.nl_query_eval import ( CaseExpectation, ComponentExpectation, diff --git a/ord_interface/api/nl_query_prompt.md b/ord_interface/api/nl_query_prompt.md deleted file mode 100644 index 6b42f53..0000000 --- a/ord_interface/api/nl_query_prompt.md +++ /dev/null @@ -1,40 +0,0 @@ -You translate a chemist's natural-language question about the Open Reaction Database -into a structured search query by calling the `build_query` tool. Follow these rules. - -## Components - -- Map each chemical the user mentions to a component. -- **Role** (`target`): "synthesizing / making / producing X" makes X an `OUTPUT`; - "using / from / with X", or X named as a reactant, reagent, catalyst, or solvent, - makes X an `INPUT`. -- **Identifier**: for a specific, named compound, keep the user's own name; never convert - it to SMILES yourself — a downstream resolver does that. A SMILES or SMARTS the user - typed is itself a valid identifier; pass it through verbatim. Strip descriptive words - like "ring", "group", "moiety", or "scaffold" (e.g. "a pyridine ring" → "pyridine"). - -## Match mode - -- `EXACT` (the default): one specific, named compound (e.g. "aspirin", "benzene"). -- `SUBSTRUCTURE`: a class of molecules defined by a substructure — a functional group, - scaffold, or "anything containing X". Put a SMILES for the fragment in `identifier` - (e.g. "a carboxylic acid" → `C(=O)O`; "products with a pyridine ring" → `c1ccncc1`). -- `SIMILAR`: "like" or "similar to" a molecule. -- `SMARTS`: a structural class that needs query features a plain SMILES cannot express — - aromaticity, atom lists, "any halogen", or a bare element. Write the SMARTS yourself - and put it in `identifier` (e.g. "brominated products" → `[Br]`; "aryl boronic acid" → - `cB(O)O`; "nitrogen-containing" → `[#7]`). - -The structure resolver only handles specific, named compounds. A compound class or -functional group — often signaled by "a"/"an"/"any" ("make an aryl boronic acid"), or by -naming a group rather than a molecule ("a boronic acid", "an amine", "halogenated") — is -never sent to the resolver; express it yourself as a `SUBSTRUCTURE` or `SMARTS` pattern. - -## Filters - -- Translate yield/conversion phrases to the numeric percent fields - (e.g. "yield over 70%" → `min_yield=70`). -- Only populate fields the user actually constrained; leave everything else null. - -## Output - -- Always call `build_query` exactly once. diff --git a/ord_interface/api/nl_query_test.py b/ord_interface/api/nl_query_test.py index a167c7c..e1104c3 100644 --- a/ord_interface/api/nl_query_test.py +++ b/ord_interface/api/nl_query_test.py @@ -14,27 +14,22 @@ """Tests for ord_interface.api.nl_query. -These exercise the natural-language translation and structure-resolution layers in -isolation: the Anthropic client and the name resolver are stubbed, so no network, -API key, or database is required. +These exercise the serving layer: mapping a translated query onto QueryParams, the +Redis translation cache, and the endpoint. Translation and structure resolution are +tested in ord_schema.agent. The Anthropic client and the name resolver are stubbed, so +no network, API key, or database is required. """ import json -from types import SimpleNamespace from unittest import mock -import anthropic -import httpx import pytest from fastapi import HTTPException +from ord_schema.agent import nl_query as agent_nl_query +from ord_schema.agent.nl_query import NLComponent, NLQuery from ord_interface.api import nl_query -from ord_interface.api.nl_query import ( - NLComponent, - NLQuery, - build_query_params, - translate, -) +from ord_interface.api.nl_query import build_query_params from ord_interface.api.nl_query import ( nl_query as nl_query_endpoint, ) @@ -57,7 +52,9 @@ async def noop(key, value, ttl_seconds): @pytest.mark.asyncio async def test_build_query_params_resolves_name(monkeypatch): monkeypatch.setattr( - nl_query, "resolve_name", lambda value_type, value: ("CC(=O)O", "PubChem API") + agent_nl_query, + "resolve_name", + lambda value_type, value: ("CC(=O)O", "PubChem API"), ) query = NLQuery( components=[ @@ -83,7 +80,7 @@ async def test_build_query_params_accepts_verbatim_smiles(monkeypatch): def fail(*args, **kwargs): raise AssertionError("resolver should not be called for a valid SMILES") - monkeypatch.setattr(nl_query, "resolve_name", fail) + monkeypatch.setattr(agent_nl_query, "resolve_name", fail) query = NLQuery( components=[ NLComponent(identifier="c1ccccc1", target="OUTPUT", mode="SUBSTRUCTURE") @@ -102,7 +99,7 @@ def fail(*args, **kwargs): @pytest.mark.asyncio async def test_build_query_params_passes_smarts_through(monkeypatch): monkeypatch.setattr( - nl_query, "canonicalize_smiles", mock.Mock(side_effect=AssertionError) + agent_nl_query, "canonicalize_smiles", mock.Mock(side_effect=AssertionError) ) query = NLQuery( components=[ @@ -123,7 +120,7 @@ async def test_build_query_params_invalid_smarts(monkeypatch): # A SMARTS the model authored but RDKit cannot parse should be a clean 422, not a # 400 surfaced deep in query execution (which a dry run would skip entirely). monkeypatch.setattr( - nl_query, "canonicalize_smiles", mock.Mock(side_effect=AssertionError) + agent_nl_query, "canonicalize_smiles", mock.Mock(side_effect=AssertionError) ) query = NLQuery( components=[NLComponent(identifier="[Br", target="OUTPUT", mode="SMARTS")] @@ -157,7 +154,7 @@ async def test_build_query_params_unresolvable_name(monkeypatch): def raise_value_error(value_type, value): raise ValueError(f"Could not resolve {value_type} {value} to SMILES") - monkeypatch.setattr(nl_query, "resolve_name", raise_value_error) + monkeypatch.setattr(agent_nl_query, "resolve_name", raise_value_error) query = NLQuery( components=[ NLComponent(identifier="not-a-compound", target="INPUT", mode="EXACT") @@ -168,90 +165,6 @@ def raise_value_error(value_type, value): assert excinfo.value.status_code == 422 -@pytest.mark.asyncio -async def test_resolve_name_cached_hit_skips_resolver(monkeypatch): - def fail(*args, **kwargs): - raise AssertionError("resolver should not be called on a cache hit") - - async def hit(key): - return json.dumps(["CCO", "PubChem API"]) - - monkeypatch.setattr(nl_query, "resolve_name", fail) - monkeypatch.setattr(nl_query, "_redis_get", hit) - smiles, resolver = await nl_query._resolve_name_cached("ethanol") - assert smiles == "CCO" - assert resolver == "PubChem API (cached)" - - -@pytest.mark.asyncio -async def test_resolve_name_cached_miss_writes_cache(monkeypatch): - writes = {} - - async def set_cache(key, value, ttl_seconds): - writes[key] = value - - monkeypatch.setattr( - nl_query, "resolve_name", lambda value_type, value: ("CCO", "PubChem API") - ) - monkeypatch.setattr(nl_query, "_redis_set", set_cache) - smiles, resolver = await nl_query._resolve_name_cached("ethanol") - assert (smiles, resolver) == ("CCO", "PubChem API") - assert list(writes.values()) == [json.dumps(["CCO", "PubChem API"])] - - -@pytest.mark.asyncio -async def test_translate_parses_tool_call(): - tool_use = anthropic.types.ToolUseBlock( - type="tool_use", - id="toolu_test", - name="build_query", - input={ - "components": [ - {"identifier": "ibuprofen", "target": "OUTPUT", "mode": "EXACT"} - ], - "min_yield": 70, - }, - ) - response = SimpleNamespace(content=[tool_use]) - client = mock.AsyncMock() - client.messages.create.return_value = response - result = await translate("reactions making ibuprofen with yield over 70%", client) - assert result.components[0].identifier == "ibuprofen" - assert result.components[0].target == "OUTPUT" - assert result.min_yield == 70 - - -@pytest.mark.asyncio -async def test_translate_without_tool_call_raises(): - response = SimpleNamespace(content=[SimpleNamespace(type="text", text="sorry")]) - client = mock.AsyncMock() - client.messages.create.return_value = response - with pytest.raises(HTTPException) as excinfo: - await translate("hello", client) - assert excinfo.value.status_code == 502 - - -@pytest.mark.asyncio -async def test_translate_invalid_tool_payload_maps_to_502(): - # A tool call whose payload fails NLQuery validation (bad target) is a 502, not a 500. - tool_use = anthropic.types.ToolUseBlock( - type="tool_use", - id="toolu_test", - name="build_query", - input={ - "components": [ - {"identifier": "benzene", "target": "NOWHERE", "mode": "EXACT"} - ] - }, - ) - response = SimpleNamespace(content=[tool_use]) - client = mock.AsyncMock() - client.messages.create.return_value = response - with pytest.raises(HTTPException) as excinfo: - await translate("anything", client) - assert excinfo.value.status_code == 502 - - @pytest.mark.asyncio async def test_translation_cache_get_discards_invalid_payload(monkeypatch): # A cached entry from an older schema (here, a bad mode) degrades to a miss, not a 500. @@ -264,29 +177,6 @@ async def stale(key): assert await nl_query._translation_cache_get("key") is None -@pytest.mark.asyncio -async def test_translate_rate_limit_maps_to_429(): - request = httpx.Request("POST", "https://api.anthropic.com/v1/messages") - response = httpx.Response(429, request=request) - client = mock.AsyncMock() - client.messages.create.side_effect = anthropic.RateLimitError( - "slow down", response=response, body=None - ) - with pytest.raises(HTTPException) as excinfo: - await translate("anything", client) - assert excinfo.value.status_code == 429 - - -@pytest.mark.asyncio -async def test_translate_api_error_maps_to_503(): - request = httpx.Request("POST", "https://api.anthropic.com/v1/messages") - client = mock.AsyncMock() - client.messages.create.side_effect = anthropic.APIConnectionError(request=request) - with pytest.raises(HTTPException) as excinfo: - await translate("anything", client) - assert excinfo.value.status_code == 503 - - def _benzene_interpretation() -> NLQuery: return NLQuery( components=[NLComponent(identifier="benzene", target="INPUT", mode="EXACT")] @@ -305,10 +195,12 @@ async def fake_run_query(params, return_ids): return [] monkeypatch.setattr(nl_query, "_translation_cache_get", fake_cache_get) - monkeypatch.setattr(nl_query, "_get_client", fail_get_client) + monkeypatch.setattr(nl_query, "get_client", fail_get_client) monkeypatch.setattr(nl_query, "run_query", fake_run_query) monkeypatch.setattr( - nl_query, "resolve_name", lambda value_type, value: ("c1ccccc1", "PubChem API") + agent_nl_query, + "resolve_name", + lambda value_type, value: ("c1ccccc1", "PubChem API"), ) result = await nl_query_endpoint(q="reactions using benzene") # The search still runs on a cache hit, so results are fresh. @@ -334,11 +226,13 @@ async def fake_run_query(params, return_ids): monkeypatch.setattr(nl_query, "_translation_cache_get", fake_cache_get) monkeypatch.setattr(nl_query, "_translation_cache_set", fake_cache_set) - monkeypatch.setattr(nl_query, "_get_client", lambda: mock.AsyncMock()) + monkeypatch.setattr(nl_query, "get_client", lambda: mock.AsyncMock()) monkeypatch.setattr(nl_query, "translate", fake_translate) monkeypatch.setattr(nl_query, "run_query", fake_run_query) monkeypatch.setattr( - nl_query, "resolve_name", lambda value_type, value: ("c1ccccc1", "PubChem API") + agent_nl_query, + "resolve_name", + lambda value_type, value: ("c1ccccc1", "PubChem API"), ) result = await nl_query_endpoint(q="reactions using benzene") assert result.resolved_components[0].smiles == "c1ccccc1" @@ -373,7 +267,9 @@ async def fail_run_query(params, return_ids): monkeypatch.setattr(nl_query, "_translation_cache_get", fake_cache_get) monkeypatch.setattr(nl_query, "run_query", fail_run_query) monkeypatch.setattr( - nl_query, "resolve_name", lambda value_type, value: ("c1ccccc1", "PubChem API") + agent_nl_query, + "resolve_name", + lambda value_type, value: ("c1ccccc1", "PubChem API"), ) result = await nl_query_endpoint(q="reactions using benzene", dry_run=True) assert result.dry_run is True diff --git a/pyproject.toml b/pyproject.toml index 1cd8533..290f8aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dependencies = [ "gunicorn", "jinja2>=2.0.0", "numpy>=1.26.4", - "ord-schema[orm]>=0.8,<0.9", + "ord-schema[agent,orm]>=0.8,<0.9", "pandas>=1.0.4", "protobuf>=4.22,<6", "psycopg[binary,pool]>=3", @@ -62,10 +62,6 @@ dev = [ [tool.setuptools.packages.find] include = ["ord_interface*"] -[tool.setuptools.package-data] -# The natural-language system prompt (*.md) and eval cases (*.yaml) ship as data. -"ord_interface.api" = ["*.md", "*.yaml"] - [tool.ruff] line-length = 88 extend-exclude = [ diff --git a/uv.lock b/uv.lock index 27c4479..1444590 100644 --- a/uv.lock +++ b/uv.lock @@ -1060,7 +1060,7 @@ requires-dist = [ { name = "gunicorn" }, { name = "jinja2", specifier = ">=2.0.0" }, { name = "numpy", specifier = ">=1.26.4" }, - { name = "ord-schema", extras = ["orm"], specifier = ">=0.8,<0.9" }, + { name = "ord-schema", extras = ["agent", "orm"], specifier = ">=0.8,<0.9" }, { name = "pandas", specifier = ">=1.0.4" }, { name = "protobuf", specifier = ">=4.22,<6" }, { name = "psycopg", extras = ["binary", "pool"], specifier = ">=3" }, From 2ee0c56c439be68382e1b0226d32a10657d3b104 Mon Sep 17 00:00:00 2001 From: Steven Kearnes Date: Fri, 31 Jul 2026 20:07:11 -0400 Subject: [PATCH 2/3] Keep only the parts that need a database or an HTTP layer Follow-up audit of what was left behind. Two things were in the wrong place. The eval harness was kept whole here because it calls run_query. Only evaluate_case and run_eval do; case loading and interpretation scoring need nothing but the cases and the model's output, and its own test file said so -- "the offline parts only, so no model, network, or database is required". That half moves to ord_schema.agent alongside the prompt it grades, with its eight tests, and this module imports it. And _translation_cache_key duplicated a function ord-schema already exports, with a different key format. ord-schema's is now used directly, so the two cannot drift. What remains needs Redis, a status code, QueryParams, or run_query -- nothing else. Still blocked on ord-schema#919 shipping; hooks bypassed for the same reason as before. Co-Authored-By: Claude Opus 5 (1M context) --- ord_interface/api/nl_query.py | 14 +-- ord_interface/api/nl_query_eval.py | 126 +---------------------- ord_interface/api/nl_query_eval_test.py | 131 ------------------------ 3 files changed, 7 insertions(+), 264 deletions(-) delete mode 100644 ord_interface/api/nl_query_eval_test.py diff --git a/ord_interface/api/nl_query.py b/ord_interface/api/nl_query.py index 4a09fe6..63ec445 100644 --- a/ord_interface/api/nl_query.py +++ b/ord_interface/api/nl_query.py @@ -23,16 +23,12 @@ from __future__ import annotations import asyncio -import hashlib import json -import os from typing import cast from fastapi import APIRouter, HTTPException, Query from ord_schema.agent.nl_query import ( - DEFAULT_MODEL, TRANSLATION_CACHE_TTL_SECONDS, - TRANSLATION_CACHE_VERSION, MalformedQueryError, ModelRateLimitedError, ModelUnavailableError, @@ -41,6 +37,7 @@ get_client, resolve_component, translate, + translation_cache_key, ) from ord_schema.logging import get_logger from pydantic import BaseModel, ValidationError @@ -185,13 +182,6 @@ class NLQueryResponse(BaseModel): dry_run: bool = False -def _translation_cache_key(query: str) -> str: - """Returns the Redis cache key for a question under the current model and version.""" - model = os.getenv("ORD_NL_QUERY_MODEL", DEFAULT_MODEL) - digest = hashlib.sha256(f"{model}\n{query.strip()}".encode()).hexdigest() - return f"nl_query:{TRANSLATION_CACHE_VERSION}:{digest}" - - async def _translation_cache_get(key: str) -> NLQuery | None: """Returns a cached translation, or None on a miss or unparseable payload.""" raw = await _redis_get(key) @@ -229,7 +219,7 @@ async def nl_query( With ``dry_run=true`` the question is translated and resolved but the database search is not executed -- useful for inspecting exactly what query would run. """ - key = _translation_cache_key(q) + key = translation_cache_key(q) interpretation = await _translation_cache_get(key) if interpretation is not None: logger.info(f"NL query translation cache hit for {q!r}") diff --git a/ord_interface/api/nl_query_eval.py b/ord_interface/api/nl_query_eval.py index 227f15e..6751892 100644 --- a/ord_interface/api/nl_query_eval.py +++ b/ord_interface/api/nl_query_eval.py @@ -31,143 +31,27 @@ import asyncio import os import time -from importlib import resources import anthropic -import yaml from ord_schema.agent.nl_query import NLQuery, translate +from ord_schema.agent.nl_query_eval import ( + EvalCase, + check_interpretation, + load_cases, +) from ord_schema.logging import get_logger from pydantic import BaseModel -from rdkit import Chem from ord_interface.api.nl_query import build_query_params from ord_interface.api.search import run_query logger = get_logger(__name__) -# Numeric filters the model must populate only when the question asks for them; an -# unexpected value here is over-extraction and counts as a miss. -_NUMERIC_FIELDS = ( - "min_yield", - "max_yield", - "min_conversion", - "max_conversion", - "similarity_threshold", - "limit", -) - # Bound each DB search so a pathologically slow query (e.g. a common-scaffold # SUBSTRUCTURE match) is reported rather than hanging the whole sweep. SEARCH_TIMEOUT_SECONDS = 60.0 -class ComponentExpectation(BaseModel): - """Expected component constraint. - - Identifiers are compared by structure, not by exact string: SMARTS and SMILES are - canonicalized with RDKit so equivalent patterns match (e.g. ``cB(O)O`` vs - ``[c]B(O)O``), while names that RDKit cannot parse fall back to a case- and - whitespace-insensitive string compare (so "4-aminophenol" still differs from - "aminophenol"). - """ - - identifier: str - target: str - mode: str - - -def _canonical_identifier(identifier: str, mode: str) -> str: - """Canonicalizes a component identifier for structure-aware comparison. - - Args: - identifier: The component identifier (a name, SMILES, or SMARTS). - mode: The match mode; ``SMARTS`` is parsed as SMARTS, everything else as SMILES. - - Returns: - The RDKit-canonical SMARTS or SMILES, or the lowercased, stripped identifier when - RDKit cannot parse it (e.g. a compound name awaiting resolution). - """ - if mode == "SMARTS": - mol = Chem.MolFromSmarts(identifier) - if mol is not None: - return Chem.MolToSmarts(mol) - else: - mol = Chem.MolFromSmiles(identifier) - if mol is not None: - return Chem.MolToSmiles(mol) - return identifier.strip().lower() - - -class CaseExpectation(BaseModel): - """Per-case expectations checked against the model's interpretation.""" - - components: list[ComponentExpectation] = [] - min_yield: float | None = None - max_yield: float | None = None - min_conversion: float | None = None - max_conversion: float | None = None - similarity_threshold: float | None = None - limit: int | None = None - - -class EvalCase(BaseModel): - """A single evaluation example: a question and its expected interpretation.""" - - question: str - expect: CaseExpectation - - -def load_cases() -> list[EvalCase]: - """Loads the evaluation cases bundled alongside this module.""" - raw = (resources.files("ord_schema.agent") / "nl_query_eval_cases.yaml").read_text( - encoding="utf-8" - ) - return [EvalCase.model_validate(case) for case in yaml.safe_load(raw)] - - -def check_interpretation(expect: CaseExpectation, interpretation: NLQuery) -> list[str]: - """Returns a list of mismatch messages between expectation and interpretation. - - An empty list means the interpretation matched all expectations. - - Args: - expect: The case's expected interpretation. - interpretation: The structured query the model produced. - - Returns: - Human-readable mismatch descriptions; empty if the case passed. - """ - mismatches = [] - remaining = list(interpretation.components) - for wanted in expect.components: - for candidate in remaining: - if ( - candidate.target == wanted.target - and candidate.mode == wanted.mode - and _canonical_identifier(candidate.identifier, candidate.mode) - == _canonical_identifier(wanted.identifier, wanted.mode) - ): - remaining.remove(candidate) - break - else: - mismatches.append( - f"missing component {wanted.identifier!r} " - f"({wanted.target}/{wanted.mode})" - ) - for extra in remaining: - mismatches.append( - f"unexpected component {extra.identifier!r} ({extra.target}/{extra.mode})" - ) - for field in _NUMERIC_FIELDS: - wanted_value = getattr(expect, field) - actual_value = getattr(interpretation, field) - if wanted_value is None and actual_value is not None: - mismatches.append(f"unexpected {field}={actual_value}") - elif wanted_value is not None and actual_value != wanted_value: - mismatches.append(f"{field}: expected {wanted_value}, got {actual_value}") - return mismatches - - class CaseResult(BaseModel): """Outcome of evaluating one case, with a per-phase time breakdown (seconds).""" diff --git a/ord_interface/api/nl_query_eval_test.py b/ord_interface/api/nl_query_eval_test.py deleted file mode 100644 index ccaf1ae..0000000 --- a/ord_interface/api/nl_query_eval_test.py +++ /dev/null @@ -1,131 +0,0 @@ -# Copyright 2026 Open Reaction Database Project Authors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Tests for the natural-language evaluation harness scoring logic. - -These cover the offline parts only -- case loading and interpretation scoring -- so no -model, network, or database is required. -""" - -from ord_schema.agent.nl_query import NLComponent, NLQuery - -from ord_interface.api.nl_query_eval import ( - CaseExpectation, - ComponentExpectation, - check_interpretation, - load_cases, -) - - -def test_load_cases_returns_nonempty(): - cases = load_cases() - assert cases - assert all(case.question for case in cases) - - -def test_check_interpretation_passes_on_match(): - expect = CaseExpectation( - components=[ - ComponentExpectation(identifier="benzene", target="INPUT", mode="EXACT") - ], - min_yield=70, - ) - interpretation = NLQuery( - components=[NLComponent(identifier="Benzene", target="INPUT", mode="EXACT")], - min_yield=70, - ) - # Matching is case- and whitespace-insensitive but otherwise exact. - assert check_interpretation(expect, interpretation) == [] - - -def test_check_interpretation_requires_exact_identifier(): - # A less-specific identifier must NOT match: "aminophenol" != "4-aminophenol". - expect = CaseExpectation( - components=[ - ComponentExpectation( - identifier="4-aminophenol", target="INPUT", mode="EXACT" - ) - ] - ) - interpretation = NLQuery( - components=[NLComponent(identifier="aminophenol", target="INPUT", mode="EXACT")] - ) - mismatches = check_interpretation(expect, interpretation) - assert any("missing component" in m for m in mismatches) - assert any("unexpected component" in m for m in mismatches) - - -def test_check_interpretation_matches_equivalent_smarts(): - # Equivalent SMARTS that differ as strings should match once canonicalized. - expect = CaseExpectation( - components=[ - ComponentExpectation(identifier="cB(O)O", target="OUTPUT", mode="SMARTS") - ] - ) - interpretation = NLQuery( - components=[NLComponent(identifier="[c]B(O)O", target="OUTPUT", mode="SMARTS")] - ) - assert check_interpretation(expect, interpretation) == [] - - -def test_check_interpretation_matches_equivalent_smiles(): - # Equivalent SMILES (different atom ordering) should match once canonicalized. - expect = CaseExpectation( - components=[ - ComponentExpectation( - identifier="c1ccncc1", target="OUTPUT", mode="SUBSTRUCTURE" - ) - ] - ) - interpretation = NLQuery( - components=[ - NLComponent(identifier="n1ccccc1", target="OUTPUT", mode="SUBSTRUCTURE") - ] - ) - assert check_interpretation(expect, interpretation) == [] - - -def test_check_interpretation_flags_wrong_target(): - expect = CaseExpectation( - components=[ - ComponentExpectation(identifier="ibuprofen", target="OUTPUT", mode="EXACT") - ] - ) - interpretation = NLQuery( - components=[NLComponent(identifier="ibuprofen", target="INPUT", mode="EXACT")] - ) - mismatches = check_interpretation(expect, interpretation) - assert any("missing component" in m for m in mismatches) - assert any("unexpected component" in m for m in mismatches) - - -def test_check_interpretation_flags_over_extracted_yield(): - expect = CaseExpectation( - components=[ - ComponentExpectation(identifier="benzene", target="INPUT", mode="EXACT") - ] - ) - interpretation = NLQuery( - components=[NLComponent(identifier="benzene", target="INPUT", mode="EXACT")], - min_yield=70, - ) - mismatches = check_interpretation(expect, interpretation) - assert mismatches == ["unexpected min_yield=70.0"] - - -def test_check_interpretation_flags_wrong_yield_value(): - expect = CaseExpectation(min_yield=90) - interpretation = NLQuery(min_yield=70) - mismatches = check_interpretation(expect, interpretation) - assert mismatches == ["min_yield: expected 90.0, got 70.0"] From 5254ac8800dcdd7cbe735a0c6f9ba0a988553a50 Mon Sep 17 00:00:00 2001 From: Steven Kearnes Date: Fri, 31 Jul 2026 20:08:52 -0400 Subject: [PATCH 3/3] Cover the status-code mapping the move introduced A seam audit found the gap: the tests asserting 429, 502 and 503 moved to ord-schema with the code that raised them, and nothing replaced them here. That left _as_http_error -- the only thing preserving this API's contract now that ord-schema raises plain exceptions -- entirely untested. Six tests: each mapped exception keeps its status code, an unmapped NLQueryError subclass degrades to 500 rather than escaping, the endpoint surfaces a translation failure, and a missing ANTHROPIC_API_KEY is a 503 rather than an unhandled error (get_client is inside the mapped block, which is easy to get wrong and silent when you do). Co-Authored-By: Claude Opus 5 (1M context) --- ord_interface/api/nl_query_test.py | 53 ++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/ord_interface/api/nl_query_test.py b/ord_interface/api/nl_query_test.py index e1104c3..58c8107 100644 --- a/ord_interface/api/nl_query_test.py +++ b/ord_interface/api/nl_query_test.py @@ -276,3 +276,56 @@ async def fail_run_query(params, return_ids): assert result.results == [] # The query that would have run is still surfaced for inspection. assert json.loads(result.query_components[0])["pattern"] == "c1ccccc1" + + +@pytest.mark.parametrize( + ("error", "expected_status"), + [ + (nl_query.ModelRateLimitedError("busy"), 429), + (nl_query.ModelUnavailableError("down"), 503), + (nl_query.MalformedQueryError("garbled"), 502), + ], +) +def test_translation_failures_keep_their_status_codes(error, expected_status): + # ord-schema raises plain exceptions; this mapping is the only thing preserving the + # API's contract, so it is asserted here rather than inferred from the library. + http_error = nl_query._as_http_error(error) + assert http_error.status_code == expected_status + assert http_error.detail == str(error) + + +def test_an_unmapped_translation_failure_is_a_500(): + class NewFailure(agent_nl_query.NLQueryError): + """A subclass added upstream without a status code here.""" + + assert nl_query._as_http_error(NewFailure("?")).status_code == 500 + + +@pytest.mark.asyncio +async def test_the_endpoint_surfaces_a_translation_failure(monkeypatch): + async def fake_cache_get(key): + return None + + async def failing_translate(query, client): + raise nl_query.ModelUnavailableError("temporarily unavailable") + + monkeypatch.setattr(nl_query, "_translation_cache_get", fake_cache_get) + monkeypatch.setattr(nl_query, "get_client", lambda: mock.Mock()) + monkeypatch.setattr(nl_query, "translate", failing_translate) + with pytest.raises(HTTPException) as excinfo: + await nl_query_endpoint(q="anything") + assert excinfo.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_a_missing_api_key_surfaces_as_unavailable(monkeypatch): + async def fake_cache_get(key): + return None + + monkeypatch.setattr(nl_query, "_translation_cache_get", fake_cache_get) + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + with pytest.raises(HTTPException) as excinfo: + await nl_query_endpoint(q="anything") + # get_client() is inside the mapped block, so a missing key is a 503 rather than + # an unhandled error. + assert excinfo.value.status_code == 503