diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index e27dd5e0..bb47d02b 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -164,6 +164,11 @@ jobs: --baseline-only - name: Start sciencebeam-parser + env: + # An llm profile reads this; every other profile ignores it. Passed as a + # bare `-e NAME` so the value never appears in a command line, and so an + # unset secret leaves it unset in the container rather than empty. + OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} run: | PROFILE_ENV=() if [ -n "${{ env.BENCHMARK_PROFILE }}" ]; then @@ -171,6 +176,7 @@ jobs: fi docker run -d --name sciencebeam-parser -p 8080:8070 \ -e SCIENCEBEAM_PARSER__PRELOAD_ON_STARTUP=true \ + -e OPENROUTER_API_KEY \ "${PROFILE_ENV[@]}" ${{ steps.image_tag.outputs.value }} - name: Wait for parser @@ -182,7 +188,10 @@ jobs: sleep 5 done docker logs sciencebeam-parser >&2 - echo "Parser never became healthy" >&2; exit 1 + echo "Parser never became healthy" >&2 + echo "For an llm profile, PRELOAD_ON_STARTUP validates the endpoint and" >&2 + echo "model id at startup, so a missing OPENROUTER_API_KEY fails here." >&2 + exit 1 - name: Run benchmark env: diff --git a/Makefile b/Makefile index a91e968d..40cd1167 100644 --- a/Makefile +++ b/Makefile @@ -13,6 +13,12 @@ NOT_SLOW_PYTEST_ARGS = -m 'not slow' SCIENCEBEAM_PARSER_PORT = 8080 +# Phoenix renders the llm engine's OpenInference spans. 6006 serves both its UI +# and its OTLP collector. +# Phoenix is the collector used in development; the engine only knows OTLP. +PHOENIX_PORT = 6006 +OTEL_EXPORTER_OTLP_ENDPOINT = http://localhost:$(PHOENIX_PORT) + # Seconds to wait for the parser API on startup. Cold starts re-download pdfalto # + GROBID lexicons, so allow several minutes. API_WAIT_TIMEOUT ?= 300 @@ -109,7 +115,8 @@ dev-install: --dev \ --extra cpu \ --extra delft \ - --extra cv + --extra cv \ + --extra telemetry dev-venv: venv-create dev-install @@ -175,6 +182,14 @@ dev-start: $(PYTHON) -m sciencebeam_parser.service.server --port=$(SCIENCEBEAM_PARSER_PORT) +# The host parser with tracing pointed at the collector. Which profile is active +# is separate and set the usual way, with SCIENCEBEAM_PARSER__PROFILE. Start the +# collector first with docker-start-telemetry. +dev-start-with-telemetry: + OTEL_EXPORTER_OTLP_ENDPOINT=$(OTEL_EXPORTER_OTLP_ENDPOINT) \ + $(MAKE) dev-start + + dev-start-debug: FLASK_ENV=development \ SCIENCEBEAM_PARSER__LOGGING__HANDLERS__LOG_FILE__LEVEL=DEBUG \ @@ -426,6 +441,24 @@ docker-logs: $(DOCKER_COMPOSE) logs -f +# Phoenix renders the llm engine's OpenInference spans. Behind a compose profile, +# so a plain docker-start does not bring an observability server up with it. +docker-start-telemetry: + $(DOCKER_COMPOSE) --profile telemetry up -d phoenix + @echo "Phoenix UI: $(OTEL_EXPORTER_OTLP_ENDPOINT)" + + +# stop + rm rather than `down`, which would take the rest of the stack with it. +# The named volume survives either way, so traces persist across restarts. +docker-stop-telemetry: + $(DOCKER_COMPOSE) --profile telemetry stop phoenix + $(DOCKER_COMPOSE) --profile telemetry rm -f phoenix + + +docker-logs-telemetry: + $(DOCKER_COMPOSE) --profile telemetry logs -f phoenix + + docker-end-to-end-pdfalto: docker-start-and-wait-for-api $(DOCKER_DEV_RUN) curl --fail --show-error --silent \ --form "file=@$(EXAMPLE_PDF_DOCUMENT);filename=$(EXAMPLE_PDF_DOCUMENT)" \ diff --git a/benchmarks/run.py b/benchmarks/run.py index 5106df30..5536608b 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -5,7 +5,7 @@ import subprocess import time from pathlib import Path -from typing import Iterable, List, Optional, Tuple, Union +from typing import Iterable, List, Optional, Sequence, Tuple, Union import httpx import yaml @@ -292,6 +292,39 @@ def run_benchmark( # pylint: disable=too-many-arguments,too-many-positional-arg LOGGER.info("Only one summary available; skipping comparison report") +# Corpora that must not be sent to a third-party model. Provider zero-retention +# does not cover the intermediary, and these manuscripts are not redistributable, +# so an LLM profile and one of these together is refused rather than warned about. +RESTRICTED_CORPORA_FOR_LLM = frozenset({"plos-manuscripts"}) + + +def check_llm_profile_corpora( + config: dict, profile: Optional[str], include: Optional[Sequence[str]] +) -> None: + if not profile: + return + sequence_model_profiles = config.get("sequence_model_profiles") + if not isinstance(sequence_model_profiles, dict): + return + profile_models = sequence_model_profiles.get(profile) + if not isinstance(profile_models, dict): + return + uses_llm = any( + isinstance(model_config, dict) and model_config.get("engine") == "llm" + for model_config in profile_models.values() + ) + if not uses_llm: + return + restricted = sorted(RESTRICTED_CORPORA_FOR_LLM.intersection(include or ())) + if restricted: + raise SystemExit( + f"refusing to run profile {profile!r}, which uses the llm engine, over " + f"{', '.join(restricted)}: those manuscripts are not redistributable and " + "provider zero-retention does not cover the intermediary. Drop the corpus " + "or point the engine at a self-hosted endpoint." + ) + + def main(argv=None) -> None: parser = argparse.ArgumentParser( description="Run benchmark: fetch gold, run baselines, predict, score, compare" @@ -330,6 +363,8 @@ def main(argv=None) -> None: with open(args.config, encoding="utf-8") as f: config = yaml.safe_load(f) + check_llm_profile_corpora(config, args.profile, args.include_corpus) + if args.predictions_repo: store: PredictionsStore = RepoPredictionsStore(Path(args.predictions_repo)) else: diff --git a/benchmarks/tests/llm_guard_test.py b/benchmarks/tests/llm_guard_test.py new file mode 100644 index 00000000..09b11e53 --- /dev/null +++ b/benchmarks/tests/llm_guard_test.py @@ -0,0 +1,44 @@ +import pytest + +from benchmarks.run import check_llm_profile_corpora + + +CONFIG = { + 'sequence_model_profiles': { + 'grobid_crf_0_9_0': { + 'citation': {'engine': 'wapiti', 'path': 'x'}, + }, + 'llm_references': { + 'citation': {'engine': 'llm', 'task': 'citation'}, + }, + }, +} + + +class TestCheckLlmProfileCorpora: + def test_should_allow_a_crf_profile_with_the_private_corpus(self): + check_llm_profile_corpora(CONFIG, 'grobid_crf_0_9_0', ['plos-manuscripts']) + + def test_should_allow_an_llm_profile_without_the_private_corpus(self): + check_llm_profile_corpora(CONFIG, 'llm_references', ['biorxiv']) + + def test_should_allow_an_llm_profile_with_no_opt_in_corpora(self): + check_llm_profile_corpora(CONFIG, 'llm_references', None) + + def test_should_refuse_an_llm_profile_with_the_private_corpus(self): + with pytest.raises(SystemExit, match='not redistributable'): + check_llm_profile_corpora(CONFIG, 'llm_references', ['plos-manuscripts']) + + def test_should_name_the_profile_and_corpus_when_refusing(self): + with pytest.raises(SystemExit) as excinfo: + check_llm_profile_corpora( + CONFIG, 'llm_references', ['biorxiv', 'plos-manuscripts'] + ) + assert 'llm_references' in str(excinfo.value) + assert 'plos-manuscripts' in str(excinfo.value) + + def test_should_ignore_an_unknown_profile(self): + check_llm_profile_corpora(CONFIG, 'nonexistent', ['plos-manuscripts']) + + def test_should_ignore_no_profile(self): + check_llm_profile_corpora(CONFIG, None, ['plos-manuscripts']) diff --git a/doc/llm_engine.md b/doc/llm_engine.md new file mode 100644 index 00000000..7619d25d --- /dev/null +++ b/doc/llm_engine.md @@ -0,0 +1,260 @@ +# LLM engine (experimental) + +A third sequence-model engine alongside `wapiti` and `delft`, serving the `reference_segmenter` and +`citation` models. It is **opt-in**: the shipped default profile stays `grobid_crf`, and a default +install acquires no network dependency and no credential requirement. + +## Using it + +```sh +export OPENROUTER_API_KEY=... # or SCIENCEBEAM_LLM_API_KEY +export SCIENCEBEAM_PARSER__PROFILE=llm_reference_segmenter +``` + +Override the profile by environment rather than editing `profile:` in `config.yml`. The shipped +default is asserted by test, so changing it there fails the build. Env keys use the +`SCIENCEBEAM_PARSER__` prefix with `__` between levels, so a single setting can be overridden the +same way — for example +`SCIENCEBEAM_PARSER__SEQUENCE_MODEL_PROFILES__LLM_REFERENCE_SEGMENTER__CITATION__MODEL`. + +Three profiles, all extending `grobid_crf_0_9_0` so every other model stays on wapiti: + +| profile | replaces | +| --- | --- | +| `llm_reference_segmenter` | `reference_segmenter` | +| `llm_citation` | `citation` | +| `llm_references` | both | + +One per model matters for attribution: when a run fails, the per-model profiles say which model did +it without having to read a stack trace. + +## Configuration + +```yaml +reference_segmenter: + engine: 'llm' + task: 'reference_segmenter' # selects the prompt and the feature layout + response_shape: 'lines' # line numbers where each reference begins + model: 'qwen/qwen3.5-9b' + provider: 'venice' # pinned; routing fails closed without a match + prompt_version: 'lines-v1' # sciencebeam_parser/models/llm/prompts//.md + reasoning: 'off' # models that think by default must be told not to +citation: + engine: 'llm' + task: 'citation' + response_shape: 'values' # field values, located back in the token sequence + model: 'qwen/qwen3.5-9b' + provider: 'siliconflow' + prompt_version: 'values-v1' + reasoning: 'off' +``` + +`response_shape` is configuration rather than a fixed choice, because the best shape differs by +model and by task and moves with each new checkpoint. Comparing shapes is therefore defining a +second profile and running the benchmark, not building a second evaluation route. + +Also accepted: `endpoint` (any OpenAI-compatible base URL, so a self-hosted vLLM works), +`temperature`, `timeout_seconds`, `max_output_tokens`, `max_attempts`, `extra_body`, +`max_references_per_request`, `record_trace_content`, `warn_input_lines`, `max_input_lines`. + +### What the segmenter is told to skip + +The region it receives is whatever segmentation labelled ``, which is sometimes a table +or body text (see +`.project-notes/references-region-includes-non-reference-content.md`). The prompt therefore says to +report only actual bibliographic references and to return an empty list if none of the lines are +references — deciding what is a reference is the segmenter's job, so this is not the same as papering +over the upstream defect in the citation prompt. + +An empty answer is accepted rather than rejected, since "no references in this region" is a valid +thing for the model to conclude. + +### Response shapes for the reference segmenter + +`lines` returns a line number per reference. `evidence` returns the line number **and** the first +words on it — the number is still the answer, the words are only a check on it, so a wrong quote +costs a check rather than the reference. + +The pilot behind this found `evidence` worth about 0.19 `partial_list` on a 12B and 0.003 on the 9B +shipped here, at roughly four times the output tokens, which is why `lines` is the default. On one +real reference list both found 33 references against a gold of 33, in 5.3s and 6.9s. `evidence` earns +its cost on a weaker model, not on a capable one. + +A quote is accepted against the line it names or the one below it, since models name the line +holding the reference number while quoting the words underneath. Anything else increments +`sciencebeam.evidence_mismatches` on the span and logs a warning; set `evidence_mismatch_raises` to +make the check load-bearing instead. It is off by default because a mismatch is a well-formed answer +whose evidence disagrees, not a protocol violation — and in the pilot a 12B mismatched on 158 of 210 +references while the 9B mismatched on 1 of 118. + +### Batching (citation) + +`processor.py` hands the engine every reference of a document at once, so the citation model batches +them: `max_references_per_request` (default 10) references per call, each numbered in the prompt, +with one entry per reference required in the response. Requiring an answer for every reference sent +makes a merged or omitted reference fail validation rather than score, and values are located within +their own reference's tokens only — which also catches the model attributing one reference's author +to another, a mistake a flat field list would have matched against the whole document and labelled +silently. Lowering the bound is the lever if it happens often. + +Both settings are in `config.yml` rather than only in code, since they are the knobs worth turning: + +```yaml +max_references_per_request: 10 # references per call +max_concurrent_requests: 4 # calls in flight at once +``` + +Measured on 10 references: one batched call took 22s against 48s for ten single calls, at token +accuracy 0.906 against 0.904 — twice as fast at no cost in accuracy. On 33 references in 4 batches, +concurrency took 111s down to 58s with accuracy unchanged. + +**Raising the batch size is the wrong lever.** Decode cost is linear in output tokens however they +are grouped, so a bigger batch does not generate less — it generates the same amount in one long +stream that cannot be overlapped, and long generations are what draw provider timeouts (one cost +about five minutes in retry). More, smaller batches running concurrently is the direction that helps. + +Lower `max_concurrent_requests` if the provider answers with 429s; the client retries them with +backoff, but a rate-limited provider can make concurrency a net loss. Note also that spans from +worker threads are siblings rather than nested, and that the service may already handle documents +concurrently, which multiplies with this. + +## What it guarantees + +No text reaches a document that was not in the source. Under `lines` the model returns line numbers +and never text at all. Under `values` it returns text, and every value is located back in the token +sequence — one that cannot be found, or that a previous field already claimed, raises. Either way +`Model._iter_flat_label_model_data_lists_to` independently rejects any result whose tokens are not +the input tokens. + +The `citation` label vocabulary is read from the model's own label map rather than restated in the +prompt source, so it cannot drift from the labels the extractor understands. + +Every request enforces zero data retention — `zdr`, `data_collection: deny`, +`allow_fallbacks: false`, `require_parameters: true`, and `only: [provider]` when pinned. A `:free` +model id is refused at load, because that tier requires allowing training on prompts. + +A response that cannot be decoded raises. There is no fallback to a CRF engine and no partial +labelling: a score is only meaningful if every label came from the model under test. + +**A response the engine cannot parse raises; a claim the engine cannot honour is dropped and +counted.** Bad JSON, a missing reference entry, an index out of range are the first. An invented +value, or two fields over one span, are the second — dropped with a warning naming the reference, +label and text, counted on `sciencebeam.dropped_fields`. + +Dropping satisfies the guarantee rather than weakening it: a discarded value never becomes a label, +so no text reaches the document that was not in the source. Raising would be stronger than the +guarantee needs, and it destroys every reference in a document over one invented field — which is +common, because the region handed over is often not a reference list. Set `dropped_field_raises` for +a run that wants strictness. + +A response cut off at the output limit raises `LlmTruncatedResponseError` naming +`finish_reason`, the completion token count and the task, rather than surfacing as a JSON parse +error. Raise `max_output_tokens`, or send less per request. + +## Tracing (optional) + +Spans follow OpenTelemetry's GenAI semantic conventions — `gen_ai.operation.name`, +`gen_ai.request.model`, `gen_ai.usage.input_tokens` and so on — so they are meaningful in any OTLP +backend. [OpenInference](https://github.com/Arize-ai/openinference) names are emitted alongside +them, so a backend that reads those rather than the GenAI conventions still renders the span as an +LLM call; [Phoenix](https://phoenix.arize.com/) is the one used in development. + +```sh +make dev-install # includes the telemetry extra +make docker-start-telemetry # Phoenix on http://localhost:6006 +export SCIENCEBEAM_PARSER__PROFILE=llm_references # or llm_citation, llm_reference_segmenter +make dev-start-with-telemetry # the host parser, endpoint already set +``` + +Phoenix runs as a compose service behind the `telemetry` profile, so a plain `docker-start` does +not bring an observability server up with it. The image is pinned in +`docker-compose.override.yml`; bump it deliberately rather than tracking `latest`. Traces persist in a named volume across restarts; +`make docker-stop-telemetry` removes the container and keeps them, and +`make docker-logs-telemetry` follows its logs. + +`dev-start-with-telemetry` is the host parser with the collector endpoint already set; it does not +choose a profile, which is set the usual way and is not specific to this engine. A parser running +inside compose would point at `http://phoenix:6006` instead, since `localhost` there is the +container. + +The endpoint variable is what turns emission on, so a parser already running has to be restarted to +start tracing. + +Without the extra installed, or without a collector endpoint set, nothing is emitted and the engine +behaves identically. An existing tracer provider is left alone rather than replaced. + +Configuration is plain OTLP: `OTEL_EXPORTER_OTLP_ENDPOINT` or +`OTEL_EXPORTER_OTLP_TRACES_ENDPOINT`, read by the exporter itself. Nothing in the engine names a +backend — Phoenix is only what listens in development, and any OTLP collector works in its place. +`PHOENIX_COLLECTOR_ENDPOINT` is *not* read: it is Phoenix's own variable, and honouring it would put +one backend's configuration into the engine. + +The response body is attached to the span even when it fails to decode, which is the point: a +truncated or malformed response is visible rather than inferred from an exception. + +**A span carrying prompt text is a copy of manuscript text.** Sending it to a collector on localhost +is not a new disclosure when the same text is already going to the model, but sending it anywhere +else is. Set `record_trace_content: false` in the model config to keep the metrics and drop the +text. + +## Input size + +`sciencebeam.input_lines` and `sciencebeam.input_tokens` are on every span, because what the engine +receives is whatever the *segmentation* model labelled `` — not necessarily a reference +list. A region of a thousand lines is a mislabelled region, and it degrades the `wapiti` path +identically, so it is worth being able to see and filter on. + +Above `warn_input_lines` (default 300) the engine logs a warning naming the count. `max_input_lines` +(default 0, off) raises instead, for a run where failing fast is wanted — off by default because +raising would fail exactly the documents where the CRF path produces poor output, which flatters a +comparison rather than informing it. + +## Choosing a provider + +The same model id served by two providers is not the same service. For +`qwen/qwen3.5-9b`, measured on 33 references in 4 batches: + +| provider | quantisation | seconds | token-acc | reported uptime (30m / 1d) | +| --- | --- | --- | --- | --- | +| Venice | fp8 | **12** | 0.860 | 100% / 100% | +| SiliconFlow | fp8 | 39 | 0.867 | 96% / 99% | + +Venice is shipped: 3.3× faster at the same quantisation and price, for about 0.007 token accuracy — +inside the run-to-run spread of most things measured here, and latency has been the binding +constraint throughout. On the segmenter the same document took 2.1s against 3.9s, with identical +output. + +Worth re-checking rather than trusting: OpenRouter reports `uptime_last_30m` and `uptime_last_1d` per +endpoint, and they move. `GET /api/v1/models/{author}/{slug}/endpoints` lists them alongside +quantisation and per-provider parameter support. Note that latency and throughput often come back +`None` there, so those have to be measured rather than read. + +Not every provider is reachable under the zero-retention pin — Parasail answered 429 immediately — +so a candidate has to be tried, not just looked up. + +## CI + +The benchmark workflow passes `OPENROUTER_API_KEY` from the `benchmark` environment into the parser +container as a bare `-e OPENROUTER_API_KEY`, so the value never appears in a command line and an +unset secret leaves it unset rather than empty. Every non-LLM profile ignores it, so it is passed +unconditionally. + +Select the engine the same way as any other configuration: a `profile:llm_references` label on the +PR, or the `profile` input on `workflow_dispatch`. + +`SCIENCEBEAM_PARSER__PRELOAD_ON_STARTUP=true` is already set there, which means an LLM profile +validates its endpoint and model id while the container starts. A missing key or an unreachable +endpoint therefore fails at "Wait for parser" with the reason in the container logs, rather than +part-way through a run. + +**PLOS cannot be combined with an LLM profile.** `benchmarks/run.py` refuses it, naming the profile +and the corpus, because provider zero-retention does not cover OpenRouter itself and those +manuscripts are not redistributable. This matters most on `main`, where the workflow adds +`--include-corpus plos-manuscripts` automatically — so an LLM profile on `main` fails rather than +quietly sending private manuscripts to a third party. Point the engine at a self-hosted endpoint if +that corpus needs covering. + +Nothing is traced in CI: no collector endpoint is set, so the engine emits no spans. + +Unit tests reach no network. The engine's tests are fixture-driven and the client is a `Protocol`, so +the ordinary CI job needs no secret at all. diff --git a/docker-compose.override.yml b/docker-compose.override.yml index 22fe6f4a..e7cbfaa9 100644 --- a/docker-compose.override.yml +++ b/docker-compose.override.yml @@ -36,5 +36,26 @@ services: - ./benchmarks:/opt/sciencebeam_parser/benchmarks - huggingface-cache:/root/.cache/huggingface + # Tracing for the llm engine, dev-only and opt-in: the profile keeps it out of + # a plain `docker compose up`. 6006 serves both the UI and the OTLP collector. + # A parser running in compose would point at http://phoenix:6006; one running on + # the host at http://localhost:6006. + phoenix: + # Pinned: a dev tool whose version moves under you makes "it worked + # yesterday" unanswerable, and the trace schema is what the spans are read + # through. Tags also exist unprefixed (20.3.0) and as -nonroot / -debug. + image: arizephoenix/phoenix:version-20.3.0 + profiles: + - telemetry + ports: + - "6006:6006" + environment: + # Phoenix keeps traces in a sqlite db under its working dir; without a + # volume they are lost on restart, which defeats comparing runs. + - PHOENIX_WORKING_DIR=/mnt/phoenix + volumes: + - phoenix-data:/mnt/phoenix + volumes: huggingface-cache: + phoenix-data: diff --git a/pyproject.toml b/pyproject.toml index aa08b970..54b888be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,12 @@ delft = [ cv = [ "layoutparser==0.3.2", ] +# Tracing for the llm engine. Optional: without it, and without a collector +# endpoint, the engine emits nothing and behaves identically. +telemetry = [ + "opentelemetry-sdk>=1.29", + "opentelemetry-exporter-otlp-proto-http>=1.29", +] [dependency-groups] benchmark = [ diff --git a/sciencebeam_parser/models/llm/__init__.py b/sciencebeam_parser/models/llm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/sciencebeam_parser/models/llm/client.py b/sciencebeam_parser/models/llm/client.py new file mode 100644 index 00000000..5cf8889b --- /dev/null +++ b/sciencebeam_parser/models/llm/client.py @@ -0,0 +1,157 @@ +import logging +import time +from typing import Any, Dict, Mapping, Optional, Protocol + +import httpx + +from sciencebeam_parser.models.llm.config import LlmEngineConfig, get_api_key + + +LOGGER = logging.getLogger(__name__) + +RETRY_STATUS_CODES = frozenset({408, 429, 500, 502, 503, 504}) + + +def get_error_status_code(response_json: Mapping[str, Any]) -> Optional[int]: + """OpenRouter returns some upstream failures as http 200 with an error body, + so the status code alone does not say whether a call is worth retrying.""" + error = response_json.get('error') + if not isinstance(error, dict): + return None + code = error.get('code') + return code if isinstance(code, int) else 0 + + +class LlmRequestError(RuntimeError): + pass + + +class LlmCompletionClient(Protocol): + def validate_configuration(self) -> None: + ... + + def get_completion( + self, prompt: str, response_schema: Mapping[str, Any] + ) -> Mapping[str, Any]: + ... + + +class LlmClient: + def __init__(self, config: LlmEngineConfig): + self.config = config + + def _headers(self) -> Dict[str, str]: + return {'Authorization': f'Bearer {get_api_key()}'} + + def _request_body(self, prompt: str, response_schema: Mapping[str, Any]) -> Dict[str, Any]: + body: Dict[str, Any] = { + 'model': self.config.model, + 'temperature': self.config.temperature, + 'max_tokens': self.config.max_output_tokens, + 'messages': [{'role': 'user', 'content': prompt}], + 'response_format': { + 'type': 'json_schema', + 'json_schema': { + 'name': 'sciencebeam_labels', + 'strict': True, + 'schema': response_schema, + }, + }, + 'provider': self.config.provider_routing, + **self.config.extra_body, + } + if self.config.reasoning == 'off': + body['reasoning'] = {'enabled': False} + return body + + def validate_configuration(self) -> None: + """Fails at load rather than at first request, and spends no tokens.""" + url = f'{self.config.endpoint.rstrip("/")}/models' + try: + with httpx.Client(timeout=self.config.timeout_seconds) as client: + response = client.get(url, headers=self._headers()) + except httpx.HTTPError as exc: + raise LlmRequestError(f'{self.config.endpoint} is not reachable: {exc}') from exc + if response.status_code != 200: + raise LlmRequestError( + f'{url} returned {response.status_code}: {response.text[:200]}' + ) + model_ids = { + entry.get('id') for entry in response.json().get('data', []) + if isinstance(entry, dict) + } + if model_ids and self.config.model not in model_ids: + raise LlmRequestError( + f'{self.config.endpoint} does not offer model {self.config.model!r}' + ) + LOGGER.info( + 'llm engine configured: model=%r provider=%r prompt=%r shape=%r', + self.config.model, self.config.provider, self.config.prompt_version, + self.config.response_shape + ) + + def get_completion(self, prompt: str, response_schema: Mapping[str, Any]) -> Mapping[str, Any]: + return self._post_with_retry(prompt, response_schema) + + def _post_with_retry( + self, prompt: str, response_schema: Mapping[str, Any] + ) -> Mapping[str, Any]: + url = f'{self.config.endpoint.rstrip("/")}/chat/completions' + body = self._request_body(prompt, response_schema) + last_error = '' + for attempt in range(self.config.max_attempts): + if attempt: + time.sleep(2 ** attempt) + try: + with httpx.Client(timeout=self.config.timeout_seconds) as client: + response = client.post(url, headers=self._headers(), json=body) + except httpx.HTTPError as exc: + last_error = f'{type(exc).__name__}: {exc}' + continue + if response.status_code in RETRY_STATUS_CODES: + last_error = f'http {response.status_code}: {response.text[:200]}' + continue + if response.status_code != 200: + raise LlmRequestError( + f'http {response.status_code}: {response.text[:200]}' + ) + response_json = response.json() + error_code = get_error_status_code(response_json) + if error_code is None: + return response_json + last_error = f'error body {error_code}: {str(response_json)[:200]}' + if error_code in RETRY_STATUS_CODES: + continue + raise LlmRequestError(last_error) + raise LlmRequestError( + f'giving up after {self.config.max_attempts} attempts: {last_error}' + ) + + +class LlmTruncatedResponseError(LlmRequestError): + pass + + +def get_response_content(response_json: Mapping[str, Any]) -> str: + choices = response_json.get('choices') + if not choices: + raise LlmRequestError(f'response has no choices: {str(response_json)[:200]}') + choice = choices[0] + content = choice.get('message', {}).get('content') + finish_reason = choice.get('finish_reason') or choice.get('native_finish_reason') + if finish_reason == 'length': + completion_tokens = ( + response_json.get('usage', {}).get('completion_tokens') + ) + raise LlmTruncatedResponseError( + 'response hit the output token limit' + f' (finish_reason={finish_reason!r},' + f' completion_tokens={completion_tokens},' + f' chars={len(content or "")});' + ' raise max_output_tokens or reduce the input per request' + ) + if not content: + raise LlmRequestError( + f'response content is empty (finish_reason={finish_reason!r})' + ) + return content diff --git a/sciencebeam_parser/models/llm/config.py b/sciencebeam_parser/models/llm/config.py new file mode 100644 index 00000000..83270ee1 --- /dev/null +++ b/sciencebeam_parser/models/llm/config.py @@ -0,0 +1,74 @@ +import os +from dataclasses import dataclass, field, fields +from typing import Any, Dict, Mapping, Optional + + +API_KEY_ENV_NAMES = ('SCIENCEBEAM_LLM_API_KEY', 'OPENROUTER_API_KEY') + +DEFAULT_ENDPOINT = 'https://openrouter.ai/api/v1' + + +class LlmConfigError(ValueError): + pass + + +@dataclass(frozen=True) +class LlmEngineConfig: + task: str + model: str + prompt_version: str + response_shape: str = 'lines' + endpoint: str = DEFAULT_ENDPOINT + provider: Optional[str] = None + reasoning: str = '' + temperature: float = 0.0 + timeout_seconds: float = 300.0 + max_output_tokens: int = 8000 + max_attempts: int = 4 + record_trace_content: bool = True + warn_input_lines: int = 300 + max_input_lines: int = 0 + max_references_per_request: int = 10 + max_concurrent_requests: int = 4 + evidence_mismatch_raises: bool = False + dropped_field_raises: bool = False + extra_body: Dict[str, Any] = field(default_factory=dict) + + @staticmethod + def from_model_config(config: Mapping[str, Any]) -> 'LlmEngineConfig': + for required in ('task', 'model', 'prompt_version'): + if not config.get(required): + raise LlmConfigError(f'llm engine requires {required!r} in the model config') + model = config['model'] + if model.endswith(':free'): + raise LlmConfigError( + f'refusing model id {model!r}: the free tier requires allowing training on' + ' prompts, which the zero-retention requirement forbids' + ) + known = {field_.name for field_ in fields(LlmEngineConfig)} + return LlmEngineConfig(**{ + key: value for key, value in config.items() + if key in known + }) + + @property + def provider_routing(self) -> Dict[str, Any]: + routing: Dict[str, Any] = { + 'zdr': True, + 'data_collection': 'deny', + 'allow_fallbacks': False, + 'require_parameters': True, + } + if self.provider: + routing['only'] = [self.provider] + return routing + + +def get_api_key() -> str: + for name in API_KEY_ENV_NAMES: + value = os.environ.get(name) + if value: + return value + raise LlmConfigError( + 'no api key: set one of ' + ', '.join(API_KEY_ENV_NAMES) + ) diff --git a/sciencebeam_parser/models/llm/decode.py b/sciencebeam_parser/models/llm/decode.py new file mode 100644 index 00000000..6cda8f32 --- /dev/null +++ b/sciencebeam_parser/models/llm/decode.py @@ -0,0 +1,266 @@ +import json +import re +from typing import Any, List, Mapping, Sequence, Tuple + + +LINE_START = 'LINESTART' + +LABEL_ONLY_LINE = re.compile(r'^[\[(]?\d{1,3}[\])]?[.)]?$') + +WORD_SEPARATOR = re.compile(r'[^0-9A-Za-zÀ-ɏ]+') + + +def iter_words(text: str) -> List[str]: + """Words with punctuation removed, so a quote can be compared with tokens. + + One implementation, imported by the value shapes too: two of these drifting + apart is how a match silently stops matching. + """ + return [word for word in WORD_SEPARATOR.split(text) if word] + + +LINES_RESPONSE_SCHEMA: Mapping[str, Any] = { + 'type': 'object', + 'additionalProperties': False, + 'required': ['starts'], + 'properties': { + 'starts': { + 'type': 'array', + 'items': {'type': 'integer'}, + }, + }, +} + + +EVIDENCE_RESPONSE_SCHEMA: Mapping[str, Any] = { + 'type': 'object', + 'additionalProperties': False, + 'required': ['references'], + 'properties': { + 'references': { + 'type': 'array', + 'items': { + 'type': 'object', + 'additionalProperties': False, + 'required': ['line', 'starts_with'], + 'properties': { + 'line': {'type': 'integer'}, + 'starts_with': {'type': 'string', 'maxLength': 48}, + }, + }, + }, + }, +} + + +class LlmResponseError(ValueError): + pass + + +class LlmInputTooLargeError(ValueError): + pass + + +def get_line_numbers(line_status_values: Sequence[str]) -> List[int]: + line_numbers: List[int] = [] + current = -1 + for index, status in enumerate(line_status_values): + if status == LINE_START or index == 0: + current += 1 + line_numbers.append(current) + return line_numbers + + +def get_lines(tokens: Sequence[str], line_numbers: Sequence[int]) -> List[List[str]]: + lines: List[List[str]] = [] + for token, line_number in zip(tokens, line_numbers): + while len(lines) <= line_number: + lines.append([]) + lines[line_number].append(token) + return lines + + +def render_numbered_lines(tokens: Sequence[str], line_numbers: Sequence[int]) -> str: + return '\n'.join( + f'{number}\t' + ' '.join(line_tokens) + for number, line_tokens in enumerate(get_lines(tokens, line_numbers)) + ) + + +def parse_line_starts(content: str, line_count: int) -> List[int]: + try: + payload = json.loads(content) + except ValueError as exc: + raise LlmResponseError(f'response is not json: {exc}') from exc + if not isinstance(payload, dict) or 'starts' not in payload: + raise LlmResponseError('response has no "starts"') + starts = payload['starts'] + if not isinstance(starts, list): + raise LlmResponseError('"starts" is not a list') + resolved: List[int] = [] + for value in starts: + if isinstance(value, bool) or not isinstance(value, int): + raise LlmResponseError(f'line number is not an integer: {value!r}') + if not 0 <= value < line_count: + raise LlmResponseError( + f'line number {value} out of range for {line_count} lines' + ) + resolved.append(value) + if resolved != sorted(set(resolved)): + raise LlmResponseError(f'line numbers are not strictly ascending: {resolved}') + return resolved + + +def snap_starts_to_label_lines( + line_starts: Sequence[int], + lines: Sequence[Sequence[str]] +) -> List[int]: + existing = set(line_starts) + snapped: List[int] = [] + for start in line_starts: + previous = start - 1 + if ( + previous >= 0 + and previous not in existing + and previous not in snapped + and LABEL_ONLY_LINE.match(''.join(lines[previous])) + ): + snapped.append(previous) + continue + snapped.append(start) + return sorted(snapped) + + +def parse_evidence_line_starts( + content: str, + line_count: int +) -> Tuple[List[int], List[str]]: + """Line numbers are the answer; the quoted words are only evidence for them. + + Keeping the payload an index means a wrong quote costs a check rather than + the reference, which is the failure mode the anchor shape had. + """ + try: + payload = json.loads(content) + except ValueError as exc: + raise LlmResponseError(f'response is not json: {exc}') from exc + if not isinstance(payload, dict) or 'references' not in payload: + raise LlmResponseError('response has no "references"') + entries = payload['references'] + if not isinstance(entries, list): + raise LlmResponseError('"references" is not a list') + starts: List[int] = [] + claims: List[str] = [] + for entry in entries: + if not isinstance(entry, dict) or 'line' not in entry: + raise LlmResponseError(f'reference entry is malformed: {entry!r}') + line = entry['line'] + if isinstance(line, bool) or not isinstance(line, int): + raise LlmResponseError(f'line number is not an integer: {line!r}') + if not 0 <= line < line_count: + raise LlmResponseError( + f'line number {line} out of range for {line_count} lines' + ) + starts.append(line) + claims.append(str(entry.get('starts_with') or '')) + if starts != sorted(set(starts)): + raise LlmResponseError(f'line numbers are not strictly ascending: {starts}') + return starts, claims + + +def count_evidence_mismatches( + starts: Sequence[int], + claims: Sequence[str], + lines: Sequence[Sequence[str]] +) -> int: + """A quote is accepted against the line it names or the one below it. + + Models name the line holding the reference number while quoting the words on + the line under it, which is the training convention rather than an error. + """ + mismatches = 0 + for line, claim in zip(starts, claims): + wanted = iter_words(claim) + if not wanted: + continue + candidates = [] + for offset in (0, 1): + if 0 <= line + offset < len(lines): + candidates.append( + iter_words(' '.join(lines[line + offset]))[:len(wanted)] + ) + if wanted not in candidates: + mismatches += 1 + return mismatches + + +def decode_evidence_response( + content: str, + tokens: Sequence[str], + line_status_values: Sequence[str] +) -> Tuple[List[Tuple[str, str]], int]: + if len(tokens) != len(line_status_values): + raise LlmResponseError( + f'token count {len(tokens)} does not match feature rows' + f' {len(line_status_values)}' + ) + line_numbers = get_line_numbers(line_status_values) + lines = get_lines(tokens, line_numbers) + starts, claims = parse_evidence_line_starts(content, line_count=len(lines)) + mismatches = count_evidence_mismatches(starts, claims, lines) + snapped = snap_starts_to_label_lines(starts, lines) + labels = iter_labels_for_line_starts(tokens, line_numbers, snapped) + return list(zip(tokens, labels)), mismatches + + +def iter_labels_for_line_starts( + tokens: Sequence[str], + line_numbers: Sequence[int], + line_starts: Sequence[int] +) -> List[str]: + lines = get_lines(tokens, line_numbers) + starts = set(line_starts) + labels: List[str] = [] + started = False + for token_index, line_number in enumerate(line_numbers): + is_line_start = token_index == 0 or line_numbers[token_index - 1] != line_number + if line_number in starts and is_line_start: + started = True + if LABEL_ONLY_LINE.match(''.join(lines[line_number])): + labels.append('B-