diff --git a/.github/workflows/python-CI.yml b/.github/workflows/python-CI.yml index b90cd3e5859..2a3281bb3d2 100644 --- a/.github/workflows/python-CI.yml +++ b/.github/workflows/python-CI.yml @@ -37,6 +37,7 @@ jobs: uv_lock: ${{ steps.filter.outputs.uv_lock }} json_canonicalization_schema: ${{ steps.filter.outputs.json_canonicalization_schema }} filter_dsl: ${{ steps.filter.outputs.filter_dsl }} + datagen_tooling: ${{ steps.filter.outputs.datagen_tooling }} steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: @@ -83,6 +84,9 @@ jobs: - "js/app/src/pages/project/sessionFilterDSL.ts" - "src/phoenix/trace/dsl/**" - "scripts/ci/check_filter_dsl_snippets.py" + datagen_tooling: + - "scripts/datagen/**" + - "src/phoenix/experimental/datagen/**" - name: Print Filters env: IPYNB: ${{ steps.filter.outputs.ipynb }} @@ -563,6 +567,35 @@ jobs: timeout-minutes: 60 run: uvx tox run -e unit_tests -- -ra --reruns 5 --db postgresql -n 16 --dist loadscope --postgresql-exec /usr/lib/postgresql/14/bin/pg_ctl + datagen-tooling-tests: + name: Datagen Tooling Tests + runs-on: ubuntu-latest + needs: changes + if: ${{ needs.changes.outputs.datagen_tooling == 'true' && github.event_name == 'pull_request' }} + steps: + - name: Checkout repository + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + sparse-checkout: | + requirements/ + scripts/datagen/ + src/ + packages/ + evals/ + - name: Set up Python + uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: "3.10" + - name: Set up `uv` + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + version: "0.12.5" + - name: Sync dependencies + run: uv sync --frozen + - name: Run datagen tooling tests + run: uv run pytest scripts/datagen/tests -ra + integration-tests: name: Integration Tests runs-on: ${{ matrix.os }} @@ -749,6 +782,7 @@ jobs: - check-lockfile - type-check - unit-tests + - datagen-tooling-tests - integration-tests - test-migrations - test-json-canonicalization-schema diff --git a/scripts/datagen/README.md b/scripts/datagen/README.md new file mode 100644 index 00000000000..9e7b58af1af --- /dev/null +++ b/scripts/datagen/README.md @@ -0,0 +1,205 @@ +# Trace corpus recorders + +These scripts record application traffic through real OpenInference instrumenters. +The resulting corpus contains raw OTLP protobuf JSON requests and the fragment rows used by the +Phoenix datagen composer. A **fragment** is one replayable unit — a conversation turn or an agent +episode — pointing at the recorded traces it produced. Recording frameworks remain outside Phoenix +runtime dependencies. + +## Fixed inputs + +`recorder_fixtures.json` contains the application inputs for every retained recorder: + +- a stable fragment ID; +- an archetype (which kind of application produced the trace — plain chat, RAG, tool agent, + graph agent, guardrails, or structured extraction) and a domain (its subject area, such as + customer support or coding); +- direct prompts, turns, documents, or expected structured values. + +The fixtures contain no sampling weights or generated text. Each recorder receives a +`RecorderFixture`, appends OTLP requests to `traces.jsonl`, and returns the trace IDs it emitted. +`record_fixture` then appends the matching fragment row (`fragment_id`, `archetype`, `domain`, +`trace_ids`) to `fragments.jsonl`. + +The fixture set includes multiple examples for plain chat, RAG, tool agents, graph agents, +guardrails, and structured extraction. Success, blocked, redacted, conflicting-source, and +incomplete-input examples are represented directly in the app inputs. + +Tool-agent fixtures may carry `prompt_variants`: alternative phrasings of the opening prompt. +Live recording picks one phrasing per run, so repeated runs of the same task do not open with +identical text. Coding fixtures without a deterministic scripted episode are skipped by +scripted auto-selection and record live only. + +## Offline providers and tools + +`ScriptedOpenAIProvider` serves a fixed sequence of text, tool-call, or HTTP responses through an +in-process `httpx` transport. It supports buffered and streaming chat completions without a network +connection or API key. + +`local_tools` exposes deterministic document search, record lookup, status lookup, arithmetic, +and ticket creation over `tool_fixtures.json`. The file contains separate customer-support, +analytics, and coding data sets. + +## Generate varied recordings + +`organic_conditions.json` defines authored input variations: degraded versions of the base +fixture inputs that let response-quality issues arise naturally rather than by script. Each +condition names a base fixture, a unique output fragment ID, an intensity, and one payload per +intensity level. Document edits, replacements at existing fixture-input paths, and +matched local-tool result overlays are applied before the application runs. Input replacements +cannot add or remove structure. Keep every condition fragment ID distinct from the IDs in +`recorder_fixtures.json` and from other conditions. Intensity selects the level: + +| Intensity | Level | +| --------- | ----- | +| below 0.2 | `subtle` | +| below 0.5 | `moderate` | +| 0.5 and above | `strong` | + +All recorder commands accept `--condition` and `--append`. With neither flag, a recorder uses its +fixed fixtures and resets the output directory. `--condition` runs the one fixture with that +condition's edits applied; `--append` preserves existing rows so multiple conditions and +recorders can share a recording directory. + +Plain chat, RAG, tool agent, and structured extraction also accept `--provider scripted|live`. +Scripted is the default. Live recording requires an explicit `--model`, reads `OPENAI_API_KEY`, and +uses `OPENAI_BASE_URL` when it is set. For example: + +```console +export OPENAI_API_KEY="..." +uv run --script scripts/datagen/tool_agent.py \ + --output-dir dist/datagen/recording \ + --condition support-stale-delivery-status \ + --provider live \ + --model gpt-5.4 \ + --append +``` + +Graph and guardrail recorders are deterministic applications and therefore expose condition and +append controls without provider or model options. + +Live plain-chat conversations run until the simulated user closes them. A target turn count +(drawn per fixture, or set with `--target-turns`) controls when the simulator is told to wrap +up once its current concern is addressed; the conversation ends at that natural closing +message, with a hard cap at twice the target. Live model aliases: `luna` and `terra` resolve +to their provider model IDs with tool-calling options applied. + +A live run records every instrumented invocation that emits trace IDs. Responses are not compared +with fixture-authored answers, and incomplete responses or traced application errors are retained. +A run fails only when it emits no trace IDs. Review or evaluate quality after recording; keep +ambiguous outcomes in the recorded set. + +### Recording playbook (all recorders, one directory) + +Choose one live model for the batch and run these commands in order. The first command starts a new +recording and captures every scripted tool fixture, including the longer coding sessions. Each later +command appends either all base fixtures for one recorder or its authored condition. The resulting +recording contains more than two dozen fragments across all six archetypes. + +```console +export DATAGEN_MODEL="gpt-5-mini" + +uv run --script scripts/datagen/tool_agent.py \ + --output-dir dist/datagen/recording + +uv run --script scripts/datagen/openai_chat_sessions.py \ + --output-dir dist/datagen/recording \ + --provider live --model "$DATAGEN_MODEL" --append +uv run --script scripts/datagen/openai_chat_sessions.py \ + --output-dir dist/datagen/recording \ + --condition support-late-express-ambiguity \ + --provider live --model "$DATAGEN_MODEL" --append + +uv run --script scripts/datagen/llama_index_rag.py \ + --output-dir dist/datagen/recording \ + --provider live --model "$DATAGEN_MODEL" --append +uv run --script scripts/datagen/llama_index_rag.py \ + --output-dir dist/datagen/recording \ + --condition research-fleet-delivery-pressure \ + --provider live --model "$DATAGEN_MODEL" --append + +uv run --script scripts/datagen/tool_agent.py \ + --output-dir dist/datagen/recording \ + --provider live --model "$DATAGEN_MODEL" --append +uv run --script scripts/datagen/tool_agent.py \ + --output-dir dist/datagen/recording \ + --condition support-stale-delivery-status \ + --provider live --model "$DATAGEN_MODEL" --append + +uv run --script scripts/datagen/structured_extraction.py \ + --output-dir dist/datagen/recording \ + --provider live --model "$DATAGEN_MODEL" --append +uv run --script scripts/datagen/structured_extraction.py \ + --output-dir dist/datagen/recording \ + --condition analytics-corrected-refund-export \ + --provider live --model "$DATAGEN_MODEL" --append + +uv run --script scripts/datagen/graph_multi_agent.py \ + --output-dir dist/datagen/recording --append +uv run --script scripts/datagen/guardrailed_app.py \ + --output-dir dist/datagen/recording --append +``` + +Keep every command result that reports trace IDs, including responses that are incomplete, +ambiguous, or accompanied by a traced application error. If a command reports no trace IDs, fix +that recorder before continuing so later `--append` calls do not hide the missing fragment. + +Anyone — or any coding agent — running this playbook can choose the conditions, models, run +count, and command order. To use Codex with ChatGPT subscription access, authenticate once and ask +the non-interactive command to inspect the playbook and invoke recorder commands: + +```console +codex login +codex exec 'Read scripts/datagen/README.md and scripts/datagen/organic_conditions.json. Choose varied conditions and models, run the applicable recorders into dist/datagen/recording with --append, and retain every run that emits trace IDs.' +``` + +`codex exec` chooses and runs commands in this workflow; it is not a provider implemented by the +recorders. Direct `--provider live` recorder calls use API credentials from the environment. + +## Recorder environments + +Every framework recorder has a PEP 723 dependency block and must be run with `uv run --script`. +This keeps recorder dependencies out of the Phoenix package and pins the instrumenter stack used +to create stored traces. Shared modules imported by those entry points use only the Python standard +library unless their dependency is declared in every importing script. + +Each JSONL line in `traces.jsonl` is one protobuf-JSON `ExportTraceServiceRequest`. A single trace +may span multiple rows. + +Tests for the packaging pipeline (conditions, packer, fetcher roundtrip) live in `tests/` next to +this file and run with `uv run pytest scripts/datagen/tests`. They are separate from the Phoenix +unit test suite: CI runs them in the Datagen Tooling Tests job when files under `scripts/datagen/` +or `src/phoenix/experimental/datagen/` change. Recorder behavior has no unit tests; verify recorders by +generating a corpus. + +## Package a corpus + +After all selected fixtures and conditions have been recorded into one directory, package the +recording. The printed statistics include `opening_diversity_by_domain` — distinct opening +inputs per domain — so low seed variety is visible before publication: + +```console +uv run python -m scripts.datagen.corpus \ + --archive dist/datagen/corpus.tar.gz +``` + +The archive contains only `fragments.jsonl` and `traces.jsonl`. + +Validate or stage the archive for manual publication: + +```console +uv run python -m scripts.datagen.publish validate \ + --archive dist/datagen/corpus.tar.gz + +uv run python -m scripts.datagen.publish prepare-archive \ + --archive dist/datagen/corpus.tar.gz \ + --output-dir dist/datagen-publication +``` + +Preparation prints the exact commands for uploading the digest-addressed archive first and the +public pointer second. + +## Freshness + +Re-record and review the corpus whenever a pinned instrumenter version changes. This keeps stored +span shapes aligned with the framework and instrumenter versions declared by each recorder. diff --git a/scripts/datagen/conditions.py b/scripts/datagen/conditions.py new file mode 100644 index 00000000000..496a6690487 --- /dev/null +++ b/scripts/datagen/conditions.py @@ -0,0 +1,445 @@ +"""Authored input conditions for trace corpus recorders.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from math import isfinite +from pathlib import Path +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Literal, cast + +if TYPE_CHECKING or __package__: + from scripts.datagen.fake_tools import ( + ToolError, + ToolPatchOperation, + ToolResultOverlay, + load_fixture_sets, + validate_result_overlays, + ) + from scripts.datagen.recording import RecorderFixture, load_fixtures +else: + from fake_tools import ( + ToolError, + ToolPatchOperation, + ToolResultOverlay, + load_fixture_sets, + validate_result_overlays, + ) + from recording import RecorderFixture, load_fixtures + +Strength = Literal["subtle", "moderate", "strong"] + + +class ConditionError(ValueError): + """Raised when an authored condition cannot be materialized.""" + + +@dataclass(frozen=True) +class ConditionedFixture: + fixture: RecorderFixture + tool_fixture_set: Mapping[str, Any] | None + tool_result_overlays: tuple[ToolResultOverlay, ...] + + +@dataclass(frozen=True) +class _DocumentEdit: + target: str + document_id: str + operation: str + source: str | None = None + replacement: str | None = None + text: str | None = None + + +@dataclass(frozen=True) +class _InputReplacement: + path: str + value: Any + + +@dataclass(frozen=True) +class _Payload: + input_replacements: tuple[_InputReplacement, ...] + document_edits: tuple[_DocumentEdit, ...] + tool_overlays: tuple[ToolResultOverlay, ...] + + +@dataclass(frozen=True) +class _Condition: + condition_id: str + fixture_id: str + fragment_id: str + intensity: float + strengths: Mapping[Strength, _Payload] + + +def strength_for_intensity(intensity: float) -> Strength: + """Map an authored intensity to its condition strength.""" + if isinstance(intensity, bool) or not isinstance(intensity, (int, float)): + raise ConditionError("condition intensity must be a number between 0 and 1") + numeric = float(intensity) + if not isfinite(numeric) or not 0 <= numeric <= 1: + raise ConditionError("condition intensity must be a number between 0 and 1") + if numeric < 0.2: + return "subtle" + if numeric < 0.5: + return "moderate" + return "strong" + + +def materialize_condition( + condition_id: str, + path: Path | None = None, + *, + fixtures: Sequence[RecorderFixture] | None = None, + fixture_sets: Mapping[str, Mapping[str, Any]] | None = None, +) -> ConditionedFixture: + """Materialize one condition without changing its base fixture data.""" + source = path or Path(__file__).with_name("organic_conditions.json") + conditions = _load_condition_file(source) + available_fixtures = tuple(fixtures) if fixtures is not None else load_fixtures() + fixture_by_id = {fixture.fragment_id: fixture for fixture in available_fixtures} + if len(fixture_by_id) != len(available_fixtures): + raise ConditionError("base recorder fixture IDs must be unique") + available_fixture_sets = fixture_sets if fixture_sets is not None else load_fixture_sets() + + output_ids = [condition.fragment_id for condition in conditions] + if len(set(output_ids)) != len(output_ids): + raise ConditionError(f"condition fragment IDs in {source} must be unique") + base_ids = set(fixture_by_id) + if collisions := base_ids & set(output_ids): + raise ConditionError( + f"condition fragment IDs in {source} collide with base fixtures: {sorted(collisions)}" + ) + + selected: tuple[_Condition, RecorderFixture, _Payload] | None = None + for condition in conditions: + try: + fixture = fixture_by_id[condition.fixture_id] + except KeyError as error: + raise ConditionError( + f"condition {condition.condition_id!r} names unknown fixture " + f"{condition.fixture_id!r}" + ) from error + for payload in condition.strengths.values(): + _materialize_payload(fixture, payload, available_fixture_sets) + if condition.condition_id == condition_id: + strength = strength_for_intensity(condition.intensity) + selected = (condition, fixture, condition.strengths[strength]) + + if selected is None: + raise ConditionError(f"unknown condition {condition_id!r} in {source}") + condition, fixture, payload = selected + inputs, tool_fixture_set = _materialize_payload(fixture, payload, available_fixture_sets) + conditioned = RecorderFixture( + fragment_id=condition.fragment_id, + archetype=fixture.archetype, + domain=fixture.domain, + inputs=inputs, + ) + frozen_fixture_set = ( + cast(Mapping[str, Any], _freeze_json(tool_fixture_set)) + if tool_fixture_set is not None + else None + ) + return ConditionedFixture(conditioned, frozen_fixture_set, payload.tool_overlays) + + +def _load_condition_file(source: Path) -> tuple[_Condition, ...]: + try: + value = json.loads(source.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ConditionError(f"unable to load conditions from {source}: {error}") from error + if not isinstance(value, list) or not value: + raise ConditionError(f"conditions in {source} must be a non-empty array") + conditions = tuple(_parse_condition(item, source) for item in value) + ids = [condition.condition_id for condition in conditions] + if len(set(ids)) != len(ids): + raise ConditionError(f"condition IDs in {source} must be unique") + return conditions + + +def _parse_condition(value: Any, source: Path) -> _Condition: + raw = _object(value, "condition", source) + _exact_fields( + raw, + {"condition_id", "fixture_id", "fragment_id", "intensity", "strengths"}, + "condition", + source, + ) + condition_id = _string(raw, "condition_id", "condition", source) + fixture_id = _string(raw, "fixture_id", condition_id, source) + fragment_id = _string(raw, "fragment_id", condition_id, source) + intensity = raw["intensity"] + strength_for_intensity(intensity) + strengths = _object(raw["strengths"], f"condition {condition_id!r}.strengths", source) + _exact_fields( + strengths, + {"subtle", "moderate", "strong"}, + f"condition {condition_id!r}.strengths", + source, + ) + payloads: dict[Strength, _Payload] = { + strength: _parse_payload( + strengths[strength], f"condition {condition_id!r}.{strength}", source + ) + for strength in ("subtle", "moderate", "strong") + } + return _Condition( + condition_id, + fixture_id, + fragment_id, + float(intensity), + MappingProxyType(payloads), + ) + + +def _parse_payload(value: Any, field: str, source: Path) -> _Payload: + raw = _object(value, field, source) + unknown = set(raw) - {"input_replacements", "document_edits", "tool_overlays"} + if unknown: + raise ConditionError(f"{field} in {source} has unknown fields: {sorted(unknown)}") + replacements_value = raw.get("input_replacements", []) + edits_value = raw.get("document_edits", []) + overlays_value = raw.get("tool_overlays", []) + if ( + not isinstance(replacements_value, list) + or not isinstance(edits_value, list) + or not isinstance(overlays_value, list) + ): + raise ConditionError( + f"{field} in {source} replacements, edits, and overlays must be arrays" + ) + replacements = tuple( + _parse_input_replacement(item, f"{field}.input_replacements[{index}]", source) + for index, item in enumerate(replacements_value) + ) + edits = tuple( + _parse_document_edit(item, f"{field}.document_edits[{index}]", source) + for index, item in enumerate(edits_value) + ) + overlays = tuple( + _parse_tool_overlay(item, f"{field}.tool_overlays[{index}]", source) + for index, item in enumerate(overlays_value) + ) + if not replacements and not edits and not overlays: + raise ConditionError(f"{field} in {source} must define a replacement, edit, or overlay") + return _Payload(replacements, edits, overlays) + + +def _parse_input_replacement(value: Any, field: str, source: Path) -> _InputReplacement: + raw = _object(value, field, source) + _exact_fields(raw, {"path", "value"}, field, source) + return _InputReplacement(_string(raw, "path", field, source), raw["value"]) + + +def _parse_document_edit(value: Any, field: str, source: Path) -> _DocumentEdit: + raw = _object(value, field, source) + operation = raw.get("operation") + if operation == "replace_once": + expected = {"target", "document_id", "operation", "source", "replacement"} + _exact_fields(raw, expected, field, source) + return _DocumentEdit( + _choice(raw, "target", {"fixture", "tool_corpus"}, field, source), + _string(raw, "document_id", field, source), + operation, + source=_string(raw, "source", field, source), + replacement=_string(raw, "replacement", field, source, allow_empty=True), + ) + if operation == "append": + expected = {"target", "document_id", "operation", "text"} + _exact_fields(raw, expected, field, source) + return _DocumentEdit( + _choice(raw, "target", {"fixture", "tool_corpus"}, field, source), + _string(raw, "document_id", field, source), + operation, + text=_string(raw, "text", field, source), + ) + raise ConditionError(f"{field}.operation in {source} must be replace_once or append") + + +def _parse_tool_overlay(value: Any, field: str, source: Path) -> ToolResultOverlay: + raw = _object(value, field, source) + _exact_fields(raw, {"tool_name", "match_arguments", "operations"}, field, source) + match_arguments = _object(raw["match_arguments"], f"{field}.match_arguments", source) + operations_value = raw["operations"] + if not isinstance(operations_value, list) or not operations_value: + raise ConditionError(f"{field}.operations in {source} must be a non-empty array") + try: + operations = tuple( + _parse_tool_operation(item, f"{field}.operations[{index}]", source) + for index, item in enumerate(operations_value) + ) + return ToolResultOverlay( + _string(raw, "tool_name", field, source), match_arguments, operations + ) + except ToolError as error: + raise ConditionError(f"invalid {field} in {source}: {error}") from error + + +def _parse_tool_operation(value: Any, field: str, source: Path) -> ToolPatchOperation: + raw = _object(value, field, source) + operation = raw.get("operation") + expected = {"operation", "path"} if operation == "remove" else {"operation", "path", "value"} + _exact_fields(raw, expected, field, source) + if operation not in {"add", "replace", "remove"}: + raise ConditionError(f"{field}.operation in {source} must be add, replace, or remove") + path = _string(raw, "path", field, source) + if operation == "remove": + return ToolPatchOperation(operation, path) + return ToolPatchOperation(operation, path, raw["value"]) + + +def _materialize_payload( + fixture: RecorderFixture, + payload: _Payload, + fixture_sets: Mapping[str, Mapping[str, Any]], +) -> tuple[dict[str, Any], dict[str, Any] | None]: + inputs = cast(dict[str, Any], _json_copy(dict(fixture.inputs))) + base_tool_fixture_set = fixture_sets.get(fixture.domain) + tool_fixture_set = ( + cast(dict[str, Any], _json_copy(base_tool_fixture_set)) + if base_tool_fixture_set is not None + else None + ) + for replacement in payload.input_replacements: + _apply_input_replacement(inputs, replacement) + for edit in payload.document_edits: + if edit.target == "fixture": + documents = inputs.get("documents") + else: + documents = tool_fixture_set.get("documents") if tool_fixture_set is not None else None + if not isinstance(documents, list): + raise ConditionError(f"fixture {fixture.fragment_id!r} has no {edit.target} documents") + _apply_document_edit(documents, edit) + if payload.tool_overlays: + if tool_fixture_set is None: + raise ConditionError(f"fixture {fixture.fragment_id!r} has no local tool fixture set") + try: + validate_result_overlays(tool_fixture_set, payload.tool_overlays) + except ToolError as error: + raise ConditionError( + f"fixture {fixture.fragment_id!r} has invalid tool overlays: {error}" + ) from error + return inputs, tool_fixture_set + + +def _apply_input_replacement(inputs: dict[str, Any], replacement: _InputReplacement) -> None: + if not replacement.path.startswith("/"): + raise ConditionError(f"input replacement path {replacement.path!r} must start with '/'") + parts = tuple( + part.replace("~1", "/").replace("~0", "~") + for part in replacement.path.removeprefix("/").split("/") + ) + current: Any = inputs + for part in parts[:-1]: + if isinstance(current, dict) and part in current: + current = current[part] + elif isinstance(current, list) and part.isdecimal() and int(part) < len(current): + current = current[int(part)] + else: + raise ConditionError(f"input replacement path {replacement.path!r} does not exist") + leaf = parts[-1] + if isinstance(current, dict) and leaf in current: + current[leaf] = _replacement_value(current[leaf], replacement) + elif isinstance(current, list) and leaf.isdecimal() and int(leaf) < len(current): + index = int(leaf) + current[index] = _replacement_value(current[index], replacement) + else: + raise ConditionError(f"input replacement path {replacement.path!r} does not exist") + + +def _replacement_value(current: Any, replacement: _InputReplacement) -> Any: + if isinstance(current, (dict, list)) or isinstance(replacement.value, (dict, list)): + raise ConditionError( + f"input replacement path {replacement.path!r} must replace a scalar value" + ) + return _json_copy(replacement.value) + + +def _apply_document_edit(documents: list[Any], edit: _DocumentEdit) -> None: + matches = [ + document + for document in documents + if isinstance(document, dict) + and edit.document_id in (document.get("id"), document.get("source"), document.get("name")) + ] + if len(matches) != 1: + raise ConditionError( + f"document {edit.document_id!r} in {edit.target} matched {len(matches)} times" + ) + document = matches[0] + content = document.get("text") + if not isinstance(content, str): + raise ConditionError(f"document {edit.document_id!r} must contain text") + if edit.operation == "append": + document["text"] = content + cast(str, edit.text) + return + source = cast(str, edit.source) + count = content.count(source) + if count != 1: + raise ConditionError( + f"replace_once source for document {edit.document_id!r} matched {count} times" + ) + document["text"] = content.replace(source, cast(str, edit.replacement), 1) + + +def _object(value: Any, field: str, source: Path) -> dict[str, Any]: + if not isinstance(value, dict): + raise ConditionError(f"{field} in {source} must be an object") + return value + + +def _exact_fields(value: Mapping[str, Any], expected: set[str], field: str, source: Path) -> None: + if set(value) != expected: + raise ConditionError(f"{field} in {source} must define exactly {sorted(expected)}") + + +def _string( + value: Mapping[str, Any], + name: str, + field: str, + source: Path, + *, + allow_empty: bool = False, +) -> str: + item = value.get(name) + if not isinstance(item, str) or (not allow_empty and not item): + qualifier = "a string" if allow_empty else "a non-empty string" + raise ConditionError(f"{field}.{name} in {source} must be {qualifier}") + return item + + +def _choice( + value: Mapping[str, Any], + name: str, + choices: set[str], + field: str, + source: Path, +) -> str: + item = value.get(name) + if item not in choices: + raise ConditionError(f"{field}.{name} in {source} must be one of {sorted(choices)}") + return cast(str, item) + + +def _json_copy(value: Any) -> Any: + return json.loads(json.dumps(_plain_json(value))) + + +def _freeze_json(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze_json(item) for key, item in value.items()}) + if isinstance(value, (list, tuple)): + return tuple(_freeze_json(item) for item in value) + return value + + +def _plain_json(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _plain_json(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain_json(item) for item in value] + return value diff --git a/scripts/datagen/corpus.py b/scripts/datagen/corpus.py new file mode 100644 index 00000000000..2861e5ac658 --- /dev/null +++ b/scripts/datagen/corpus.py @@ -0,0 +1,277 @@ +"""Build a replay corpus from recorded fragment and trace rows.""" + +from __future__ import annotations + +import argparse +import gzip +import io +import json +import os +import sys +import tarfile +import tempfile +from collections import Counter +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from statistics import fmean, median +from typing import Any, Iterator, Mapping, Sequence, TextIO + +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, +) +from opentelemetry.proto.trace.v1.trace_pb2 import Span + +from phoenix.experimental.datagen.loader import Corpus, CorpusError, load_corpus + +_ARCHIVE_MEMBERS = ("fragments.jsonl", "traces.jsonl") +_FRAGMENT_FIELDS = ("fragment_id", "archetype", "domain", "trace_ids") +_SPAN_KIND = "openinference.span.kind" +_SESSION_ID = "session.id" +_INPUT_VALUE = "input.value" + + +@dataclass(frozen=True) +class CorpusPackage: + path: Path + sha256: str + size_bytes: int + fragment_count: int + trace_count: int + span_count: int + span_kind_counts: Mapping[str, int] + span_kind_shares: Mapping[str, float] + spans_per_trace: Mapping[str, int | float] + tool_span_count: int + tool_span_share: float + llm_turns_by_session: Mapping[str, int] + llm_turns_per_session: Mapping[str, int | float] + opening_diversity_by_domain: Mapping[str, Mapping[str, int]] + + +class CorpusArchiveError(ValueError): + """Raised when recorded rows cannot form a corpus archive.""" + + +def package_corpus(source: Path, destination: Path) -> CorpusPackage: + """Package recorded rows atomically and verify the completed archive.""" + fragments_bytes = _project_fragments(_read_bytes(source / "fragments.jsonl")) + traces_bytes = _read_bytes(source / "traces.jsonl") + corpus = _write_archive_atomic( + destination, + { + "fragments.jsonl": fragments_bytes, + "traces.jsonl": traces_bytes, + }, + ) + archive_bytes = destination.read_bytes() + statistics = _corpus_statistics(corpus) + return CorpusPackage( + path=destination, + sha256=sha256(archive_bytes).hexdigest(), + size_bytes=len(archive_bytes), + fragment_count=len(corpus.fragments), + trace_count=len(corpus.requests), + **statistics, + ) + + +def _corpus_statistics(corpus: Corpus) -> dict[str, Any]: + spans_per_trace = [sum(1 for _ in _iter_spans(request)) for request in corpus.requests] + span_kind_counts: Counter[str] = Counter() + llm_turns_by_session: Counter[str] = Counter() + for request in corpus.requests: + for span in _iter_spans(request): + span_kind = _string_attribute(span, _SPAN_KIND) or "UNKNOWN" + span_kind_counts[span_kind] += 1 + if session_id := _string_attribute(span, _SESSION_ID): + llm_turns_by_session.setdefault(session_id, 0) + if span_kind == "LLM": + llm_turns_by_session[session_id] += 1 + + openings_by_domain: dict[str, set[str]] = {} + fragments_by_domain: Counter[str] = Counter() + for fragment in corpus.fragments: + fragments_by_domain[fragment.domain] += 1 + request = corpus.requests_by_trace_id.get(fragment.trace_ids[0]) + if request is None: + continue + for span in _iter_spans(request): + if span.parent_span_id: + continue + if opening := _string_attribute(span, _INPUT_VALUE): + openings_by_domain.setdefault(fragment.domain, set()).add(opening[:200]) + break + + span_count = sum(span_kind_counts.values()) + span_kind_counts.setdefault("TOOL", 0) + span_kind_shares = { + span_kind: count / span_count for span_kind, count in span_kind_counts.items() + } + tool_span_count = span_kind_counts["TOOL"] + return { + "span_count": span_count, + "span_kind_counts": dict(span_kind_counts), + "span_kind_shares": span_kind_shares, + "spans_per_trace": _distribution(spans_per_trace), + "tool_span_count": tool_span_count, + "tool_span_share": tool_span_count / span_count, + "llm_turns_by_session": dict(llm_turns_by_session), + "llm_turns_per_session": _distribution(list(llm_turns_by_session.values())), + "opening_diversity_by_domain": { + domain: { + "fragments": fragments_by_domain[domain], + "distinct_openings": len(openings_by_domain.get(domain, set())), + } + for domain in sorted(fragments_by_domain) + }, + } + + +def _iter_spans(request: ExportTraceServiceRequest) -> Iterator[Span]: + for resource_spans in request.resource_spans: + for scope_spans in resource_spans.scope_spans: + yield from scope_spans.spans + + +def _string_attribute(span: Span, key: str) -> str | None: + attribute = next((attribute for attribute in span.attributes if attribute.key == key), None) + if attribute is None or attribute.value.WhichOneof("value") != "string_value": + return None + return attribute.value.string_value or None + + +def _distribution(values: Sequence[int]) -> dict[str, int | float]: + if not values: + return {} + return { + "min": min(values), + "median": median(values), + "mean": fmean(values), + "max": max(values), + } + + +def _read_bytes(path: Path) -> bytes: + try: + return path.read_bytes() + except OSError as error: + raise CorpusArchiveError(f"unable to read recorded corpus file {path}: {error}") from error + + +def _project_fragments(content: bytes) -> bytes: + try: + lines = content.decode("utf-8").splitlines() + except UnicodeDecodeError as error: + raise CorpusArchiveError(f"fragments.jsonl is not valid UTF-8: {error}") from error + documents = [] + for line_number, line in enumerate(lines, start=1): + if not line.strip(): + continue + try: + value = json.loads(line) + except json.JSONDecodeError as error: + raise CorpusArchiveError( + f"invalid fragments.jsonl entry at line {line_number}: {error}" + ) from error + if not isinstance(value, Mapping): + raise CorpusArchiveError( + f"fragments.jsonl entry at line {line_number} must be a JSON object" + ) + documents.append({field: value[field] for field in _FRAGMENT_FIELDS if field in value}) + if not documents: + raise CorpusArchiveError("fragments.jsonl contains no fragments") + return b"".join(_canonical_bytes(document) + b"\n" for document in documents) + + +def _canonical_bytes(value: Mapping[str, Any]) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + + +def _write_archive_atomic(destination: Path, files: Mapping[str, bytes]) -> Corpus: + destination.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + dir=destination.parent, prefix=f".{destination.name}.", suffix=".tmp" + ) + temporary = Path(temporary_name) + try: + with os.fdopen(descriptor, "wb") as raw: + with gzip.GzipFile(filename="", mode="wb", fileobj=raw, mtime=0) as compressed: + with tarfile.open( + fileobj=compressed, mode="w", format=tarfile.PAX_FORMAT + ) as archive: + for filename in _ARCHIVE_MEMBERS: + content = files[filename] + info = tarfile.TarInfo(filename) + info.size = len(content) + info.mtime = 0 + info.mode = 0o644 + info.uid = 0 + info.gid = 0 + info.uname = "" + info.gname = "" + archive.addfile(info, fileobj=io.BytesIO(content)) + raw.flush() + os.fsync(raw.fileno()) + try: + corpus = load_corpus(temporary) + except CorpusError as error: + raise CorpusArchiveError(f"invalid corpus: {error}") from error + os.replace(temporary, destination) + directory_descriptor = os.open(destination.parent, os.O_RDONLY) + try: + os.fsync(directory_descriptor) + finally: + os.close(directory_descriptor) + return corpus + except Exception: + temporary.unlink(missing_ok=True) + raise + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source", type=Path, help="directory containing recorded corpus rows") + parser.add_argument("--archive", type=Path, required=True) + return parser + + +def command( + argv: Sequence[str] | None = None, + *, + stdout: TextIO = sys.stdout, + stderr: TextIO = sys.stderr, +) -> int: + args = build_parser().parse_args(argv) + try: + package = package_corpus(args.source, args.archive) + except (CorpusArchiveError, OSError, ValueError) as error: + print(json.dumps({"error": type(error).__name__, "message": str(error)}), file=stderr) + return 2 + print(json.dumps(_package_document(package), indent=2, sort_keys=True), file=stdout) + return 0 + + +def _package_document(package: CorpusPackage) -> dict[str, Any]: + return { + "archive": str(package.path), + "sha256": package.sha256, + "size_bytes": package.size_bytes, + "fragment_count": package.fragment_count, + "trace_count": package.trace_count, + "span_count": package.span_count, + "span_kind_counts": dict(package.span_kind_counts), + "span_kind_shares": dict(package.span_kind_shares), + "spans_per_trace": dict(package.spans_per_trace), + "tool_span_count": package.tool_span_count, + "tool_span_share": package.tool_span_share, + "llm_turns_by_session": dict(package.llm_turns_by_session), + "llm_turns_per_session": dict(package.llm_turns_per_session), + "opening_diversity_by_domain": { + domain: dict(counts) for domain, counts in package.opening_diversity_by_domain.items() + }, + } + + +if __name__ == "__main__": + raise SystemExit(command()) diff --git a/scripts/datagen/fake_tools.py b/scripts/datagen/fake_tools.py new file mode 100644 index 00000000000..0150b11204b --- /dev/null +++ b/scripts/datagen/fake_tools.py @@ -0,0 +1,841 @@ +"""Small deterministic tool set for offline trace recorders.""" + +from __future__ import annotations + +import ast +import json +import math +import operator +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from dataclasses import field as dataclass_field +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Any, TypeAlias, cast + +JSON: TypeAlias = None | bool | int | float | str | list["JSON"] | dict[str, "JSON"] +ToolResult: TypeAlias = dict[str, JSON] + +_WORD = re.compile(r"[a-z0-9]+") + + +class ToolError(ValueError): + """Raised when local tool data or arguments are invalid.""" + + +@dataclass +class CodingRepository: + files: dict[str, str] + search_page_size: int + read_chunk_lines: int + transient_failures: dict[str, int] + tests: dict[str, tuple[str, str]] + search_cursors: dict[str, int] = dataclass_field(default_factory=dict) + read_cursors: dict[str, int] = dataclass_field(default_factory=dict) + test_runs: dict[str, int] = dataclass_field(default_factory=dict) + + @classmethod + def from_fixture_set(cls, fixture_set: Mapping[str, Any]) -> CodingRepository | None: + if fixture_set.get("name") != "coding_agent": + return None + repository = fixture_set.get("repository") + if not isinstance(repository, Mapping): + raise ToolError("coding_agent fixtures must define a repository") + raw_files = repository.get("files") + raw_failures = repository.get("transient_failures") + raw_tests = repository.get("tests") + if not isinstance(raw_files, list) or not raw_files: + raise ToolError("coding_agent repository files must be a non-empty array") + if not isinstance(raw_failures, Mapping): + raise ToolError("coding_agent transient_failures must be an object") + if not isinstance(raw_tests, list) or not raw_tests: + raise ToolError("coding_agent repository tests must be a non-empty array") + + files: dict[str, str] = {} + for item in raw_files: + if not isinstance(item, Mapping): + raise ToolError("coding_agent repository files must be objects") + path = item.get("path") + content = item.get("content") + if not isinstance(path, str) or not path or not isinstance(content, str): + raise ToolError("coding_agent repository files require path and content strings") + if path in files: + raise ToolError(f"duplicate coding_agent repository path {path!r}") + files[path] = content + + failures: dict[str, int] = {} + for operation, count in raw_failures.items(): + if not isinstance(operation, str) or not isinstance(count, int) or count < 0: + raise ToolError("coding_agent transient failures require non-negative counts") + failures[operation] = count + + tests: dict[str, tuple[str, str]] = {} + for item in raw_tests: + if not isinstance(item, Mapping): + raise ToolError("coding_agent repository tests must be objects") + name = item.get("name") + path = item.get("path") + contains = item.get("contains") + if not all(isinstance(value, str) and value for value in (name, path, contains)): + raise ToolError("coding_agent repository tests require name, path, and contains") + if cast(str, path) not in files: + raise ToolError(f"coding_agent test path {path!r} does not exist") + tests[cast(str, name)] = (cast(str, path), cast(str, contains)) + + search_page_size = repository.get("search_page_size", 2) + read_chunk_lines = repository.get("read_chunk_lines", 2) + if not isinstance(search_page_size, int) or search_page_size < 1: + raise ToolError("coding_agent search_page_size must be a positive integer") + if not isinstance(read_chunk_lines, int) or read_chunk_lines < 1: + raise ToolError("coding_agent read_chunk_lines must be a positive integer") + return cls(files, search_page_size, read_chunk_lines, failures, tests) + + def consume_failure(self, operation: str) -> None: + remaining = self.transient_failures.get(operation, 0) + if remaining: + self.transient_failures[operation] = remaining - 1 + raise ToolError(f"transient repository failure during {operation}; retry the call") + + +ToolHandler: TypeAlias = Callable[ + [Mapping[str, Any], Mapping[str, Any], CodingRepository | None], ToolResult +] + + +_MISSING = object() + + +@dataclass(frozen=True) +class ToolPatchOperation: + operation: str + path: str + value: Any = _MISSING + + def __post_init__(self) -> None: + if self.operation not in {"add", "replace", "remove"}: + raise ToolError(f"unknown tool overlay operation {self.operation!r}") + _json_pointer_tokens(self.path) + if self.operation == "remove": + if self.value is not _MISSING: + raise ToolError("remove tool overlay operations may not define a value") + elif self.value is _MISSING: + raise ToolError(f"{self.operation} tool overlay operations require a value") + else: + object.__setattr__(self, "value", _freeze_json(self.value)) + + +@dataclass(frozen=True) +class ToolResultOverlay: + tool_name: str + match_arguments: Mapping[str, Any] + operations: tuple[ToolPatchOperation, ...] + + def __post_init__(self) -> None: + if not self.tool_name: + raise ToolError("tool overlay names must be non-empty") + if not isinstance(self.match_arguments, Mapping): + raise ToolError("tool overlay match_arguments must be an object") + if not self.operations: + raise ToolError("tool overlays must define at least one operation") + frozen_arguments = _freeze_json(self.match_arguments) + object.__setattr__(self, "match_arguments", cast(Mapping[str, Any], frozen_arguments)) + object.__setattr__(self, "operations", tuple(self.operations)) + + +@dataclass(frozen=True) +class ToolSpec: + name: str + description: str + parameters: Mapping[str, Any] + handler: ToolHandler + coding_only: bool = False + + def model_schema(self) -> dict[str, JSON]: + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": _json_copy(dict(self.parameters)), + }, + } + + +class ToolRegistry: + def __init__(self, specs: Sequence[ToolSpec]) -> None: + by_name = {spec.name: spec for spec in specs} + if len(by_name) != len(specs): + raise ToolError("tool names must be unique") + self._specs = MappingProxyType(by_name) + + def model_schemas(self, *, include_coding: bool) -> tuple[dict[str, JSON], ...]: + return tuple( + spec.model_schema() + for spec in self._specs.values() + if include_coding or not spec.coding_only + ) + + def invoke( + self, + name: str, + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + result_overlays: Sequence[ToolResultOverlay] = (), + repository: CodingRepository | None = None, + ) -> ToolResult: + try: + spec = self._specs[name] + except KeyError as error: + raise ToolError(f"unknown tool {name!r}") from error + validated = _validate_arguments(spec, arguments) + result = spec.handler(validated, fixture_set, repository) + return _apply_result_overlays(name, validated, result, result_overlays) + + +@dataclass(frozen=True) +class LocalTools: + fixture_set: Mapping[str, Any] + registry: ToolRegistry + result_overlays: tuple[ToolResultOverlay, ...] = () + repository: CodingRepository | None = None + + @property + def schemas(self) -> tuple[dict[str, JSON], ...]: + return self.registry.model_schemas(include_coding=self.repository is not None) + + def invoke(self, name: str, arguments: Mapping[str, Any]) -> ToolResult: + return self.registry.invoke( + name, + arguments, + self.fixture_set, + self.result_overlays, + self.repository, + ) + + +def load_fixture_sets(path: Path | None = None) -> Mapping[str, Mapping[str, Any]]: + source = path or Path(__file__).with_name("tool_fixtures.json") + try: + value = json.loads(source.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ToolError(f"unable to load tool fixtures from {source}: {error}") from error + if not isinstance(value, dict) or not value: + raise ToolError(f"tool fixtures in {source} must be a non-empty object") + fixture_sets: dict[str, Mapping[str, Any]] = {} + for name, fixture_set in value.items(): + if not isinstance(name, str) or not isinstance(fixture_set, dict): + raise ToolError(f"invalid tool fixture set in {source}") + if fixture_set.get("name") != name: + raise ToolError(f"tool fixture set {name!r} must repeat its name") + for field in ("documents", "records", "statuses"): + if not isinstance(fixture_set.get(field), list): + raise ToolError(f"tool fixture set {name!r} must define {field}") + fixture_sets[name] = MappingProxyType(fixture_set) + return MappingProxyType(fixture_sets) + + +def local_tools( + name: str, + *, + fixture_set: Mapping[str, Any] | None = None, + result_overlays: Sequence[ToolResultOverlay] = (), +) -> LocalTools: + if fixture_set is None: + try: + fixture_set = load_fixture_sets()[name] + except KeyError as error: + raise ToolError(f"unknown tool fixture set {name!r}") from error + if fixture_set.get("name") != name: + raise ToolError(f"tool fixture set must be named {name!r}") + overlays = tuple(result_overlays) + validate_result_overlays(fixture_set, overlays) + repository = CodingRepository.from_fixture_set(fixture_set) + return LocalTools(fixture_set, DEFAULT_REGISTRY, overlays, repository) + + +def validate_result_overlays( + fixture_set: Mapping[str, Any], + overlays: Sequence[ToolResultOverlay], +) -> None: + occupied_paths: list[tuple[str, Mapping[str, Any], str]] = [] + for overlay in overlays: + try: + spec = DEFAULT_REGISTRY._specs[overlay.tool_name] + except KeyError as error: + raise ToolError(f"unknown tool {overlay.tool_name!r} in result overlay") from error + matching_arguments = _validation_arguments(spec, overlay.match_arguments, fixture_set) + for operation in overlay.operations: + for tool_name, match_arguments, path in occupied_paths: + if ( + tool_name == overlay.tool_name + and path == operation.path + and _argument_matches_overlap(match_arguments, overlay.match_arguments) + ): + raise ToolError( + f"tool overlays collide at {overlay.tool_name!r} {operation.path!r}" + ) + occupied_paths.append((overlay.tool_name, overlay.match_arguments, operation.path)) + for arguments in matching_arguments: + repository = CodingRepository.from_fixture_set(fixture_set) + if repository is not None: + repository.transient_failures.clear() + result = spec.handler(arguments, fixture_set, repository) + for operation in overlay.operations: + _apply_json_pointer_operation(result, operation) + + +def build_registry() -> ToolRegistry: + return ToolRegistry( + ( + ToolSpec( + name="document_search", + description="Search local reference documents for relevant passages.", + parameters=_object_schema( + { + "query": {"type": "string", "minLength": 1}, + "limit": {"type": "integer", "minimum": 1, "maximum": 5}, + }, + required=("query",), + ), + handler=_document_search, + ), + ToolSpec( + name="record_lookup", + description="Look up a local record by identifier.", + parameters=_object_schema( + {"record_id": {"type": "string", "minLength": 1}}, + required=("record_id",), + ), + handler=_record_lookup, + ), + ToolSpec( + name="safe_arithmetic", + description="Calculate an expression using basic arithmetic.", + parameters=_object_schema( + {"expression": {"type": "string", "minLength": 1, "maxLength": 128}}, + required=("expression",), + ), + handler=_safe_arithmetic, + ), + ToolSpec( + name="status_lookup", + description="Look up the current status of a local item.", + parameters=_object_schema( + {"status_id": {"type": "string", "minLength": 1}}, + required=("status_id",), + ), + handler=_status_lookup, + ), + ToolSpec( + name="ticket_creation", + description="Create a deterministic local ticket.", + parameters=_object_schema( + { + "title": {"type": "string", "minLength": 1, "maxLength": 120}, + "description": {"type": "string", "minLength": 1, "maxLength": 1000}, + "priority": {"type": "string", "enum": ["low", "medium", "high"]}, + }, + required=("title", "description", "priority"), + ), + handler=_ticket_creation, + ), + ToolSpec( + name="repository_search", + description="Search repository paths and contents, continuing from the last page.", + parameters=_object_schema( + {"query": {"type": "string", "minLength": 1}}, + required=("query",), + ), + handler=_repository_search, + coding_only=True, + ), + ToolSpec( + name="read_file", + description="Read the next chunk of a repository file.", + parameters=_object_schema( + {"path": {"type": "string", "minLength": 1}}, + required=("path",), + ), + handler=_read_file, + coding_only=True, + ), + ToolSpec( + name="edit_file", + description="Replace one exact string in a repository file.", + parameters=_object_schema( + { + "path": {"type": "string", "minLength": 1}, + "old": {"type": "string", "minLength": 1}, + "new": {"type": "string", "minLength": 1}, + }, + required=("path", "old", "new"), + ), + handler=_edit_file, + coding_only=True, + ), + ToolSpec( + name="run_tests", + description="Run one focused repository test against the current files.", + parameters=_object_schema( + {"test": {"type": "string", "minLength": 1}}, + required=("test",), + ), + handler=_run_tests, + coding_only=True, + ), + ) + ) + + +def _validate_arguments(spec: ToolSpec, arguments: Mapping[str, Any]) -> dict[str, Any]: + result = _validate_argument_subset(spec, arguments) + required = set(spec.parameters["required"]) + missing = required - set(arguments) + if missing: + raise ToolError(f"{spec.name} is missing arguments: {sorted(missing)}") + return result + + +def _validate_argument_subset(spec: ToolSpec, arguments: Mapping[str, Any]) -> dict[str, Any]: + if not isinstance(arguments, Mapping): + raise ToolError(f"{spec.name} arguments must be an object") + properties = spec.parameters["properties"] + unknown = set(arguments) - set(properties) + if unknown: + raise ToolError(f"{spec.name} has unknown arguments: {sorted(unknown)}") + result = dict(arguments) + for name, value in result.items(): + _validate_value(spec.name, name, value, properties[name]) + return result + + +def _validation_arguments( + spec: ToolSpec, + selector: Mapping[str, Any], + fixture_set: Mapping[str, Any], +) -> tuple[dict[str, Any], ...]: + validated_selector = _validate_argument_subset(spec, selector) + if spec.name == "document_search": + candidates = [{"query": "local guidance"}] + elif spec.name == "record_lookup": + candidates = [{"record_id": str(item["id"])} for item in fixture_set["records"]] + [ + {"record_id": "__missing_record__"} + ] + elif spec.name == "status_lookup": + candidates = [{"status_id": str(item["id"])} for item in fixture_set["statuses"]] + [ + {"status_id": "__missing_status__"} + ] + elif spec.name == "safe_arithmetic": + candidates = [{"expression": "0"}] + elif spec.name == "repository_search": + candidates = [{"query": "repository"}] + elif spec.name == "read_file": + candidates = [{"path": str(item["path"])} for item in fixture_set["repository"]["files"]] + elif spec.name == "edit_file": + candidates = [ + { + "path": str(item["path"]), + "old": str(item["content"]), + "new": str(item["content"]), + } + for item in fixture_set["repository"]["files"] + ] + elif spec.name == "run_tests": + candidates = [{"test": str(item["name"])} for item in fixture_set["repository"]["tests"]] + else: + candidates = [ + { + "title": "Local request", + "description": "Validate the authored result shape.", + "priority": "low", + } + ] + results: list[dict[str, Any]] = [] + seen: set[str] = set() + for candidate in candidates: + arguments = _validate_arguments(spec, {**candidate, **validated_selector}) + key = json.dumps(arguments, sort_keys=True, separators=(",", ":")) + if key not in seen: + results.append(arguments) + seen.add(key) + return tuple(results) + + +def _object_schema(properties: Mapping[str, Any], *, required: Sequence[str]) -> Mapping[str, Any]: + return MappingProxyType( + { + "type": "object", + "properties": dict(properties), + "required": list(required), + "additionalProperties": False, + } + ) + + +def _validate_value(tool: str, name: str, value: Any, schema: Mapping[str, Any]) -> None: + expected = schema["type"] + valid = { + "string": lambda item: isinstance(item, str), + "integer": lambda item: isinstance(item, int) and not isinstance(item, bool), + }[expected](value) + if not valid: + raise ToolError(f"{tool}.{name} must be a {expected}") + if isinstance(value, str): + if len(value) < schema.get("minLength", 0): + raise ToolError(f"{tool}.{name} is too short") + if len(value) > schema.get("maxLength", math.inf): + raise ToolError(f"{tool}.{name} is too long") + if "enum" in schema and value not in schema["enum"]: + raise ToolError(f"{tool}.{name} must be one of {schema['enum']}") + if isinstance(value, int) and not isinstance(value, bool): + if value < schema.get("minimum", -math.inf) or value > schema.get("maximum", math.inf): + raise ToolError(f"{tool}.{name} is outside its allowed range") + + +def _document_search( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del repository + query_terms = set(_WORD.findall(str(arguments["query"]).lower())) + documents = fixture_set["documents"] + ranked = sorted( + documents, + key=lambda document: ( + -len(query_terms & set(_WORD.findall(str(document["text"]).lower()))), + str(document["id"]), + ), + ) + return {"documents": [_json_copy(document) for document in ranked[: arguments.get("limit", 3)]]} + + +def _record_lookup( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del repository + record = next( + (value for value in fixture_set["records"] if value["id"] == str(arguments["record_id"])), + None, + ) + return { + "found": record is not None, + "record": _json_copy(record) if record is not None else None, + } + + +def _status_lookup( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del repository + status = next( + (value for value in fixture_set["statuses"] if value["id"] == str(arguments["status_id"])), + None, + ) + return { + "found": status is not None, + "status": _json_copy(status) if status is not None else None, + } + + +def _safe_arithmetic( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del fixture_set, repository + expression = str(arguments["expression"]) + try: + result = _evaluate_arithmetic(ast.parse(expression, mode="eval").body) + except (SyntaxError, ArithmeticError, ValueError) as error: + raise ToolError(f"invalid arithmetic expression: {error}") from error + if not math.isfinite(float(result)) or abs(result) > 1_000_000_000_000: + raise ToolError("arithmetic result is outside the allowed range") + return {"expression": expression, "result": result} + + +def _ticket_creation( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del fixture_set, repository + encoded = json.dumps(arguments, sort_keys=True, separators=(",", ":")).encode() + return { + "ticket_id": f"TKT-{sha256(encoded).hexdigest()[:12].upper()}", + "state": "created", + "priority": str(arguments["priority"]), + } + + +def _repository_search( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del fixture_set + state = _require_repository(repository) + query = str(arguments["query"]) + query_terms = set(_WORD.findall(query.lower())) + matches = [] + for path, content in sorted(state.files.items()): + for line_number, line in enumerate(content.splitlines(), start=1): + searchable = set(_WORD.findall(f"{path} {line}".lower())) + if query_terms & searchable: + matches.append({"path": path, "line": line_number, "text": line}) + cursor_key = " ".join(sorted(query_terms)) + start = state.search_cursors.get(cursor_key, 0) + end = min(start + state.search_page_size, len(matches)) + state.search_cursors[cursor_key] = end + has_more = end < len(matches) + return { + "matches": cast(JSON, matches[start:end]), + "cursor": end, + "has_more": has_more, + } + + +def _read_file( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del fixture_set + state = _require_repository(repository) + path = str(arguments["path"]) + try: + content = state.files[path] + except KeyError as error: + raise ToolError(f"repository path {path!r} does not exist") from error + state.consume_failure(f"read_file:{path}") + lines = content.splitlines() + start = state.read_cursors.get(path, 0) + end = min(start + state.read_chunk_lines, len(lines)) + state.read_cursors[path] = end + return { + "path": path, + "start_line": start + 1, + "end_line": end, + "content": "\n".join(lines[start:end]), + "cursor": end, + "has_more": end < len(lines), + } + + +def _edit_file( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del fixture_set + state = _require_repository(repository) + path = str(arguments["path"]) + old = str(arguments["old"]) + new = str(arguments["new"]) + try: + content = state.files[path] + except KeyError as error: + raise ToolError(f"repository path {path!r} does not exist") from error + occurrences = content.count(old) + if occurrences != 1: + raise ToolError( + f"edit_file expected one occurrence of {old!r} in {path!r}, found {occurrences}" + ) + state.files[path] = content.replace(old, new, 1) + state.read_cursors[path] = 0 + return {"path": path, "changed": True, "replacements": 1} + + +def _run_tests( + arguments: Mapping[str, Any], + fixture_set: Mapping[str, Any], + repository: CodingRepository | None, +) -> ToolResult: + del fixture_set + state = _require_repository(repository) + test = str(arguments["test"]) + try: + path, expected = state.tests[test] + except KeyError as error: + raise ToolError(f"unknown repository test {test!r}") from error + run = state.test_runs.get(test, 0) + 1 + state.test_runs[test] = run + passed = expected in state.files[path] + return { + "test": test, + "run": run, + "passed": passed, + "summary": "1 passed" if passed else "1 failed", + "failure": None if passed else f"{path} does not contain {expected!r}", + } + + +def _require_repository(repository: CodingRepository | None) -> CodingRepository: + if repository is None: + raise ToolError("repository tools require the coding_agent fixture set") + return repository + + +_BINARY_OPERATORS: Mapping[type[ast.operator], Callable[[float, float], float]] = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, + ast.Mod: operator.mod, +} +_UNARY_OPERATORS: Mapping[type[ast.unaryop], Callable[[float], float]] = { + ast.UAdd: operator.pos, + ast.USub: operator.neg, +} + + +def _evaluate_arithmetic(node: ast.expr) -> int | float: + if ( + isinstance(node, ast.Constant) + and isinstance(node.value, (int, float)) + and not isinstance(node.value, bool) + ): + return node.value + if isinstance(node, ast.BinOp) and type(node.op) in _BINARY_OPERATORS: + return _BINARY_OPERATORS[type(node.op)]( + _evaluate_arithmetic(node.left), + _evaluate_arithmetic(node.right), + ) + if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPERATORS: + return _UNARY_OPERATORS[type(node.op)](_evaluate_arithmetic(node.operand)) + raise ValueError("only numeric literals and +, -, *, /, //, % are allowed") + + +def _argument_matches_overlap(left: Mapping[str, Any], right: Mapping[str, Any]) -> bool: + return all(left[key] == right[key] for key in left.keys() & right.keys()) + + +def _apply_result_overlays( + tool_name: str, + arguments: Mapping[str, Any], + result: ToolResult, + overlays: Sequence[ToolResultOverlay], +) -> ToolResult: + patched = cast(ToolResult, _json_copy(result)) + for overlay in overlays: + if overlay.tool_name != tool_name or not all( + arguments.get(name) == value for name, value in overlay.match_arguments.items() + ): + continue + for operation in overlay.operations: + _apply_json_pointer_operation(patched, operation) + return patched + + +def _apply_json_pointer_operation( + result: ToolResult, + operation: ToolPatchOperation, +) -> None: + tokens = _json_pointer_tokens(operation.path) + parent: Any = result + for token in tokens[:-1]: + if isinstance(parent, dict) and token in parent: + parent = parent[token] + elif isinstance(parent, list): + parent = parent[_list_index(token, len(parent), allow_end=False)] + else: + raise ToolError(f"tool overlay path {operation.path!r} does not exist") + token = tokens[-1] + if isinstance(parent, dict): + _patch_mapping(parent, token, operation) + elif isinstance(parent, list): + _patch_sequence(parent, token, operation) + else: + raise ToolError(f"tool overlay path {operation.path!r} has a scalar parent") + + +def _patch_mapping( + parent: dict[str, JSON], + token: str, + operation: ToolPatchOperation, +) -> None: + if operation.operation == "add": + parent[token] = _json_copy(operation.value) + return + if token not in parent: + raise ToolError(f"tool overlay path component {token!r} does not exist") + if operation.operation == "remove": + del parent[token] + else: + parent[token] = _json_copy(operation.value) + + +def _patch_sequence( + parent: list[JSON], + token: str, + operation: ToolPatchOperation, +) -> None: + if operation.operation == "add": + if token == "-": + parent.append(_json_copy(operation.value)) + else: + parent.insert( + _list_index(token, len(parent), allow_end=True), + _json_copy(operation.value), + ) + return + index = _list_index(token, len(parent), allow_end=False) + if operation.operation == "remove": + del parent[index] + else: + parent[index] = _json_copy(operation.value) + + +def _json_pointer_tokens(path: str) -> list[str]: + if not isinstance(path, str) or not path.startswith("/") or path == "/": + raise ToolError("tool overlay paths must be non-root JSON Pointers") + tokens = [] + for raw in path[1:].split("/"): + if re.search(r"~(?![01])", raw): + raise ToolError(f"tool overlay path {path!r} has an invalid escape") + tokens.append(raw.replace("~1", "/").replace("~0", "~")) + return tokens + + +def _list_index(token: str, length: int, *, allow_end: bool) -> int: + if not token.isdigit() or (len(token) > 1 and token.startswith("0")): + raise ToolError(f"tool overlay list index {token!r} is invalid") + index = int(token) + limit = length if allow_end else length - 1 + if index > limit: + raise ToolError(f"tool overlay list index {index} is out of range") + return index + + +def _json_copy(value: Any) -> Any: + return json.loads(json.dumps(_plain_json(value))) + + +def _freeze_json(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze_json(item) for key, item in value.items()}) + if isinstance(value, (list, tuple)): + return tuple(_freeze_json(item) for item in value) + try: + json.dumps(value) + except (TypeError, ValueError) as error: + raise ToolError(f"tool overlay values must be JSON-compatible: {error}") from error + return value + + +def _plain_json(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _plain_json(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_plain_json(item) for item in value] + return value + + +DEFAULT_REGISTRY = build_registry() diff --git a/scripts/datagen/graph_multi_agent.py b/scripts/datagen/graph_multi_agent.py new file mode 100644 index 00000000000..711a8b11d77 --- /dev/null +++ b/scripts/datagen/graph_multi_agent.py @@ -0,0 +1,179 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "langchain-core==1.5.6", +# "openinference-instrumentation==0.1.57", +# "openinference-instrumentation-langchain==0.1.70", +# "opentelemetry-exporter-otlp-proto-common==1.44.0", +# "opentelemetry-sdk==1.44.0", +# "protobuf==7.35.1", +# ] +# /// +"""Record fixed multi-agent graph fixtures through LangChain callbacks.""" + +from __future__ import annotations + +import argparse +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +from langchain_core.runnables import RunnableLambda +from openinference.instrumentation import ( + OITracer, + TraceConfig, + get_attributes_from_context, + using_session, +) +from openinference.instrumentation.langchain import LangChainInstrumentor +from openinference.semconv.trace import OpenInferenceMimeTypeValues, OpenInferenceSpanKindValues +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +if TYPE_CHECKING or __package__: + from scripts.datagen.conditions import materialize_condition + from scripts.datagen.recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + trace_ids, + ) +else: + from conditions import materialize_condition + from recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + trace_ids, + ) + + +class OpenInferenceContextSpanProcessor(SpanProcessor): + """Apply the active session to spans started by LangChain callbacks.""" + + def on_start(self, span: Span, parent_context: Any = None) -> None: + span.set_attributes(dict(get_attributes_from_context())) + + def on_end(self, span: ReadableSpan) -> None: + pass + + def shutdown(self) -> None: + pass + + +class GraphMultiAgentRecorder: + def __init__(self, exporter: SpanCaptureExporter, tracer: OITracer) -> None: + self._exporter = exporter + self._tracer = tracer + + def record(self, fixture: RecorderFixture, traces_path: Path) -> tuple[str, ...]: + prompt = fixture.inputs.get("prompt") + documents = fixture.inputs.get("documents") + if not isinstance(prompt, str) or not isinstance(documents, list) or not documents: + raise ValueError(f"fixture {fixture.fragment_id!r} has invalid graph inputs") + checkpoint = self._exporter.checkpoint() + + def research(state: Mapping[str, Any]) -> dict[str, Any]: + evidence = " ".join( + str(document.get("text", "")) + for document in documents + if isinstance(document, dict) + ) + return {**state, "evidence": evidence} + + def write(state: Mapping[str, Any]) -> dict[str, Any]: + if "ack-timeout" in fixture.fragment_id: + raise RuntimeError("writer could not validate the acknowledgement boundary") + return {**state, "answer": f"{state['prompt']}: {state['evidence']}"} + + researcher = RunnableLambda(research).with_config({"run_name": "research_agent"}) + writer = RunnableLambda(write).with_config({"run_name": "writer_agent"}) + + def supervise(state: Mapping[str, Any]) -> dict[str, Any]: + return cast(dict[str, Any], writer.invoke(researcher.invoke(state))) + + graph = RunnableLambda(supervise).with_config({"run_name": "supervisor_agent"}) + try: + with using_session(fixture.fragment_id): + with self._tracer.start_as_current_span( + "coordinate_research_request", + openinference_span_kind=OpenInferenceSpanKindValues.AGENT, + ) as root_span: + root_span.set_input(prompt, mime_type=OpenInferenceMimeTypeValues.TEXT.value) + try: + result = graph.invoke({"prompt": prompt}) + except RuntimeError as error: + if "ack-timeout" not in fixture.fragment_id: + raise + output = f"Unable to complete the request: {error}" + else: + output = str(result["answer"]) + root_span.set_output( + output, + mime_type=OpenInferenceMimeTypeValues.TEXT.value, + ) + finally: + spans = self._exporter.spans_since(checkpoint) + if spans: + append_spans(traces_path, spans) + return trace_ids(spans) + + +def record( + output_dir: Path, + *, + fixtures: Sequence[RecorderFixture] | None = None, + condition: str | None = None, + append: bool = False, +) -> tuple[dict[str, Any], ...]: + """Record every selected graph fixture into a corpus directory.""" + if condition is not None and fixtures is not None: + raise ValueError("condition and fixtures cannot be selected together") + if condition is not None: + conditioned = materialize_condition(condition) + if conditioned.fixture.archetype != "graph_multi_agent": + raise ValueError(f"condition {condition!r} does not select a graph fixture") + selected_fixtures: Sequence[RecorderFixture] = (conditioned.fixture,) + else: + selected_fixtures = fixtures_for("graph_multi_agent", fixtures=fixtures) + prepare_recording(output_dir, append=append) + exporter = SpanCaptureExporter() + provider = TracerProvider( + resource=Resource.create({"service.name": "datagen.graph_multi_agent"}) + ) + provider.add_span_processor(OpenInferenceContextSpanProcessor()) + provider.add_span_processor(SimpleSpanProcessor(cast(Any, exporter))) + tracer = OITracer(provider.get_tracer(__name__), TraceConfig()) + instrumentor = LangChainInstrumentor() + instrumentor.instrument(tracer_provider=provider) + fragments = [] + try: + recorder = GraphMultiAgentRecorder(exporter, tracer) + for fixture in selected_fixtures: + fragments.append(record_fixture(fixture, output_dir, recorder.record)) + finally: + instrumentor.uninstrument() + provider.shutdown() + return tuple(fragments) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--condition") + parser.add_argument("--append", action="store_true") + args = parser.parse_args() + fragments = record(args.output_dir, condition=args.condition, append=args.append) + print(f"Recorded {len(fragments)} graph fragments in {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/datagen/guardrailed_app.py b/scripts/datagen/guardrailed_app.py new file mode 100644 index 00000000000..cf687c815f7 --- /dev/null +++ b/scripts/datagen/guardrailed_app.py @@ -0,0 +1,150 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "guardrails-ai==0.5.0", +# "openinference-instrumentation==0.1.57", +# "openinference-instrumentation-guardrails==0.1.16", +# "opentelemetry-exporter-otlp-proto-common==1.44.0", +# "opentelemetry-sdk==1.44.0", +# "protobuf==7.35.1", +# ] +# /// +"""Record fixed policy outcomes through the Guardrails instrumentor.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +if TYPE_CHECKING or __package__: + from scripts.datagen.conditions import materialize_condition + from scripts.datagen.recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + trace_ids, + validate_recording, + ) +else: + from conditions import materialize_condition + from recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + trace_ids, + validate_recording, + ) + + +def record( + output_dir: Path, + *, + fixtures: Sequence[RecorderFixture] | None = None, + condition: str | None = None, + append: bool = False, +) -> tuple[dict[str, Any], ...]: + """Record every selected guardrail fixture into a corpus directory.""" + from guardrails import Guard # type: ignore[import-not-found] + from guardrails.validators import ( # type: ignore[import-not-found] + FailResult, + PassResult, + Validator, + register_validator, + ) + from openinference.instrumentation import using_session + from openinference.instrumentation.guardrails import ( # type: ignore[import-not-found] + GuardrailsInstrumentor, + ) + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + + @register_validator(name="datagen/fixed-policy", data_type="string") + class FixturePolicy(Validator): # type: ignore[misc] + def validate(self, value: Any, metadata: dict[str, Any]) -> Any: + del value + outcome = metadata.get("outcome") + if outcome == "allowed": + return PassResult() + if outcome == "redacted": + return FailResult( + error_message="sensitive detail removed", + fix_value="[redacted by policy]", + ) + return FailResult(error_message="request blocked by policy") + + if condition is not None and fixtures is not None: + raise ValueError("condition and fixtures cannot be selected together") + if condition is not None: + conditioned = materialize_condition(condition) + if conditioned.fixture.archetype != "guardrailed": + raise ValueError(f"condition {condition!r} does not select a guardrail fixture") + selected_fixtures: Sequence[RecorderFixture] = (conditioned.fixture,) + else: + selected_fixtures = fixtures_for("guardrailed", fixtures=fixtures) + prepare_recording(output_dir, append=append) + exporter = SpanCaptureExporter() + provider = TracerProvider(resource=Resource.create({"service.name": "datagen.guardrailed"})) + provider.add_span_processor(SimpleSpanProcessor(cast(Any, exporter))) + instrumentor = GuardrailsInstrumentor() + instrumentor.instrument(tracer_provider=provider) + + def adapter(fixture: RecorderFixture, traces_path: Path) -> tuple[str, ...]: + text = fixture.inputs.get("text") + outcome = fixture.inputs.get("outcome") + if not isinstance(text, str) or outcome not in {"allowed", "blocked", "redacted"}: + raise ValueError(f"fixture {fixture.fragment_id!r} has invalid policy inputs") + on_fail = {"allowed": "noop", "blocked": "exception", "redacted": "fix"}[outcome] + checkpoint = exporter.checkpoint() + try: + with using_session(fixture.fragment_id): + try: + Guard().use(FixturePolicy(on_fail=on_fail)).validate( + text, + metadata={"outcome": outcome}, + ) + except Exception: + if outcome != "blocked": + raise + finally: + spans = exporter.spans_since(checkpoint) + if spans: + append_spans(traces_path, spans) + return trace_ids(spans) + + fragments = [] + try: + for fixture in selected_fixtures: + fragments.append(record_fixture(fixture, output_dir, adapter)) + finally: + instrumentor.uninstrument() + provider.shutdown() + validate_recording( + output_dir / "traces.jsonl", + required_span_kinds=("GUARDRAIL",), + recorder_name="Guardrails instrumenter", + ) + return tuple(fragments) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--condition") + parser.add_argument("--append", action="store_true") + args = parser.parse_args() + fragments = record(args.output_dir, condition=args.condition, append=args.append) + print(f"Recorded {len(fragments)} guardrail fragments in {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/datagen/llama_index_rag.py b/scripts/datagen/llama_index_rag.py new file mode 100644 index 00000000000..f65cc917efc --- /dev/null +++ b/scripts/datagen/llama_index_rag.py @@ -0,0 +1,170 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "llama-index-core==0.14.23", +# "llama-index-llms-openai==0.7.10", +# "llama-index-postprocessor-cohere-rerank==0.9.0", +# "openinference-instrumentation==0.1.57", +# "openinference-instrumentation-llama-index==4.4.5", +# "opentelemetry-exporter-otlp-proto-common==1.44.0", +# "opentelemetry-sdk==1.44.0", +# "protobuf==7.35.1", +# ] +# /// +"""Record fixed RAG fixtures through the LlamaIndex instrumentor.""" + +from __future__ import annotations + +import argparse +import os +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, cast + +from openinference.instrumentation import using_session +from openinference.instrumentation.llama_index import ( # type: ignore[import-not-found] + LlamaIndexInstrumentor, +) +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +if TYPE_CHECKING or __package__: + from scripts.datagen.conditions import materialize_condition + from scripts.datagen.rag import build_rag_engine + from scripts.datagen.recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + validate_recording, + ) +else: + from conditions import materialize_condition + from rag import build_rag_engine + from recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + validate_recording, + ) + +Provider = Literal["scripted", "live"] + + +def record( + output_dir: Path, + *, + fixtures: Sequence[RecorderFixture] | None = None, + condition: str | None = None, + append: bool = False, + provider: Provider = "scripted", + model: str | None = None, + live_llm: Any = None, +) -> tuple[dict[str, Any], ...]: + """Record every selected RAG fixture into a corpus directory.""" + model = resolve_live_model(model) + if provider not in ("scripted", "live"): + raise ValueError(f"unknown RAG provider {provider!r}") + if condition is not None and fixtures is not None: + raise ValueError("condition and fixtures cannot be selected together") + if provider == "live" and not model: + raise ValueError("live RAG recording requires an explicit model") + if provider == "live" and live_llm is None: + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + raise ValueError("live RAG recording requires OPENAI_API_KEY") + from llama_index.llms.openai import OpenAI # type: ignore[import-not-found] + + llm_args: dict[str, Any] = {"model": model, "api_key": api_key} + if base_url := os.environ.get("OPENAI_BASE_URL"): + llm_args["api_base"] = base_url + live_llm = OpenAI(**llm_args) + if condition is not None: + conditioned = materialize_condition(condition) + if conditioned.fixture.archetype != "rag": + raise ValueError(f"condition {condition!r} does not select a RAG fixture") + selected_fixtures: Sequence[RecorderFixture] = (conditioned.fixture,) + else: + selected_fixtures = fixtures_for("rag", fixtures=fixtures) + prepare_recording(output_dir, append=append) + exporter = SpanCaptureExporter() + tracer_provider = TracerProvider(resource=Resource.create({"service.name": "datagen.rag"})) + tracer_provider.add_span_processor(SimpleSpanProcessor(cast(Any, exporter))) + instrumentor = LlamaIndexInstrumentor() + instrumentor.instrument(tracer_provider=tracer_provider) + + def adapter(fixture: RecorderFixture, traces_path: Path) -> tuple[str, ...]: + questions = fixture.inputs.get("questions") + documents = fixture.inputs.get("documents") + if ( + not isinstance(questions, list) + or not all(isinstance(question, str) for question in questions) + or not isinstance(documents, list) + or not all(isinstance(document, Mapping) for document in documents) + ): + raise ValueError(f"fixture {fixture.fragment_id!r} has invalid RAG inputs") + engine = build_rag_engine( + cast(Sequence[Mapping[str, Any]], documents), + llm=live_llm if provider == "live" else None, + ) + checkpoint = exporter.checkpoint() + try: + with using_session(fixture.fragment_id): + for question in questions: + engine.query(question) + except Exception: + if provider == "scripted": + raise + finally: + spans = exporter.spans_since(checkpoint) + if spans: + append_spans(traces_path, spans) + return trace_ids(spans) + + fragments = [] + try: + for fixture in selected_fixtures: + fragments.append(record_fixture(fixture, output_dir, adapter)) + finally: + instrumentor.uninstrument() + tracer_provider.shutdown() + if provider == "scripted": + validate_recording( + output_dir / "traces.jsonl", + required_span_kinds=("CHAIN", "EMBEDDING", "RETRIEVER", "RERANKER", "LLM"), + recorder_name="LlamaIndex instrumenter", + ) + return tuple(fragments) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--condition") + parser.add_argument("--append", action="store_true") + parser.add_argument("--provider", choices=("scripted", "live"), default="scripted") + parser.add_argument("--model") + args = parser.parse_args() + fragments = record( + args.output_dir, + condition=args.condition, + append=args.append, + provider=args.provider, + model=args.model, + ) + print(f"Recorded {len(fragments)} RAG fragments in {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/datagen/mock_openai_provider.py b/scripts/datagen/mock_openai_provider.py new file mode 100644 index 00000000000..0c46d15d803 --- /dev/null +++ b/scripts/datagen/mock_openai_provider.py @@ -0,0 +1,211 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "httpx==0.28.1", +# ] +# /// +"""Deterministic OpenAI-compatible responses for offline trace recording.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from hashlib import sha256 +from http import HTTPStatus +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING or __package__: + from scripts.datagen.recording import RecorderFixture +else: + from recording import RecorderFixture # type: ignore[import-not-found,no-redef] + + +class ScriptedProviderError(ValueError): + """Raised when a response script cannot serve an OpenAI request.""" + + +class ScriptedOpenAIProvider: + """Serve a fixed response sequence through an in-process HTTP transport.""" + + def __init__(self, responses: Sequence[Mapping[str, Any]]) -> None: + if not responses: + raise ScriptedProviderError("a provider script must contain at least one response") + self._responses = tuple(dict(response) for response in responses) + self._response_index = 0 + self.requests: list[dict[str, Any]] = [] + + @classmethod + def for_fixture(cls, fixture: RecorderFixture) -> ScriptedOpenAIProvider: + turns = fixture.inputs.get("turns") + if not isinstance(turns, list) or not turns: + raise ScriptedProviderError( + f"fixture {fixture.fragment_id!r} does not contain scripted turns" + ) + responses = [] + for turn in turns: + if not isinstance(turn, dict) or not isinstance(turn.get("assistant"), str): + raise ScriptedProviderError( + f"fixture {fixture.fragment_id!r} has an invalid scripted turn" + ) + responses.append({"content": turn["assistant"]}) + return cls(responses) + + @property + def response_index(self) -> int: + return self._response_index + + def http_client(self) -> Any: + import httpx + + return httpx.Client(transport=httpx.MockTransport(self._handle)) + + def _handle(self, request: Any) -> Any: + import httpx + + if request.url.path != "/v1/chat/completions": + return httpx.Response( + HTTPStatus.NOT_FOUND, + json={"error": {"message": "not found", "type": "invalid_request_error"}}, + request=request, + ) + try: + body = json.loads(request.content) + except (json.JSONDecodeError, TypeError): + return httpx.Response( + HTTPStatus.BAD_REQUEST, + json={"error": {"message": "invalid JSON", "type": "invalid_request_error"}}, + request=request, + ) + if not isinstance(body, dict): + return httpx.Response( + HTTPStatus.BAD_REQUEST, + json={"error": {"message": "request must be an object"}}, + request=request, + ) + self.requests.append(body) + if self._response_index >= len(self._responses): + raise ScriptedProviderError("provider received more requests than scripted responses") + response = self._responses[self._response_index] + self._response_index += 1 + + status = response.get("status", HTTPStatus.OK) + if not isinstance(status, int): + raise ScriptedProviderError("scripted response status must be an integer") + if status != HTTPStatus.OK: + error = response.get("error") + payload = ( + {"error": dict(error)} + if isinstance(error, Mapping) + else {"error": {"message": f"scripted HTTP {status}"}} + ) + return httpx.Response(status, json=payload, request=request) + + completion = _completion(body, response, self._response_index) + if body.get("stream"): + return httpx.Response( + HTTPStatus.OK, + headers={"content-type": "text/event-stream"}, + content=stream_chat_completion(completion), + request=request, + ) + return httpx.Response(HTTPStatus.OK, json=completion, request=request) + + +def _completion( + request: Mapping[str, Any], + response: Mapping[str, Any], + response_index: int, +) -> dict[str, Any]: + message: dict[str, Any] = {"role": "assistant", "content": None} + tool_call = response.get("tool_call") + if tool_call is not None: + if not isinstance(tool_call, Mapping): + raise ScriptedProviderError("tool_call must be an object") + name = tool_call.get("name") + arguments = tool_call.get("arguments") + if not isinstance(name, str) or not isinstance(arguments, Mapping): + raise ScriptedProviderError("tool_call requires a name and object arguments") + message["tool_calls"] = [ + { + "id": f"call-{response_index}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments, sort_keys=True, separators=(",", ":")), + }, + } + ] + finish_reason = "tool_calls" + else: + content = response.get("content") + if not isinstance(content, str): + raise ScriptedProviderError("scripted response requires content or tool_call") + message["content"] = content + finish_reason = "stop" + + prompt_tokens = _token_count(request.get("messages", [])) + completion_tokens = _token_count(message) + identifier = _stable_id({"request": request, "response_index": response_index}) + return { + "id": f"chatcmpl-{identifier[:24]}", + "object": "chat.completion", + "created": 0, + "model": request.get("model", "datagen-scripted"), + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": finish_reason, + } + ], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + "prompt_tokens_details": {"cached_tokens": 0}, + "completion_tokens_details": {"reasoning_tokens": 0}, + }, + } + + +def stream_chat_completion(completion: Mapping[str, Any]) -> bytes: + """Encode a text completion as an OpenAI server-sent-event stream.""" + choice = completion["choices"][0] + content = choice["message"].get("content") or "" + chunks = [ + { + "id": completion["id"], + "object": "chat.completion.chunk", + "created": completion["created"], + "model": completion["model"], + "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}], + }, + { + "id": completion["id"], + "object": "chat.completion.chunk", + "created": completion["created"], + "model": completion["model"], + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + { + "id": completion["id"], + "object": "chat.completion.chunk", + "created": completion["created"], + "model": completion["model"], + "choices": [], + "usage": completion["usage"], + }, + ] + events = [f"data: {json.dumps(chunk, separators=(',', ':'))}\n\n" for chunk in chunks] + return ("".join(events) + "data: [DONE]\n\n").encode() + + +def _token_count(value: Any) -> int: + text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) + return max(1, round(len(text.split()) * 1.35)) + + +def _stable_id(value: Any) -> str: + encoded = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode() + return sha256(encoded).hexdigest() diff --git a/scripts/datagen/openai_chat_sessions.py b/scripts/datagen/openai_chat_sessions.py new file mode 100644 index 00000000000..3a30332a54a --- /dev/null +++ b/scripts/datagen/openai_chat_sessions.py @@ -0,0 +1,376 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "httpx==0.28.1", +# "openai==3.2.0", +# "openinference-instrumentation==0.1.57", +# "openinference-instrumentation-openai==0.1.54", +# "opentelemetry-exporter-otlp-proto-common==1.44.0", +# "opentelemetry-sdk==1.44.0", +# "protobuf==7.35.1", +# ] +# /// +"""Record fixed plain-chat fixtures through the OpenAI instrumentor.""" + +from __future__ import annotations + +import argparse +import os +import random +from collections.abc import Mapping, Sequence +from dataclasses import replace +from math import log +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, cast + +from openai import OpenAI +from openinference.instrumentation import suppress_tracing, using_session +from openinference.instrumentation.openai import OpenAIInstrumentor +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +if TYPE_CHECKING or __package__: + from scripts.datagen.conditions import materialize_condition + from scripts.datagen.mock_openai_provider import ScriptedOpenAIProvider + from scripts.datagen.recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + ) +else: + from conditions import materialize_condition + from mock_openai_provider import ScriptedOpenAIProvider + from recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + ) + +Provider = Literal["scripted", "live"] + +# Live conversations run until the simulated user is satisfied. The target +# steers how much conversation happens before the wind-down instruction is +# added; the actual ending is the simulated user closing without a further +# request. The hard cap only guards against a conversation that never closes. +_TARGET_TURNS_MEDIAN = 8.0 +_TARGET_TURNS_SIGMA = 0.4 +_TARGET_TURNS_MAX = 25 + +_WIND_DOWN_SUFFIX = ( + " Your remaining concerns are nearly addressed. When the assistant has " + "answered your current question, close the conversation with a brief " + "message of thanks or acknowledgement that asks nothing further. If " + "something important is still unresolved, ask about it instead." +) + + +def _draw_target_turns(rng: random.Random) -> int: + target = round(rng.lognormvariate(log(_TARGET_TURNS_MEDIAN), _TARGET_TURNS_SIGMA)) + return min(_TARGET_TURNS_MAX, max(2, target)) + + +def _is_closing(message: str) -> bool: + """A wind-down-phase user message with no question or request ends the session.""" + return "?" not in message + + +_DISPOSITION_PROMPTS: Mapping[str, str] = { + "impatient": ( + "You are an imperfect human continuing the conversation. You are impatient and want a " + "useful answer quickly, so press for specifics and skip pleasantries. Reply with only the " + "next user message. Do not label the speaker or explain the simulation." + ), + "confused_novice": ( + "You are an imperfect human continuing the conversation. You are a confused novice who " + "may misuse terms, ask basic follow-ups, or need an earlier point clarified. Reply with " + "only the next user message. Do not label the speaker or explain the simulation." + ), + "terse_expert": ( + "You are an imperfect human continuing the conversation. You are a terse expert who uses " + "precise domain language, omits context you assume is obvious, and corrects inaccuracies " + "directly. Reply with only the next user message. Do not label the speaker or explain the " + "simulation." + ), + "chatty": ( + "You are an imperfect human continuing the conversation. You are chatty and volunteer " + "small contextual details, reactions, and side comments while pursuing the task. Reply " + "with only the next user message. Do not label the speaker or explain the simulation." + ), + "frustrated": ( + "You are an imperfect human continuing the conversation. You are frustrated by the " + "situation, show restrained annoyance, and challenge answers that do not resolve the " + "problem. Reply with only the next user message. Do not label the speaker or explain the " + "simulation." + ), + "distracted": ( + "You are an imperfect human continuing the conversation. You are distracted mid-task, so " + "you may lose the thread, revise a detail, or abruptly return to an earlier concern. Reply " + "with only the next user message. Do not label the speaker or explain the simulation." + ), +} + + +def _with_opening_variant(fixture: RecorderFixture, rng: random.Random) -> RecorderFixture: + """Pick one authored phrasing of the conversation's opening for this run.""" + variants = fixture.inputs.get("opening_variants") + turns = fixture.inputs.get("turns") + if not isinstance(variants, list) or not isinstance(turns, list) or not turns: + return fixture + first = turns[0] + if not isinstance(first, dict) or not isinstance(first.get("user"), str): + return fixture + phrasings = [first["user"], *(v for v in variants if isinstance(v, str))] + opening = rng.choice(phrasings) + new_turns = [{**first, "user": opening}, *turns[1:]] + return replace(fixture, inputs={**fixture.inputs, "turns": new_turns}) + + +def record( + output_dir: Path, + *, + fixtures: Sequence[RecorderFixture] | None = None, + condition: str | None = None, + append: bool = False, + provider: Provider = "scripted", + model: str | None = None, + live_client: OpenAI | None = None, + disposition: str | None = None, + target_turns: int | None = None, +) -> tuple[dict[str, Any], ...]: + """Record every selected plain-chat fixture into a corpus directory.""" + model = resolve_live_model(model) + if provider not in ("scripted", "live"): + raise ValueError(f"unknown plain-chat provider {provider!r}") + if condition is not None and fixtures is not None: + raise ValueError("condition and fixtures cannot be selected together") + if provider == "live" and not model: + raise ValueError("live plain-chat recording requires an explicit model") + if provider == "live" and live_client is None: + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + raise ValueError("live plain-chat recording requires OPENAI_API_KEY") + client_args: dict[str, Any] = {"api_key": api_key, "max_retries": 0} + if base_url := os.environ.get("OPENAI_BASE_URL"): + client_args["base_url"] = base_url + live_client = OpenAI(**client_args) + if condition is not None: + conditioned = materialize_condition(condition) + if conditioned.fixture.archetype != "plain_chat": + raise ValueError(f"condition {condition!r} does not select a plain-chat fixture") + selected_fixtures: Sequence[RecorderFixture] = (conditioned.fixture,) + else: + selected_fixtures = fixtures_for("plain_chat", fixtures=fixtures) + if disposition is not None and disposition not in _DISPOSITION_PROMPTS: + raise ValueError(f"unknown plain-chat disposition {disposition!r}") + disposition_prompts = ( + (_DISPOSITION_PROMPTS[disposition],) + if disposition is not None + else tuple(_DISPOSITION_PROMPTS.values()) + ) + length_rng = random.Random() + opening_rng = random.Random() + prepare_recording(output_dir, append=append) + exporter = SpanCaptureExporter() + tracer_provider = TracerProvider( + resource=Resource.create({"service.name": "datagen.plain_chat"}) + ) + tracer_provider.add_span_processor(SimpleSpanProcessor(cast(Any, exporter))) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=tracer_provider) + fragments = [] + try: + for fixture_index, fixture in enumerate(selected_fixtures): + if provider == "live" and condition is None: + fixture = _with_opening_variant(fixture, opening_rng) + fragments.append( + record_fixture( + fixture, + output_dir, + lambda selected, traces_path: _record_fixture( + selected, + traces_path, + exporter, + provider=provider, + model=model, + live_client=live_client, + disposition_prompt=disposition_prompts[ + fixture_index % len(disposition_prompts) + ], + target_turns=target_turns or _draw_target_turns(length_rng), + ), + ) + ) + finally: + instrumentor.uninstrument() + tracer_provider.shutdown() + return tuple(fragments) + + +def _record_fixture( + fixture: RecorderFixture, + traces_path: Path, + exporter: SpanCaptureExporter, + *, + provider: Provider, + model: str | None, + live_client: OpenAI | None, + disposition_prompt: str, + target_turns: int, +) -> tuple[str, ...]: + if provider == "scripted": + scripted = ScriptedOpenAIProvider.for_fixture(fixture) + client = OpenAI( + api_key="datagen-dummy-key", + base_url="https://datagen.test/v1", + http_client=cast(Any, scripted.http_client()), + max_retries=0, + ) + model_name = "datagen-scripted" + else: + client = cast(OpenAI, live_client) + model_name = cast(str, model) + turns = fixture.inputs.get("turns") + if not isinstance(turns, list) or not turns: + raise ValueError(f"fixture {fixture.fragment_id!r} has no chat turns") + messages: list[Mapping[str, Any]] = [] + checkpoint = exporter.checkpoint() + try: + with using_session(fixture.fragment_id): + if provider == "scripted": + _run_scripted_turns(fixture, client, model_name, turns, messages) + else: + _run_live_conversation( + fixture, + client, + model_name, + turns, + messages, + disposition_prompt=disposition_prompt, + target_turns=target_turns, + ) + except Exception: + if provider == "scripted": + raise + finally: + spans = exporter.spans_since(checkpoint) + if spans: + append_spans(traces_path, spans) + return trace_ids(spans) + + +def _run_scripted_turns( + fixture: RecorderFixture, + client: OpenAI, + model_name: str, + turns: Sequence[Any], + messages: list[Mapping[str, Any]], +) -> None: + for turn in turns: + if not isinstance(turn, dict): + raise ValueError(f"fixture {fixture.fragment_id!r} has an invalid turn") + user = turn.get("user") + expected = turn.get("assistant") + if not isinstance(user, str) or not isinstance(expected, str): + raise ValueError(f"fixture {fixture.fragment_id!r} has an invalid turn") + messages.append({"role": "user", "content": user}) + response = client.chat.completions.create(model=model_name, messages=cast(Any, messages)) + content = response.choices[0].message.content + if content != expected: + raise ValueError(f"fixture {fixture.fragment_id!r} returned unexpected content") + messages.append({"role": "assistant", "content": content}) + + +def _run_live_conversation( + fixture: RecorderFixture, + client: OpenAI, + model_name: str, + turns: Sequence[Any], + messages: list[Mapping[str, Any]], + *, + disposition_prompt: str, + target_turns: int, +) -> None: + opening = turns[0].get("user") if isinstance(turns[0], dict) else None + if not isinstance(opening, str): + raise ValueError(f"fixture {fixture.fragment_id!r} has an invalid opening turn") + hard_cap = max(target_turns * 2, target_turns + 3) + user: str | None = opening + for turn_index in range(hard_cap): + winding_down = turn_index + 1 >= target_turns + if turn_index > 0: + prompt = disposition_prompt + (_WIND_DOWN_SUFFIX if winding_down else "") + user = _simulate_user(client, model_name, messages, prompt) + if not isinstance(user, str): + break + closing = turn_index > 0 and winding_down and _is_closing(user) + messages.append({"role": "user", "content": user}) + response = client.chat.completions.create(model=model_name, messages=cast(Any, messages)) + content = response.choices[0].message.content + if not isinstance(content, str): + break + messages.append({"role": "assistant", "content": content}) + if closing: + break + + +def _simulate_user( + client: OpenAI, + model: str, + messages: Sequence[Mapping[str, Any]], + disposition_prompt: str, +) -> str | None: + with suppress_tracing(): + response = client.chat.completions.create( + model=model, + messages=cast( + Any, + [{"role": "system", "content": disposition_prompt}, *messages], + ), + ) + return response.choices[0].message.content + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--condition") + parser.add_argument("--append", action="store_true") + parser.add_argument("--provider", choices=("scripted", "live"), default="scripted") + parser.add_argument("--model") + parser.add_argument("--disposition", choices=tuple(_DISPOSITION_PROMPTS)) + parser.add_argument( + "--target-turns", + type=int, + help=( + "Steer live conversations toward this many user turns; the ending " + "still happens when the simulated user closes. Drawn per fixture " + "when omitted." + ), + ) + args = parser.parse_args() + fragments = record( + args.output_dir, + condition=args.condition, + append=args.append, + provider=args.provider, + model=args.model, + disposition=args.disposition, + target_turns=args.target_turns, + ) + print(f"Recorded {len(fragments)} plain-chat fragments in {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/datagen/organic_conditions.json b/scripts/datagen/organic_conditions.json new file mode 100644 index 00000000000..447bc075cc1 --- /dev/null +++ b/scripts/datagen/organic_conditions.json @@ -0,0 +1,201 @@ +[ + { + "condition_id": "support-stale-delivery-status", + "fixture_id": "support-order-and-status-tools", + "fragment_id": "support-order-and-status-tools-stale", + "intensity": 0.2, + "strengths": { + "subtle": { + "document_edits": [ + { + "target": "tool_corpus", + "document_id": "delivery-guide", + "operation": "replace_once", + "source": "after two business days without movement", + "replacement": "after three business days without movement" + } + ], + "tool_overlays": [ + { + "tool_name": "status_lookup", + "match_arguments": { + "status_id": "order-1001" + }, + "operations": [ + { + "operation": "replace", + "path": "/status/state", + "value": "delayed_review" + } + ] + } + ] + }, + "moderate": { + "document_edits": [ + { + "target": "tool_corpus", + "document_id": "delivery-guide", + "operation": "replace_once", + "source": "after two business days without movement", + "replacement": "after five business days without movement" + } + ], + "tool_overlays": [ + { + "tool_name": "status_lookup", + "match_arguments": { + "status_id": "order-1001" + }, + "operations": [ + { + "operation": "replace", + "path": "/status/state", + "value": "exception_review" + }, + { + "operation": "add", + "path": "/status/note", + "value": "The carrier history is still being reconciled." + } + ] + } + ] + }, + "strong": { + "document_edits": [ + { + "target": "tool_corpus", + "document_id": "delivery-guide", + "operation": "replace_once", + "source": "A carrier trace is available after two business days without movement.", + "replacement": "Carrier trace guidance is temporarily unavailable." + } + ], + "tool_overlays": [ + { + "tool_name": "status_lookup", + "match_arguments": { + "status_id": "order-1001" + }, + "operations": [ + { + "operation": "replace", + "path": "/status/state", + "value": "unknown" + }, + { + "operation": "remove", + "path": "/status/detail" + } + ] + } + ] + } + } + }, + { + "condition_id": "support-late-express-ambiguity", + "fixture_id": "support-late-express-order", + "fragment_id": "support-late-express-order-ambiguous", + "intensity": 0.45, + "strengths": { + "subtle": { + "input_replacements": [ + { + "path": "/turns/1/user", + "value": "Checkout promised Thursday, while the carrier now says Friday. Which estimate should support use?" + } + ] + }, + "moderate": { + "input_replacements": [ + { + "path": "/turns/2/user", + "value": "I still need the order for an event, but I also paid for express delivery. Can support protect the shipment and review only the fee?" + } + ] + }, + "strong": { + "input_replacements": [ + { + "path": "/turns/0/user", + "value": "The carrier says weather delayed my express order, but tracking shows the package never left the origin facility. What can support verify?" + }, + { + "path": "/turns/2/user", + "value": "Do not cancel it yet. I need support to separate a shipping-fee review from a replacement decision if the scan never moves." + } + ] + } + } + }, + { + "condition_id": "research-fleet-delivery-pressure", + "fixture_id": "research-electric-fleet-status", + "fragment_id": "research-electric-fleet-status-pressure", + "intensity": 0.55, + "strengths": { + "subtle": { + "input_replacements": [ + { + "path": "/questions/2", + "value": "What evidence would confirm that commissioning remains on schedule?" + } + ] + }, + "moderate": { + "input_replacements": [ + { + "path": "/questions/2", + "value": "How much confidence should decision-makers place in the target date given the different reporting dates and depot warning?" + } + ] + }, + "strong": { + "input_replacements": [ + { + "path": "/questions/0", + "value": "Reconcile the auditor and Fleet Operations counts before stating how many battery-electric buses are active or commissioning." + }, + { + "path": "/questions/2", + "value": "Could the depot constraint prevent the 2026 target even if every contracted bus arrives, and which source supports that conclusion?" + } + ] + } + } + }, + { + "condition_id": "analytics-corrected-refund-export", + "fixture_id": "analytics-refund-export-extraction", + "fragment_id": "analytics-refund-export-extraction-corrected", + "intensity": 0.35, + "strengths": { + "subtle": { + "input_replacements": [ + { + "path": "/text", + "value": "Export last month's refunds with order, refund, reason, amount, currency, and refund timestamp as UTF-8 CSV." + } + ] + }, + "moderate": { + "input_replacements": [ + { + "path": "/text", + "value": "Export last month's refunds. Correction: use one UTF-8 CSV row per refund, not one row per order, and include order ID, refund ID, reason, amount, currency, and refund timestamp." + } + ] + }, + "strong": { + "input_replacements": [ + { + "path": "/text", + "value": "Prepare last month's refund reconciliation export. The first request omitted identifiers; include order ID and refund ID, preserve separate rows for partial refunds, and output reason, amount, currency, and refund timestamp as UTF-8 CSV." + } + ] + } + } + } +] diff --git a/scripts/datagen/publish.py b/scripts/datagen/publish.py new file mode 100644 index 00000000000..d86ab7aefa0 --- /dev/null +++ b/scripts/datagen/publish.py @@ -0,0 +1,185 @@ +"""Prepare and validate the datagen corpus for manual publication.""" + +from __future__ import annotations + +import argparse +import json +import shlex +import shutil +import sys +import tempfile +from dataclasses import asdict, dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Any, Mapping, Sequence, TextIO + +from phoenix.experimental.datagen.fetcher import CorpusFetchError, fetch_corpus +from phoenix.experimental.datagen.loader import CorpusError, load_corpus + +_ARCHIVE_NAME = "corpus.tar.gz" +_BUCKET = "arize-phoenix-assets" +_PREFIX = "datagen" +_DEFAULT_OUTPUT_DIR = Path("dist/datagen-publication") + + +@dataclass(frozen=True) +class ValidatedCorpus: + archive: Path + sha256: str + size_bytes: int + fragment_count: int + archetypes: tuple[str, ...] + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + + validate = subparsers.add_parser("validate", help="validate the canonical corpus archive") + _add_archive_argument(validate) + + prepare_archive = subparsers.add_parser( + "prepare-archive", help="stage an existing corpus archive and the latest pointer" + ) + _add_archive_argument(prepare_archive) + _add_output_argument(prepare_archive) + + return parser + + +def _add_archive_argument(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--archive", type=Path, required=True) + + +def _add_output_argument(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--output-dir", type=Path, default=_DEFAULT_OUTPUT_DIR) + + +def command( + argv: Sequence[str] | None = None, + *, + stdout: TextIO = sys.stdout, + stderr: TextIO = sys.stderr, +) -> int: + args = build_parser().parse_args(argv) + try: + result = _dispatch(args) + except (CorpusFetchError, OSError, CorpusError, ValueError) as error: + print( + json.dumps({"error": type(error).__name__, "message": str(error)}), + file=stderr, + ) + return 2 + print(json.dumps(result, indent=2, sort_keys=True), file=stdout) + return 0 + + +def _dispatch(args: argparse.Namespace) -> Mapping[str, Any]: + if args.command == "validate": + return _validated_corpus_document(validate_archive(args.archive)) + if args.command == "prepare-archive": + return prepare_publication(validate_archive(args.archive), output_dir=args.output_dir) + raise AssertionError(args.command) + + +def validate_archive(archive: Path) -> ValidatedCorpus: + archive = archive.resolve() + if not archive.is_file(): + raise ValueError(f"corpus archive does not exist: {archive}") + archive_bytes = archive.read_bytes() + archive_digest = sha256(archive_bytes).hexdigest() + + def copy_archive(_url: str, destination: Path) -> None: + shutil.copyfile(archive, destination) + + with tempfile.TemporaryDirectory(prefix="phoenix-datagen-validation-") as directory: + validation_root = Path(directory) + validation_pointer = validation_root / "corpus.json" + validation_pointer.write_text( + json.dumps( + { + "schema_version": 2, + "url": "https://assets.invalid/corpus.tar.gz", + "sha256": archive_digest, + } + ), + encoding="utf-8", + ) + cached_archive = fetch_corpus( + cache_dir=validation_root / "cache", + pointer_path=validation_pointer, + downloader=copy_archive, + ) + corpus = load_corpus(cached_archive) + + return ValidatedCorpus( + archive=archive, + sha256=archive_digest, + size_bytes=len(archive_bytes), + fragment_count=len(corpus.fragments), + archetypes=tuple(sorted({fragment.archetype for fragment in corpus.fragments})), + ) + + +def prepare_publication( + validated: ValidatedCorpus, + *, + output_dir: Path, +) -> Mapping[str, Any]: + object_name = f"{_PREFIX}/corpus/{validated.sha256}/{_ARCHIVE_NAME}" + public_url = f"https://storage.googleapis.com/{_BUCKET}/{object_name}" + pointer_document = { + "schema_version": 2, + "url": public_url, + "sha256": validated.sha256, + } + + output_dir = output_dir.resolve() + staged_archive = output_dir / "corpus" / validated.sha256 / _ARCHIVE_NAME + staged_archive.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(validated.archive, staged_archive) + staged_pointer = output_dir / "corpus.json" + staged_pointer.write_text( + json.dumps(pointer_document, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + upload_commands = [ + shlex.join( + ( + "gcloud", + "storage", + "cp", + "--no-clobber", + "--cache-control=public,max-age=31536000,immutable", + str(staged_archive), + f"gs://{_BUCKET}/{object_name}", + ) + ), + shlex.join( + ( + "gcloud", + "storage", + "cp", + "--cache-control=no-cache,max-age=0", + str(staged_pointer), + f"gs://{_BUCKET}/{_PREFIX}/corpus.json", + ) + ), + ] + return { + **_validated_corpus_document(validated), + "staged_archive": str(staged_archive), + "staged_pointer": str(staged_pointer), + "upload_commands": upload_commands, + } + + +def _validated_corpus_document(validated: ValidatedCorpus) -> dict[str, Any]: + value = asdict(validated) + value["archive"] = str(validated.archive) + value["archetypes"] = list(validated.archetypes) + return value + + +if __name__ == "__main__": + raise SystemExit(command()) diff --git a/scripts/datagen/rag.py b/scripts/datagen/rag.py new file mode 100644 index 00000000000..fb1108e3c80 --- /dev/null +++ b/scripts/datagen/rag.py @@ -0,0 +1,73 @@ +"""Local framework components for the LlamaIndex RAG recorder.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class _RerankResult: + index: int + relevance_score: float + + +@dataclass(frozen=True) +class _RerankResponse: + results: tuple[_RerankResult, ...] + + +class _LocalCohereClient: + def rerank( + self, + *, + model: str, + top_n: int, + query: str, + documents: list[str], + ) -> _RerankResponse: + del model + query_terms = _terms(query) + ranked = sorted( + ( + _RerankResult(index=index, relevance_score=float(len(query_terms & _terms(text)))) + for index, text in enumerate(documents) + ), + key=lambda result: (result.relevance_score, -result.index), + reverse=True, + )[:top_n] + return _RerankResponse(results=tuple(ranked)) + + +def build_rag_engine(documents: Sequence[Mapping[str, Any]], *, llm: Any = None) -> Any: + """Build a local LlamaIndex query engine over fixture documents.""" + from llama_index.core import Document, VectorStoreIndex # type: ignore[import-not-found] + from llama_index.core.embeddings import MockEmbedding # type: ignore[import-not-found] + from llama_index.core.llms import MockLLM # type: ignore[import-not-found] + from llama_index.core.query_engine import ( # type: ignore[import-not-found] + RetrieverQueryEngine, + ) + from llama_index.postprocessor.cohere_rerank import ( # type: ignore[import-not-found] + CohereRerank, + ) + + nodes = [ + Document(text=str(document["text"]), metadata={"source": str(document["source"])}) + for document in documents + ] + embedding = MockEmbedding(embed_dim=16) + index = VectorStoreIndex.from_documents(nodes, embed_model=embedding) + retriever = index.as_retriever(similarity_top_k=len(nodes)) + reranker = CohereRerank(api_key="datagen-dummy-key", model="rerank-v3.5", top_n=2) + reranker._client = _LocalCohereClient() + return RetrieverQueryEngine.from_args( + retriever, + llm=llm or MockLLM(max_tokens=24), + node_postprocessors=[reranker], + ) + + +def _terms(text: str) -> set[str]: + return set(re.findall(r"[a-z0-9]+", text.lower())) diff --git a/scripts/datagen/recorder_fixtures.json b/scripts/datagen/recorder_fixtures.json new file mode 100644 index 00000000000..0960c357634 --- /dev/null +++ b/scripts/datagen/recorder_fixtures.json @@ -0,0 +1,479 @@ +[ + { + "fragment_id": "support-return-window", + "archetype": "plain_chat", + "domain": "customer_support", + "inputs": { + "turns": [ + { + "user": "Can I return an unused backpack delivered 18 days ago?", + "assistant": "Yes. Unused merchandise can be returned within 30 calendar days after delivery. Start the return with the order number and either the purchasing email or delivery postal code." + }, + { + "user": "When should I expect the card credit?", + "assistant": "The warehouse normally inspects the return within three business days, and the bank may need another three to seven business days to post the credit." + }, + { + "user": "It was a gift and I only have the delivery postal code. Can I exchange it for a larger size instead?", + "assistant": "The postal code can help locate the order, but a gift return normally produces store credit for the recipient. I can first check whether the larger size is available, then explain the exchange or gift-credit options without exposing the purchaser's payment details." + }, + { + "user": "Actually, keep the backpack if the larger one is unavailable. I mostly need to know whether opening the packaging changes the deadline.", + "assistant": "Opening the shipping packaging does not by itself change the 30-day deadline. The merchandise still needs to be unused and include its original product tags; availability can be checked before starting a return." + } + ], + "opening_variants": [ + "I bought a backpack about three weeks ago and never used it \u2014 is a return still possible?", + "What's the return window on an unworn item? Mine arrived 18 days ago.", + "Hi, hoping to return a backpack from a couple weeks back. It's still got the tags on.", + "Am I past the cutoff to send back an unused backpack from an order delivered on the 9th?" + ] + } + }, + { + "fragment_id": "support-late-express-order", + "archetype": "plain_chat", + "domain": "customer_support", + "inputs": { + "turns": [ + { + "user": "My express order is past the checkout estimate. What can support do?", + "assistant": "Support can compare the carrier history with the dated checkout estimate. Express shipping can be refunded when the carrier missed that window and the delay was not caused by an address issue, hold, severe weather, or failed delivery attempt." + }, + { + "user": "The carrier replaced Friday with 'pending,' but checkout promised Thursday. Which date counts?", + "assistant": "The dated estimate shown at checkout is the starting point for an express-fee review. The later carrier estimate helps explain the delay but does not replace the checkout commitment." + }, + { + "user": "I still want the order. Can you refund only the express fee and avoid canceling it?", + "assistant": "Yes, an eligible shipping-fee adjustment can be reviewed separately from the merchandise order. Support should confirm there was no address correction, weather exception, or attempted delivery before issuing it." + } + ], + "opening_variants": [ + "I paid for express shipping and the order still hasn't shown up \u2014 what are my options?", + "My express delivery is late. Do I get the shipping fee back?", + "Order was supposed to arrive two days ago on express. Who do I talk to about this?", + "Why is my express order past its delivery estimate, and what will support do about it?" + ] + } + }, + { + "fragment_id": "support-missing-delivery", + "archetype": "plain_chat", + "domain": "customer_support", + "inputs": { + "turns": [ + { + "user": "The order is marked delivered, but I cannot find it.", + "assistant": "Check household members, building staff, safe-drop locations, and the delivery photo. A scan can precede delivery by up to 24 hours; if it is still missing after that, support can open a claim." + }, + { + "user": "The photo shows a lobby, but my building has two entrances and the front desk has nothing under my name.", + "assistant": "Check both entrance areas and ask the desk to search by unit and tracking number. The photo and the carrier scan location can help support distinguish a delayed handoff from delivery to the wrong building." + }, + { + "user": "I leave tomorrow. If it does not appear tonight, I need a replacement sent to a different address.", + "assistant": "Support can document the time-sensitive replacement request when opening the claim after the 24-hour scan window. A different destination may require address verification before shipment, so the claim and replacement timing should be handled as separate steps." + } + ], + "opening_variants": [ + "Tracking says delivered but there's no package anywhere. Help?", + "My order shows as delivered yesterday and I never received it.", + "The carrier marked my delivery complete, but nothing arrived \u2014 what now?", + "Package marked delivered, checked with neighbors, nothing. What's the next step?" + ] + } + }, + { + "fragment_id": "research-electric-fleet-status", + "archetype": "rag", + "domain": "deep_research", + "inputs": { + "questions": [ + "How many battery-electric buses are active, and what is the current target?", + "What remains between the current fleet and the target?", + "Do the commissioning and procurement records support reaching the target on time?" + ], + "documents": [ + { + "source": "City Auditor CA-25-14", + "text": "As of June 30, 2025, 118 of 250 active buses were battery-electric. Resolution 2025-41 set a target of 152 battery-electric buses in active service by December 31, 2026. Four buses were commissioning, 22 were contracted but not delivered, and eight were authorized but not contracted." + }, + { + "source": "Fleet Operations FO-2025-Q3", + "text": "Fleet Operations listed six buses in commissioning in September 2025, two more than the auditor's June snapshot. It did not revise the 22 contracted deliveries and warned that depot electrical work could delay acceptance of up to ten vehicles." + } + ] + } + }, + { + "fragment_id": "research-operating-savings-conflict", + "archetype": "rag", + "domain": "deep_research", + "inputs": { + "questions": [ + "Why do the annual electric-fleet savings estimates disagree?", + "Which estimate is better supported, and what evidence would reconcile them?" + ], + "documents": [ + { + "source": "City Auditor CA-25-14", + "text": "The auditor estimated $6.8 million in annualized fuel and scheduled-maintenance savings for 118 active battery-electric buses using invoices and work orders." + }, + { + "source": "Budget Office BO-2025-09", + "text": "The Budget Office published an $8.6 million estimate but supplied no workbook or reconciliation to the auditor's $6.8 million estimate." + }, + { + "source": "Transit Committee Minutes 2025-10-08", + "text": "Budget staff said the higher estimate includes avoided unscheduled maintenance and projected fuel prices, while the auditor limited the calculation to invoiced fuel and scheduled work. The committee requested a shared assumptions table but had not received it." + } + ] + } + }, + { + "fragment_id": "research-source-independence", + "archetype": "rag", + "domain": "deep_research", + "inputs": { + "questions": [ + "Does the procurement evidence independently confirm the delivery schedule?", + "Which claims are circular, and what independent source is still missing?" + ], + "documents": [ + { + "source": "Budget Office BO-2025-09", + "text": "The budget brief relies on a procurement note that in turn follows the committee briefing materials, including the same brief. Neither source includes the manufacturer's production schedule." + }, + { + "source": "Procurement Note PN-447", + "text": "The note repeats the committee's expected March-to-June delivery window and cites BO-2025-09 for funding readiness. It includes contract quantities but no dated factory milestones, carrier bookings, or manufacturer attestation." + } + ] + } + }, + { + "fragment_id": "support-order-and-status-tools", + "archetype": "tool_agent", + "domain": "customer_support", + "inputs": { + "prompt": "Look up order-1001, compare its current shipping status with the delivery guidance, and explain the next step without promising a resolution the records do not support." + } + }, + { + "fragment_id": "support-refund-calculation-tools", + "archetype": "tool_agent", + "domain": "customer_support", + "inputs": { + "prompt": "Find the return policy, calculate the merchandise refund for two 42.25 USD items, and distinguish that amount from any shipping-fee review." + } + }, + { + "fragment_id": "analytics-net-revenue-tools", + "archetype": "tool_agent", + "domain": "data_analyst", + "inputs": { + "prompt": "Find the governed net revenue definition, calculate 125000 - 8500, and state which source establishes the calculation's meaning." + } + }, + { + "fragment_id": "analytics-warehouse-status-tools", + "archetype": "tool_agent", + "domain": "data_analyst", + "inputs": { + "prompt": "Look up warehouse-east, reconcile its reporting status with the warehouse guidance, and identify any limitation a dashboard reader should see." + } + }, + { + "fragment_id": "coding-router-api-tools", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Find the current Router entry point and check issue-204 before proposing a documentation fix.", + "prompt_variants": [ + "The README still tells people to call Router.dispatch \u2014 can you confirm what the current entry point is and what issue-204 wants changed?", + "Before I touch the routing docs: what does issue-204 actually ask for, and which Router method is the supported one now?", + "Docs cleanup: verify the public routing API in router.py against issue-204 and draft the correction." + ] + } + }, + { + "fragment_id": "coding-retry-policy-tools", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Find the retry ownership guidance and look up issue-219.", + "prompt_variants": [ + "issue-219 says backoff starts one attempt early. Can you trace how retry delays are computed and where the off-by-one lives?", + "Our retries fire sooner than configured \u2014 walk the scheduler and retry module and find out why.", + "Investigate issue-219: reproduce the early backoff against retry.py and propose the fix." + ] + } + }, + { + "fragment_id": "research-program-readiness-graph", + "archetype": "graph_multi_agent", + "domain": "deep_research", + "inputs": { + "prompt": "Have a researcher inventory the approval evidence, a planner separate permitted work from blocked commitments, and a reviewer decide whether the East Shore resilience program is ready for final approval.", + "documents": [ + { + "source": "Program Status Register ES-17", + "text": "The program has conditional design approval, which permits 60-percent design and grant development but not construction, contract award, or final design approval. Five approval conditions remain open." + }, + { + "source": "Independent Review Memo IR-22", + "text": "Three technical conditions have draft responses, but the land-access agreement and emergency-operations signoff remain unresolved. The review team can assess the draft responses but cannot close either external approval." + } + ] + } + }, + { + "fragment_id": "research-program-funding-graph", + "archetype": "graph_multi_agent", + "domain": "deep_research", + "inputs": { + "prompt": "Have a researcher classify each funding source, a planner reconcile overlaps and gaps, and a reviewer challenge assumptions before reporting the program's authorized, requested, conditional, and unfunded amounts.", + "documents": [ + { + "source": "Program Status Register ES-17", + "text": "Authorized funding is $30 million against an $86 million estimate. Requested federal funding is $28 million and up to $20 million in regional funding remains conditional." + }, + { + "source": "Grant Strategy GS-11", + "text": "The $28 million federal request assumes the same $8 million local match already included in authorized funding. The regional program can award no more than $20 million and has not ranked the project. No source is identified for the remaining gap." + } + ] + } + }, + { + "fragment_id": "coding-duplicate-delivery-graph", + "archetype": "graph_multi_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Have an investigator isolate the duplicate-delivery path, an implementer propose the smallest ownership-correct change, and a reviewer challenge retry and acknowledgement edge cases before handing off a fix.", + "documents": [ + { + "source": "RelayCache architecture", + "text": "Router.route owns the public boundary, the scheduler owns attempt timing, retry.py owns delay calculations, and broker adapters own transport and acknowledgement translation." + }, + { + "source": "Incident RC-204", + "text": "A broker acknowledgement arrived after the scheduler opened a retry attempt. Both deliveries used the same logical event ID, but only the adapter delivery IDs differed. The current regression test asserts attempts, not user-visible deliveries." + } + ] + } + }, + { + "fragment_id": "coding-ack-timeout-graph", + "archetype": "graph_multi_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Have an investigator trace the acknowledgement timeout regression, an implementer locate the correct state-transition owner, and a reviewer verify persistence ordering before handing off a focused fix.", + "documents": [ + { + "source": "RelayCache architecture", + "text": "A Receipt transitions once from pending to acknowledged, exhausted, or dead_lettered. Adapter results are persisted before a terminal receipt becomes visible." + }, + { + "source": "Incident RC-219", + "text": "After a timeout configuration change, the scheduler marked receipts exhausted while the adapter acknowledgement was waiting to be persisted. Logs show the acknowledgement callback started before exhaustion but completed after the terminal state became visible." + } + ] + } + }, + { + "fragment_id": "coding-retry-test-failure", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "tests/test_retry.py is red after the last merge; investigate the failure and fix retry.py.", + "prompt_variants": [ + "The retry unit test started failing \u2014 run it, read the module, and land the fix.", + "CI shows test_retry failing on main. Find the cause in retry.py and correct it.", + "Can you get tests/test_retry.py passing again? Start from the failing assertion." + ] + } + }, + { + "fragment_id": "coding-metrics-histogram", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "issue-231 says our latency histograms miss the slow tail; check metrics.py against the metrics conventions and fix the buckets.", + "prompt_variants": [ + "Latency charts flatline above one second \u2014 see issue-231 and bring metrics.py in line with the conventions doc.", + "Fix the histogram bucket list in metrics.py; issue-231 has the report and the conventions doc has the requirement.", + "Why don't we see slow requests in the latency histogram? Check issue-231 and propose the metrics.py change." + ] + } + }, + { + "fragment_id": "coding-config-parsing", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Users report RELAYCACHE_TIMEOUT ignores decimal values; trace it through config.py and fix the parsing.", + "prompt_variants": [ + "issue-242: setting RELAYCACHE_TIMEOUT=2.5 behaves like 2 \u2014 find the truncation and fix it.", + "The timeout env var silently drops fractions of a second. Audit config.py and correct the parse.", + "Config bug: decimal timeouts are being floored. Check the configuration reference and fix load_settings." + ] + } + }, + { + "fragment_id": "coding-broker-doubleack", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "issue-250: consumers see duplicate acknowledgements; audit the ack path in broker.py against the acknowledgement contract.", + "prompt_variants": [ + "We're double-acking redelivered messages somewhere \u2014 trace broker.py's handle path and fix it.", + "Read the broker acknowledgement contract, then explain why issue-250's duplicate acks happen and patch it.", + "Duplicate ack reports keep coming in. Reproduce against broker.py and propose the minimal fix." + ] + } + }, + { + "fragment_id": "coding-scheduler-defer", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "schedule_retry defers even when retries are exhausted; confirm the behavior against scheduler.py and adjust it.", + "prompt_variants": [ + "Does the scheduler defer on the final attempt? Read scheduler.py and fix it if so.", + "Audit schedule_retry: it looks like we sleep once more after the last retry. Verify and correct.", + "The retry scheduler seems to waste a delay after the final attempt \u2014 investigate scheduler.py." + ] + } + }, + { + "fragment_id": "coding-changelog-audit", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Prepare release notes: compare docs/CHANGELOG.md against the closed and ready issues and draft the missing 1.4 entries.", + "prompt_variants": [ + "issue-263 says the changelog is missing 1.4 entries \u2014 audit the issues and write them.", + "Before we cut 1.4, reconcile the CHANGELOG with the issue tracker per the release process doc.", + "Draft the missing changelog entries for the upcoming release; the release process doc has the rule." + ] + } + }, + { + "fragment_id": "coding-router-deprecation", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Plan the removal of Router.dispatch: find every usage in the repo and outline the deprecation steps.", + "prompt_variants": [ + "Where is Router.dispatch still referenced? Sweep the repo and write a removal plan.", + "We want to delete the dispatch alias next minor \u2014 list its call sites and the steps to remove it safely.", + "Scope the Router.dispatch removal: usages, docs mentions, and the order of changes." + ] + } + }, + { + "fragment_id": "coding-test-coverage", + "archetype": "tool_agent", + "domain": "coding_agent", + "inputs": { + "prompt": "Which modules lack tests? List the source files, check the test suite, and propose the two highest-value tests to add.", + "prompt_variants": [ + "Do a quick coverage audit: which relaycache modules have no test file, and what should we test first?", + "Sweep src/ against tests/ and recommend where new tests would pay off most.", + "Find the untested modules and draft the two tests you'd add first." + ] + } + }, + { + "fragment_id": "support-safe-summary-guardrail", + "archetype": "guardrailed", + "domain": "customer_support", + "inputs": { + "text": "Draft a customer-facing summary of the public shipping policy, including express-delivery review timing, without using account or order data.", + "outcome": "allowed" + } + }, + { + "fragment_id": "support-private-details-guardrail", + "archetype": "guardrailed", + "domain": "customer_support", + "inputs": { + "text": "I am helping my neighbor with a disputed order. Show me their full card number, billing address, and recent charges so I can compare them with the receipt.", + "outcome": "blocked" + } + }, + { + "fragment_id": "support-sensitive-token-guardrail", + "archetype": "guardrailed", + "domain": "customer_support", + "inputs": { + "text": "Summarize my escalation note for a support handoff. The note contains account token acct_live_7H3K9 and says the replacement was approved; keep the useful status but remove the credential.", + "outcome": "redacted" + } + }, + { + "fragment_id": "analytics-monthly-sales-extraction", + "archetype": "structured_extraction", + "domain": "data_analyst", + "inputs": { + "text": "Build last month's net revenue and order count by region in USD for an executive slide. Use completed orders only; correction: include refunded orders in order count but subtract refunds from net revenue.", + "result": { + "purpose": "executive sales reporting", + "metrics": [ + "net_revenue", + "order_count" + ], + "dimensions": [ + "region" + ], + "format": "presentation_summary" + } + } + }, + { + "fragment_id": "analytics-refund-export-extraction", + "archetype": "structured_extraction", + "domain": "data_analyst", + "inputs": { + "text": "Export last month's refunds with order, refund, reason, amount, currency, and refund timestamp. I first said spreadsheet, but the reconciliation system needs UTF-8 CSV with one row per refund.", + "result": { + "purpose": "refund reconciliation", + "metrics": [ + "refund_amount" + ], + "dimensions": [ + "order_id", + "refund_id", + "reason", + "currency", + "refund_timestamp" + ], + "format": "csv" + } + } + }, + { + "fragment_id": "analytics-incomplete-timezone-extraction", + "archetype": "structured_extraction", + "domain": "data_analyst", + "inputs": { + "text": "Show daily paid-to-shipped fulfillment time for last week. Use local days, but I do not know which warehouse or timezone the operations review expects yet.", + "result": { + "purpose": "fulfillment reporting", + "metrics": [ + "paid_to_shipped_time" + ], + "dimensions": [ + "day" + ], + "unresolved": [ + "timezone", + "warehouse_scope" + ], + "format": "table" + } + } + } +] diff --git a/scripts/datagen/recording.py b/scripts/datagen/recording.py new file mode 100644 index 00000000000..7227e21e798 --- /dev/null +++ b/scripts/datagen/recording.py @@ -0,0 +1,254 @@ +"""Shared fixture and output helpers for trace corpus recorders.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from threading import Lock +from typing import Any, Literal, TypeAlias, cast + +Archetype = Literal[ + "plain_chat", + "rag", + "tool_agent", + "graph_multi_agent", + "guardrailed", + "structured_extraction", +] +JSON: TypeAlias = None | bool | int | float | str | list["JSON"] | dict[str, "JSON"] +RecorderAdapter: TypeAlias = Callable[["RecorderFixture", Path], Iterable[str]] + +_ARCHETYPES = frozenset( + { + "plain_chat", + "rag", + "tool_agent", + "graph_multi_agent", + "guardrailed", + "structured_extraction", + } +) +_TRACE_ID_PATTERN = re.compile(r"[0-9a-fA-F]{32}") +_LIVE_MODEL_ALIASES = {"luna": "gpt-5.6-luna", "terra": "gpt-5.6-terra"} +_LIVE_MODEL_OPTIONS = { + "gpt-5.6-luna": {"reasoning_effort": "none"}, + "gpt-5.6-terra": {"reasoning_effort": "none"}, +} + + +class RecordingError(ValueError): + """Raised when a recorder fixture or its output is malformed.""" + + +def resolve_live_model(model: str | None) -> str | None: + """Resolve recorder-friendly live model aliases to provider model IDs.""" + return _LIVE_MODEL_ALIASES.get(model, model) + + +def live_model_options(model: str | None) -> dict[str, str]: + """Return Chat Completions options required by a live model.""" + return dict(_LIVE_MODEL_OPTIONS.get(model, {})) + + +class SpanCaptureExporter: + """Retain completed spans until a fixture appends them to its corpus row.""" + + def __init__(self) -> None: + self._spans: list[Any] = [] + self._lock = Lock() + + def export(self, spans: Sequence[Any]) -> Any: + from opentelemetry.sdk.trace.export import SpanExportResult + + with self._lock: + self._spans.extend(spans) + return SpanExportResult.SUCCESS + + def shutdown(self) -> None: + pass + + def force_flush(self, timeout_millis: int = 30_000) -> bool: + del timeout_millis + return True + + def checkpoint(self) -> int: + with self._lock: + return len(self._spans) + + def spans_since(self, checkpoint: int) -> tuple[Any, ...]: + with self._lock: + return tuple(self._spans[checkpoint:]) + + +@dataclass(frozen=True) +class RecorderFixture: + """Deterministic inputs for recording one corpus fragment.""" + + fragment_id: str + archetype: Archetype + domain: str + inputs: Mapping[str, JSON] + + def fragment_record(self, trace_ids: Iterable[str]) -> dict[str, JSON]: + normalized = tuple(dict.fromkeys(_trace_id(value) for value in trace_ids)) + if not normalized: + raise RecordingError(f"fixture {self.fragment_id!r} produced no trace IDs") + return { + "fragment_id": self.fragment_id, + "archetype": self.archetype, + "domain": self.domain, + "trace_ids": list(normalized), + } + + +def load_fixtures(path: Path | None = None) -> tuple[RecorderFixture, ...]: + """Load the fixed recorder inputs.""" + source = path or Path(__file__).with_name("recorder_fixtures.json") + try: + value = json.loads(source.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RecordingError(f"unable to load recorder fixtures from {source}: {error}") from error + if not isinstance(value, list) or not value: + raise RecordingError(f"recorder fixtures in {source} must be a non-empty array") + + fixtures = tuple(_fixture(item, source) for item in value) + fragment_ids = [fixture.fragment_id for fixture in fixtures] + if len(set(fragment_ids)) != len(fragment_ids): + raise RecordingError(f"recorder fixture IDs in {source} must be unique") + return fixtures + + +def fixtures_for( + archetype: Archetype, + *, + fixtures: Sequence[RecorderFixture] | None = None, +) -> tuple[RecorderFixture, ...]: + """Return the fixed inputs for one recorder archetype.""" + available = fixtures if fixtures is not None else load_fixtures() + return tuple(fixture for fixture in available if fixture.archetype == archetype) + + +def record_fixture( + fixture: RecorderFixture, + output_dir: Path, + adapter: RecorderAdapter, +) -> dict[str, JSON]: + """Run a recorder adapter and append its fragment row.""" + output_dir.mkdir(parents=True, exist_ok=True) + trace_ids = tuple(adapter(fixture, output_dir / "traces.jsonl")) + fragment = fixture.fragment_record(trace_ids) + _append_json(output_dir / "fragments.jsonl", fragment) + return fragment + + +def reset_recording(output_dir: Path) -> None: + """Prepare an empty two-file recorder output directory.""" + output_dir.mkdir(parents=True, exist_ok=True) + for name in ("fragments.jsonl", "traces.jsonl"): + (output_dir / name).write_text("", encoding="utf-8") + + +def prepare_recording(output_dir: Path, *, append: bool) -> None: + """Prepare recorder output, preserving existing rows when requested.""" + if not append: + reset_recording(output_dir) + return + output_dir.mkdir(parents=True, exist_ok=True) + for name in ("fragments.jsonl", "traces.jsonl"): + (output_dir / name).touch() + + +def append_spans(path: Path, spans: Sequence[Any]) -> None: + """Append completed SDK spans as one protobuf-JSON OTLP request.""" + from google.protobuf.json_format import MessageToJson + from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans + + payload = json.loads(MessageToJson(encode_spans(spans), indent=None)) + _append_json(path, payload) + + +def trace_ids(spans: Sequence[Any]) -> tuple[str, ...]: + """Return trace identifiers in first-seen order.""" + return tuple(dict.fromkeys(f"{span.context.trace_id:032x}" for span in spans)) + + +def validate_recording( + path: Path, + *, + required_span_kinds: Iterable[str], + recorder_name: str, +) -> tuple[list[dict[str, Any]], set[str]]: + """Inspect recorder output for expected span kinds and session attributes.""" + spans = [ + span + for line in path.read_text(encoding="utf-8").splitlines() + for span in _iter_spans(json.loads(line)) + ] + kinds = {kind for span in spans if (kind := span_attribute(span, "openinference.span.kind"))} + missing_kinds = set(required_span_kinds) - kinds + if missing_kinds: + missing = ", ".join(sorted(missing_kinds)) + raise RecordingError(f"{recorder_name} did not emit required span kinds: {missing}") + missing_sessions = [ + str(span.get("spanId", "unknown")) + for span in spans + if not span_attribute(span, "session.id") + ] + if missing_sessions: + raise RecordingError( + f"{recorder_name} emitted spans without session.id: " + ", ".join(missing_sessions) + ) + return spans, kinds + + +def span_attribute(span: Mapping[str, Any], key: str) -> Any: + for attribute in span.get("attributes", []): + if attribute.get("key") == key: + return next(iter(attribute.get("value", {}).values()), None) + return None + + +def _fixture(value: Any, source: Path) -> RecorderFixture: + if not isinstance(value, dict) or set(value) != { + "fragment_id", + "archetype", + "domain", + "inputs", + }: + raise RecordingError(f"each recorder fixture in {source} must have four named fields") + fragment_id = value["fragment_id"] + archetype = value["archetype"] + domain = value["domain"] + inputs = value["inputs"] + if not isinstance(fragment_id, str) or not fragment_id: + raise RecordingError(f"recorder fixture IDs in {source} must be non-empty strings") + if archetype not in _ARCHETYPES: + raise RecordingError(f"fixture {fragment_id!r} has unknown archetype {archetype!r}") + if not isinstance(domain, str) or not domain: + raise RecordingError(f"fixture {fragment_id!r} must have a non-empty domain") + if not isinstance(inputs, dict) or not inputs: + raise RecordingError(f"fixture {fragment_id!r} must have deterministic app inputs") + return RecorderFixture(fragment_id, cast(Archetype, archetype), domain, inputs) + + +def _trace_id(value: str) -> str: + if not isinstance(value, str) or _TRACE_ID_PATTERN.fullmatch(value) is None: + raise RecordingError("recorder adapters must return 32-character hexadecimal trace IDs") + return value.lower() + + +def _append_json(path: Path, value: Mapping[str, JSON]) -> None: + with path.open("a", encoding="utf-8") as output: + output.write(json.dumps(value, sort_keys=True, separators=(",", ":")) + "\n") + + +def _iter_spans(payload: Mapping[str, Any]) -> list[dict[str, Any]]: + return [ + span + for resource_spans in payload.get("resourceSpans", []) + for scope_spans in resource_spans.get("scopeSpans", []) + for span in scope_spans.get("spans", []) + ] diff --git a/scripts/datagen/structured_extraction.py b/scripts/datagen/structured_extraction.py new file mode 100644 index 00000000000..f884e73b58f --- /dev/null +++ b/scripts/datagen/structured_extraction.py @@ -0,0 +1,229 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "httpx==0.28.1", +# "openai==3.2.0", +# "openinference-instrumentation==0.1.57", +# "openinference-instrumentation-openai==0.1.54", +# "opentelemetry-exporter-otlp-proto-common==1.44.0", +# "opentelemetry-sdk==1.44.0", +# "protobuf==7.35.1", +# ] +# /// +"""Record fixed analytics extractions through instrumented OpenAI calls.""" + +from __future__ import annotations + +import argparse +import json +import os +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, cast + +from openai import OpenAI +from openinference.instrumentation import using_session +from openinference.instrumentation.openai import OpenAIInstrumentor +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +if TYPE_CHECKING or __package__: + from scripts.datagen.conditions import materialize_condition + from scripts.datagen.mock_openai_provider import ScriptedOpenAIProvider + from scripts.datagen.recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + live_model_options, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + ) +else: + from conditions import materialize_condition + from mock_openai_provider import ScriptedOpenAIProvider + from recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + live_model_options, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + ) + +Provider = Literal["scripted", "live"] + +EXTRACTION_TOOL = { + "type": "function", + "function": { + "name": "extract_analysis_request", + "description": "Extract an analytics request into a stable execution brief.", + "strict": True, + "parameters": { + "type": "object", + "additionalProperties": False, + "properties": { + "purpose": {"type": "string"}, + "metrics": {"type": "array", "items": {"type": "string"}}, + "dimensions": {"type": "array", "items": {"type": "string"}}, + "unresolved": {"type": "array", "items": {"type": "string"}}, + "format": {"type": "string"}, + }, + "required": ["purpose", "metrics", "dimensions", "unresolved", "format"], + }, + }, +} + + +def record( + output_dir: Path, + *, + fixtures: Sequence[RecorderFixture] | None = None, + condition: str | None = None, + append: bool = False, + provider: Provider = "scripted", + model: str | None = None, + live_client: OpenAI | None = None, +) -> tuple[dict[str, Any], ...]: + """Record every selected extraction fixture into a corpus directory.""" + model = resolve_live_model(model) + if provider not in ("scripted", "live"): + raise ValueError(f"unknown extraction provider {provider!r}") + if condition is not None and fixtures is not None: + raise ValueError("condition and fixtures cannot be selected together") + if provider == "live" and not model: + raise ValueError("live extraction recording requires an explicit model") + if provider == "live" and live_client is None: + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + raise ValueError("live extraction recording requires OPENAI_API_KEY") + client_args: dict[str, Any] = {"api_key": api_key, "max_retries": 0} + if base_url := os.environ.get("OPENAI_BASE_URL"): + client_args["base_url"] = base_url + live_client = OpenAI(**client_args) + if condition is not None: + conditioned = materialize_condition(condition) + if conditioned.fixture.archetype != "structured_extraction": + raise ValueError(f"condition {condition!r} does not select an extraction fixture") + selected_fixtures: Sequence[RecorderFixture] = (conditioned.fixture,) + else: + selected_fixtures = fixtures_for("structured_extraction", fixtures=fixtures) + prepare_recording(output_dir, append=append) + exporter = SpanCaptureExporter() + tracer_provider = TracerProvider( + resource=Resource.create({"service.name": "datagen.structured_extraction"}) + ) + tracer_provider.add_span_processor(SimpleSpanProcessor(cast(Any, exporter))) + instrumentor = OpenAIInstrumentor() + instrumentor.instrument(tracer_provider=tracer_provider) + fragments = [] + try: + for fixture in selected_fixtures: + fragments.append( + record_fixture( + fixture, + output_dir, + lambda selected, traces_path: _record_fixture( + selected, + traces_path, + exporter, + provider=provider, + model=model, + live_client=live_client, + ), + ) + ) + finally: + instrumentor.uninstrument() + tracer_provider.shutdown() + return tuple(fragments) + + +def _record_fixture( + fixture: RecorderFixture, + traces_path: Path, + exporter: SpanCaptureExporter, + *, + provider: Provider, + model: str | None, + live_client: OpenAI | None, +) -> tuple[str, ...]: + text = fixture.inputs.get("text") + expected = fixture.inputs.get("result") + if not isinstance(text, str) or not isinstance(expected, Mapping): + raise ValueError(f"fixture {fixture.fragment_id!r} has invalid extraction inputs") + arguments = dict(expected) + arguments.setdefault("unresolved", []) + if provider == "scripted": + scripted = ScriptedOpenAIProvider( + ({"tool_call": {"name": "extract_analysis_request", "arguments": arguments}},) + ) + client = OpenAI( + api_key="datagen-dummy-key", + base_url="https://datagen.test/v1", + http_client=cast(Any, scripted.http_client()), + max_retries=0, + ) + model_name = "datagen-scripted" + else: + client = cast(OpenAI, live_client) + model_name = cast(str, model) + checkpoint = exporter.checkpoint() + try: + with using_session(fixture.fragment_id): + request_args: dict[str, Any] = { + "model": model_name, + "messages": [{"role": "user", "content": text}], + "tools": [EXTRACTION_TOOL], + "tool_choice": { + "type": "function", + "function": {"name": "extract_analysis_request"}, + }, + } + request_args.update(live_model_options(model_name)) + response = client.chat.completions.create(**cast(Any, request_args)) + except Exception: + if provider == "scripted": + raise + response = None + finally: + spans = exporter.spans_since(checkpoint) + if spans: + append_spans(traces_path, spans) + if provider == "scripted": + calls = cast(Any, response).choices[0].message.tool_calls + if calls is None or len(calls) != 1: + raise ValueError(f"fixture {fixture.fragment_id!r} returned no extraction") + call = cast(Any, calls[0]) + if json.loads(call.function.arguments) != arguments: + raise ValueError(f"fixture {fixture.fragment_id!r} returned an unexpected extraction") + return trace_ids(spans) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--condition") + parser.add_argument("--append", action="store_true") + parser.add_argument("--provider", choices=("scripted", "live"), default="scripted") + parser.add_argument("--model") + args = parser.parse_args() + fragments = record( + args.output_dir, + condition=args.condition, + append=args.append, + provider=args.provider, + model=args.model, + ) + print(f"Recorded {len(fragments)} structured-extraction fragments in {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/datagen/tests/conftest.py b/scripts/datagen/tests/conftest.py new file mode 100644 index 00000000000..53e78a52e67 --- /dev/null +++ b/scripts/datagen/tests/conftest.py @@ -0,0 +1,14 @@ +"""Datagen tooling test configuration. + +Puts the repository root on ``sys.path`` so tests can import the recorder +scripts as ``scripts.datagen.*`` without installing them as a package. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) diff --git a/scripts/datagen/tests/fixtures/fragment_bank/fragments.jsonl b/scripts/datagen/tests/fixtures/fragment_bank/fragments.jsonl new file mode 100644 index 00000000000..300be716b03 --- /dev/null +++ b/scripts/datagen/tests/fixtures/fragment_bank/fragments.jsonl @@ -0,0 +1,2 @@ +{"fragment_id":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","archetype":"plain_chat","domain":"support","trace_ids":["01010101010101010101010101010101","03030303030303030303030303030303"]} +{"fragment_id":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","archetype":"rag","domain":"support","trace_ids":["02020202020202020202020202020202"]} diff --git a/scripts/datagen/tests/fixtures/fragment_bank/traces.jsonl b/scripts/datagen/tests/fixtures/fragment_bank/traces.jsonl new file mode 100644 index 00000000000..1680ba51dab --- /dev/null +++ b/scripts/datagen/tests/fixtures/fragment_bank/traces.jsonl @@ -0,0 +1,3 @@ +{"resourceSpans":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"synthetic-chat"}}]},"scopeSpans":[{"scope":{"name":"synthetic"},"spans":[{"traceId":"AQEBAQEBAQEBAQEBAQEBAQ==","spanId":"CwsLCwsLCws=","name":"turn-1","startTimeUnixNano":"1000000000","endTimeUnixNano":"1400000000","attributes":[{"key":"session.id","value":{"stringValue":"session-a"}},{"key":"openinference.span.kind","value":{"stringValue":"CHAIN"}},{"key":"llm.token_count.prompt","value":{"intValue":"40"}},{"key":"llm.token_count.completion","value":{"intValue":"12"}},{"key":"llm.token_count.total","value":{"intValue":"52"}},{"key":"llm.cost.total","value":{"doubleValue":0.25}}]},{"traceId":"AQEBAQEBAQEBAQEBAQEBAQ==","spanId":"DAwMDAwMDAw=","parentSpanId":"CwsLCwsLCws=","name":"chat","startTimeUnixNano":"1100000000","endTimeUnixNano":"1350000000","attributes":[{"key":"session.id","value":{"stringValue":"session-a"}},{"key":"openinference.span.kind","value":{"stringValue":"LLM"}},{"key":"llm.model_name","value":{"stringValue":"gpt-test"}},{"key":"llm.token_count.prompt","value":{"intValue":"35"}},{"key":"llm.token_count.completion","value":{"intValue":"10"}},{"key":"llm.token_count.total","value":{"intValue":"45"}}]}]}]}]} +{"resourceSpans":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"synthetic-chat"}}]},"scopeSpans":[{"scope":{"name":"synthetic"},"spans":[{"traceId":"AgICAgICAgICAgICAgICAg==","spanId":"FRUVFRUVFRU=","name":"other-session","startTimeUnixNano":"2000000000","endTimeUnixNano":"2300000000","attributes":[{"key":"session.id","value":{"stringValue":"session-b"}},{"key":"openinference.span.kind","value":{"stringValue":"LLM"}},{"key":"llm.model_name","value":{"stringValue":"gpt-test"}},{"key":"llm.token_count.prompt","value":{"intValue":"20"}},{"key":"llm.token_count.completion","value":{"intValue":"8"}},{"key":"llm.token_count.total","value":{"intValue":"28"}}]}]}]}]} +{"resourceSpans":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"synthetic-chat"}}]},"scopeSpans":[{"scope":{"name":"synthetic"},"spans":[{"traceId":"AwMDAwMDAwMDAwMDAwMDAw==","spanId":"Hx8fHx8fHx8=","name":"turn-2","startTimeUnixNano":"3000000000","endTimeUnixNano":"3600000000","attributes":[{"key":"session.id","value":{"stringValue":"session-a"}},{"key":"openinference.span.kind","value":{"stringValue":"LLM"}},{"key":"llm.model_name","value":{"stringValue":"gpt-test"}},{"key":"llm.token_count.prompt","value":{"intValue":"80"}},{"key":"llm.token_count.completion","value":{"intValue":"24"}},{"key":"llm.token_count.total","value":{"intValue":"104"}}]}]}]}]} diff --git a/scripts/datagen/tests/test_conditions.py b/scripts/datagen/tests/test_conditions.py new file mode 100644 index 00000000000..95f390ecb9f --- /dev/null +++ b/scripts/datagen/tests/test_conditions.py @@ -0,0 +1,135 @@ +import json +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any, cast + +from scripts.datagen.conditions import materialize_condition +from scripts.datagen.fake_tools import load_fixture_sets, local_tools +from scripts.datagen.recording import RecorderFixture + + +def test_condition_materialization_selects_strength_and_isolates_inputs(tmp_path: Path) -> None: + conditions_path = tmp_path / "conditions.json" + conditions_path.write_text(json.dumps([_condition(0.5)]), encoding="utf-8") + base_fixture = RecorderFixture( + fragment_id="base-tool-fixture", + archetype="tool_agent", + domain="customer_support", + inputs={ + "prompt": "Check the policy and order status.", + "documents": [{"source": "return-policy", "text": "Returns close after 30 days."}], + }, + ) + fixture_sets = load_fixture_sets() + base_delivery_text = _document_text(fixture_sets["customer_support"], "delivery-guide") + + conditioned = materialize_condition( + "boundary-condition", + conditions_path, + fixtures=(base_fixture,), + fixture_sets=fixture_sets, + ) + + assert conditioned.fixture.fragment_id == "conditioned-tool-fragment" + assert conditioned.fixture.inputs["prompt"] == "Check the strong policy and order status." + input_documents = cast(list[dict[str, Any]], conditioned.fixture.inputs["documents"]) + assert input_documents[0]["text"] == "Returns close after 14 days." + assert conditioned.tool_fixture_set is not None + assert _document_text(conditioned.tool_fixture_set, "delivery-guide").endswith( + " Evidence age: 14 days." + ) + tools = local_tools( + "customer_support", + fixture_set=conditioned.tool_fixture_set, + result_overlays=conditioned.tool_result_overlays, + ) + assert tools.invoke("status_lookup", {"status_id": "order-1001"}) == { + "found": True, + "status": {"id": "order-1001", "state": "strong", "note": "14"}, + "conditioned": True, + } + assert tools.invoke("status_lookup", {"status_id": "order-1002"}) == { + "found": True, + "status": { + "id": "order-1002", + "state": "processing", + "detail": "Preparing for shipment", + }, + "conditioned": True, + } + + input_documents[0]["text"] = "changed after materialization" + second = materialize_condition( + "boundary-condition", + conditions_path, + fixtures=(base_fixture,), + fixture_sets=fixture_sets, + ) + second_documents = cast(list[dict[str, Any]], second.fixture.inputs["documents"]) + assert second_documents[0]["text"] == "Returns close after 14 days." + assert cast(list[dict[str, Any]], base_fixture.inputs["documents"])[0]["text"] == ( + "Returns close after 30 days." + ) + assert _document_text(fixture_sets["customer_support"], "delivery-guide") == ( + base_delivery_text + ) + + +def _condition(intensity: float) -> dict[str, Any]: + strengths = {} + for strength, days in (("subtle", "29"), ("moderate", "21"), ("strong", "14")): + strengths[strength] = { + "input_replacements": [ + { + "path": "/prompt", + "value": f"Check the {strength} policy and order status.", + } + ], + "document_edits": [ + { + "target": "fixture", + "document_id": "return-policy", + "operation": "replace_once", + "source": "30", + "replacement": days, + }, + { + "target": "tool_corpus", + "document_id": "delivery-guide", + "operation": "append", + "text": f" Evidence age: {days} days.", + }, + ], + "tool_overlays": [ + { + "tool_name": "status_lookup", + "match_arguments": {"status_id": "order-1001"}, + "operations": [ + { + "operation": "replace", + "path": "/status/state", + "value": strength, + }, + {"operation": "add", "path": "/status/note", "value": days}, + {"operation": "remove", "path": "/status/detail"}, + ], + }, + { + "tool_name": "status_lookup", + "match_arguments": {}, + "operations": [{"operation": "add", "path": "/conditioned", "value": True}], + }, + ], + } + return { + "condition_id": "boundary-condition", + "fixture_id": "base-tool-fixture", + "fragment_id": "conditioned-tool-fragment", + "intensity": intensity, + "strengths": strengths, + } + + +def _document_text(fixture_set: Mapping[str, Any], document_id: str) -> str: + documents = cast(Sequence[Mapping[str, Any]], fixture_set["documents"]) + return cast(str, next(item["text"] for item in documents if item["id"] == document_id)) diff --git a/scripts/datagen/tests/test_corpus_pipeline.py b/scripts/datagen/tests/test_corpus_pipeline.py new file mode 100644 index 00000000000..acb4db638da --- /dev/null +++ b/scripts/datagen/tests/test_corpus_pipeline.py @@ -0,0 +1,80 @@ +import io +import json +import tarfile +from pathlib import Path + +from phoenix.experimental.datagen import load_corpus +from scripts.datagen.corpus import command as corpus_command +from scripts.datagen.publish import command as publish_command + + +def test_package_and_prepare_publication(tmp_path: Path) -> None: + source = Path(__file__).parent / "fixtures" / "fragment_bank" + archive = tmp_path / "corpus.tar.gz" + package_output = io.StringIO() + + assert ( + corpus_command( + [str(source), "--archive", str(archive)], + stdout=package_output, + ) + == 0 + ) + package = json.loads(package_output.getvalue()) + assert { + key: package[key] + for key in ( + "span_count", + "span_kind_counts", + "span_kind_shares", + "spans_per_trace", + "tool_span_count", + "tool_span_share", + "llm_turns_by_session", + "llm_turns_per_session", + ) + } == { + "span_count": 4, + "span_kind_counts": {"CHAIN": 1, "LLM": 3, "TOOL": 0}, + "span_kind_shares": {"CHAIN": 0.25, "LLM": 0.75, "TOOL": 0.0}, + "spans_per_trace": {"min": 1, "median": 1, "mean": 4 / 3, "max": 2}, + "tool_span_count": 0, + "tool_span_share": 0.0, + "llm_turns_by_session": {"session-a": 2, "session-b": 1}, + "llm_turns_per_session": {"min": 1, "median": 1.5, "mean": 1.5, "max": 2}, + } + + with tarfile.open(archive, "r:gz") as contents: + assert [member.name for member in contents.getmembers()] == [ + "fragments.jsonl", + "traces.jsonl", + ] + fragment_rows = contents.extractfile("fragments.jsonl") + assert fragment_rows is not None + assert all( + set(json.loads(line)) == {"fragment_id", "archetype", "domain", "trace_ids"} + for line in fragment_rows.read().decode().splitlines() + ) + + publication_output = io.StringIO() + publication_dir = tmp_path / "publication" + assert ( + publish_command( + [ + "prepare-archive", + "--archive", + str(archive), + "--output-dir", + str(publication_dir), + ], + stdout=publication_output, + ) + == 0 + ) + publication = json.loads(publication_output.getvalue()) + corpus = load_corpus(archive) + pointer = json.loads((publication_dir / "corpus.json").read_text()) + + assert package["sha256"] == publication["sha256"] == pointer["sha256"] + assert package["fragment_count"] == len(corpus.fragments) == 2 + assert publication["archetypes"] == ["plain_chat", "rag"] diff --git a/scripts/datagen/tests/test_fetcher.py b/scripts/datagen/tests/test_fetcher.py new file mode 100644 index 00000000000..aed7f508d15 --- /dev/null +++ b/scripts/datagen/tests/test_fetcher.py @@ -0,0 +1,58 @@ +import json +import shutil +from hashlib import sha256 +from pathlib import Path + +from phoenix.experimental.datagen import load_corpus +from phoenix.experimental.datagen.fetcher import fetch_corpus +from scripts.datagen.corpus import package_corpus + + +def test_fetch_corpus_caches_digest_addressed_archive(tmp_path: Path) -> None: + archive = _build_archive(tmp_path) + pointer = _write_pointer(tmp_path, archive) + downloads = 0 + + def download(_url: str, destination: Path) -> None: + nonlocal downloads + downloads += 1 + shutil.copyfile(archive, destination) + + cached = fetch_corpus( + cache_dir=tmp_path / "cache", + pointer_path=pointer, + downloader=download, + ) + corpus = load_corpus(cached) + cached_again = fetch_corpus( + cache_dir=tmp_path / "cache", + pointer_path=pointer, + downloader=download, + ) + + assert cached_again == cached + assert cached.name == sha256(archive.read_bytes()).hexdigest() + assert cached.read_bytes() == archive.read_bytes() + assert len(corpus.fragments) == 2 + assert downloads == 1 + + +def _build_archive(tmp_path: Path) -> Path: + source = Path(__file__).parent / "fixtures" / "fragment_bank" + archive = tmp_path / "corpus.tar.gz" + package_corpus(source, archive) + return archive + + +def _write_pointer(tmp_path: Path, archive: Path) -> Path: + pointer = tmp_path / "corpus.json" + pointer.write_text( + json.dumps( + { + "schema_version": 2, + "url": "https://assets.example/datagen/corpus.tar.gz", + "sha256": sha256(archive.read_bytes()).hexdigest(), + } + ) + ) + return pointer diff --git a/scripts/datagen/tool_agent.py b/scripts/datagen/tool_agent.py new file mode 100644 index 00000000000..a398c7f2b3a --- /dev/null +++ b/scripts/datagen/tool_agent.py @@ -0,0 +1,456 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "httpx==0.28.1", +# "langchain-core==1.5.6", +# "langchain-openai==1.5.1", +# "openai==3.2.0", +# "openinference-instrumentation==0.1.57", +# "openinference-instrumentation-langchain==0.1.70", +# "opentelemetry-exporter-otlp-proto-common==1.44.0", +# "opentelemetry-sdk==1.44.0", +# "protobuf==7.35.1", +# ] +# /// +"""Record fixed tool-agent fixtures through LangChain callbacks.""" + +from __future__ import annotations + +import argparse +import json +import os +import random +from collections.abc import Mapping, Sequence +from dataclasses import replace +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, cast + +from langchain_core.messages import AIMessage, BaseMessage, ToolMessage +from langchain_core.tools import BaseTool, StructuredTool +from langchain_openai import ChatOpenAI +from openinference.instrumentation import ( + OITracer, + TraceConfig, + get_attributes_from_context, + using_session, +) +from openinference.instrumentation.langchain import LangChainInstrumentor +from openinference.semconv.trace import OpenInferenceMimeTypeValues, OpenInferenceSpanKindValues +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +if TYPE_CHECKING or __package__: + from scripts.datagen.conditions import materialize_condition + from scripts.datagen.fake_tools import LocalTools, ToolError, local_tools + from scripts.datagen.mock_openai_provider import ScriptedOpenAIProvider + from scripts.datagen.recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + live_model_options, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + ) +else: + from conditions import materialize_condition + from fake_tools import LocalTools, ToolError, local_tools + from mock_openai_provider import ScriptedOpenAIProvider + from recording import ( + RecorderFixture, + SpanCaptureExporter, + append_spans, + fixtures_for, + live_model_options, + prepare_recording, + record_fixture, + resolve_live_model, + trace_ids, + ) + +MAX_TOOL_CALLS = 24 +Provider = Literal["scripted", "live"] +_ROOT_SPAN_NAMES = { + "coding_agent": "resolve_engineering_task", + "customer_support": "handle_support_request", + "data_analyst": "answer_analytics_request", +} + + +class OpenInferenceContextSpanProcessor(SpanProcessor): + """Apply the active session to spans started by LangChain callbacks.""" + + def on_start(self, span: Span, parent_context: Any = None) -> None: + span.set_attributes(dict(get_attributes_from_context())) + + def on_end(self, span: ReadableSpan) -> None: + pass + + def shutdown(self) -> None: + pass + + +class ToolAgentRecorder: + def __init__( + self, + model: ChatOpenAI, + tools: LocalTools, + exporter: SpanCaptureExporter, + tracer: OITracer, + *, + require_terminal_answer: bool, + ) -> None: + self._model = model + self._tools = tools + self._exporter = exporter + self._tracer = tracer + self._require_terminal_answer = require_terminal_answer + + def record(self, fixture: RecorderFixture, traces_path: Path) -> tuple[str, ...]: + prompt = fixture.inputs.get("prompt") + if not isinstance(prompt, str): + raise ValueError(f"fixture {fixture.fragment_id!r} has no tool-agent prompt") + tools = _bound_tools(self._tools) + model = self._model.bind_tools(tools) + checkpoint = self._exporter.checkpoint() + + def run_agent(inputs: Mapping[str, Any]) -> list[BaseMessage]: + messages: list[BaseMessage] = list(cast(Sequence[BaseMessage], inputs["messages"])) + with self._tracer.start_as_current_span( + _ROOT_SPAN_NAMES[fixture.domain], + openinference_span_kind=OpenInferenceSpanKindValues.AGENT, + ) as root_span: + root_span.set_input(prompt, mime_type=OpenInferenceMimeTypeValues.TEXT.value) + with self._tracer.start_as_current_span( + "triage", + openinference_span_kind=OpenInferenceSpanKindValues.CHAIN, + ): + reply = model.invoke(messages) + messages.append(reply) + if not reply.tool_calls: + root_span.set_output( + _message_text(reply), + mime_type=OpenInferenceMimeTypeValues.TEXT.value, + ) + return messages + with self._tracer.start_as_current_span( + "investigate", + openinference_span_kind=OpenInferenceSpanKindValues.CHAIN, + ): + for _ in range(MAX_TOOL_CALLS + 1): + for call in reply.tool_calls: + tool = _tool_by_name(tools, call["name"]) + try: + result = tool.invoke(call["args"]) + except ToolError as error: + messages.append( + ToolMessage( + content=json.dumps({"error": str(error)}), + tool_call_id=call["id"], + name=call["name"], + status="error", + ) + ) + else: + messages.append( + ToolMessage( + content=json.dumps( + result, sort_keys=True, separators=(",", ":") + ), + tool_call_id=call["id"], + name=call["name"], + ) + ) + reply = model.invoke(messages) + messages.append(reply) + if not reply.tool_calls: + root_span.set_output( + _message_text(reply), + mime_type=OpenInferenceMimeTypeValues.TEXT.value, + ) + return messages + raise RuntimeError(f"fixture {fixture.fragment_id!r} exceeded its tool-call limit") + + try: + with using_session(fixture.fragment_id): + result = run_agent({"messages": [{"role": "user", "content": prompt}]}) + except Exception: + if self._require_terminal_answer: + raise + result = [] + finally: + spans = self._exporter.spans_since(checkpoint) + if spans: + append_spans(traces_path, spans) + if self._require_terminal_answer and ( + not result or not isinstance(result[-1], AIMessage) or not result[-1].content + ): + raise RuntimeError(f"fixture {fixture.fragment_id!r} did not finish with an answer") + return trace_ids(spans) + + +def _bound_tools(tools: LocalTools) -> tuple[StructuredTool, ...]: + result = [] + for schema in tools.schemas: + function = cast(Mapping[str, Any], schema["function"]) + name = cast(str, function["name"]) + + def invoke(_name: str = name, **arguments: Any) -> Mapping[str, Any]: + return tools.invoke(_name, arguments) + + result.append( + StructuredTool.from_function( + func=invoke, + name=name, + description=cast(str, function["description"]), + args_schema=dict(cast(Mapping[str, Any], function["parameters"])), + infer_schema=False, + ) + ) + return tuple(result) + + +def _tool_by_name(tools: Sequence[BaseTool], name: str) -> BaseTool: + try: + return next(tool for tool in tools if tool.name == name) + except StopIteration as error: + raise RuntimeError(f"scripted model requested unknown tool {name!r}") from error + + +def _message_text(message: AIMessage) -> str: + if isinstance(message.content, str): + return message.content + return "\n".join( + str(block.get("text", "")) if isinstance(block, Mapping) else str(block) + for block in message.content + ) + + +def _with_prompt_variant(fixture: RecorderFixture, rng: random.Random) -> RecorderFixture: + """Pick one authored phrasing of the fixture's opening prompt for this run.""" + variants = fixture.inputs.get("prompt_variants") + prompt = fixture.inputs.get("prompt") + if not isinstance(variants, list) or not isinstance(prompt, str): + return fixture + phrasings = [prompt, *(v for v in variants if isinstance(v, str))] + return replace(fixture, inputs={**fixture.inputs, "prompt": rng.choice(phrasings)}) + + +def record( + output_dir: Path, + *, + fixtures: Sequence[RecorderFixture] | None = None, + condition: str | None = None, + append: bool = False, + provider: Provider = "scripted", + model: str | None = None, + live_model: ChatOpenAI | None = None, +) -> tuple[dict[str, Any], ...]: + """Record every selected tool-agent fixture into a corpus directory.""" + model = resolve_live_model(model) + if provider not in ("scripted", "live"): + raise ValueError(f"unknown tool-agent provider {provider!r}") + if condition is not None and fixtures is not None: + raise ValueError("condition and fixtures cannot be selected together") + if provider == "live" and not model: + raise ValueError("live tool-agent recording requires an explicit model") + if provider == "live" and live_model is None: + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + raise ValueError("live tool-agent recording requires OPENAI_API_KEY") + model_args: dict[str, Any] = { + "model": model, + "api_key": api_key, + "max_retries": 0, + } + model_args.update(live_model_options(model)) + if base_url := os.environ.get("OPENAI_BASE_URL"): + model_args["base_url"] = base_url + live_model = ChatOpenAI(**model_args) + conditioned_tools: LocalTools | None = None + if condition is not None: + conditioned = materialize_condition(condition) + if conditioned.fixture.archetype != "tool_agent": + raise ValueError(f"condition {condition!r} does not select a tool-agent fixture") + selected_fixtures: Sequence[RecorderFixture] = (conditioned.fixture,) + conditioned_tools = local_tools( + conditioned.fixture.domain, + fixture_set=conditioned.tool_fixture_set, + result_overlays=conditioned.tool_result_overlays, + ) + else: + selected_fixtures = fixtures_for("tool_agent", fixtures=fixtures) + if provider == "scripted" and fixtures is None: + # Auto-selection in scripted mode keeps only fixtures that have a + # deterministic episode; the rest record live. + selected_fixtures = tuple( + fixture + for fixture in selected_fixtures + if fixture.domain != "coding_agent" + or fixture.fragment_id in _SCRIPTED_CODING_EPISODES + ) + prepare_recording(output_dir, append=append) + exporter = SpanCaptureExporter() + tracer_provider = TracerProvider( + resource=Resource.create({"service.name": "datagen.tool_agent"}) + ) + tracer_provider.add_span_processor(OpenInferenceContextSpanProcessor()) + tracer_provider.add_span_processor(SimpleSpanProcessor(cast(Any, exporter))) + tracer = OITracer(tracer_provider.get_tracer(__name__), TraceConfig()) + instrumentor = LangChainInstrumentor() + instrumentor.instrument(tracer_provider=tracer_provider) + fragments = [] + variant_rng = random.Random() + try: + for fixture in selected_fixtures: + if provider == "live": + fixture = _with_prompt_variant(fixture, variant_rng) + if provider == "scripted": + scripted = ScriptedOpenAIProvider(_responses_for(fixture)) + selected_model = ChatOpenAI( + model="datagen-scripted", + api_key="datagen-dummy-key", + base_url="https://datagen.test/v1", + http_client=scripted.http_client(), + max_retries=0, + temperature=0, + ) + else: + selected_model = cast(ChatOpenAI, live_model) + recorder = ToolAgentRecorder( + selected_model, + conditioned_tools or local_tools(fixture.domain), + exporter, + tracer, + require_terminal_answer=provider == "scripted", + ) + fragments.append(record_fixture(fixture, output_dir, recorder.record)) + finally: + instrumentor.uninstrument() + tracer_provider.shutdown() + return tuple(fragments) + + +# Coding fixtures with a deterministic scripted episode; the rest are +# recorded live, where the model chooses its own tool calls. +_SCRIPTED_CODING_EPISODES = frozenset({"coding-router-api-tools", "coding-retry-policy-tools"}) + + +def _responses_for(fixture: RecorderFixture) -> tuple[dict[str, Any], ...]: + prompt = str(fixture.inputs.get("prompt", "")) + if "calculate" in prompt: + expression = "42.25 * 2" if "42.25" in prompt else "125000 - 8500" + calls = ( + {"name": "document_search", "arguments": {"query": prompt, "limit": 1}}, + {"name": "safe_arithmetic", "arguments": {"expression": expression}}, + ) + elif fixture.fragment_id == "coding-router-api-tools": + calls = ( + {"name": "repository_search", "arguments": {"query": "Router"}}, + {"name": "repository_search", "arguments": {"query": "Router"}}, + {"name": "record_lookup", "arguments": {"record_id": "issue-204"}}, + {"name": "read_file", "arguments": {"path": "README.md"}}, + {"name": "read_file", "arguments": {"path": "README.md"}}, + {"name": "read_file", "arguments": {"path": "README.md"}}, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/router.py"}, + }, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/router.py"}, + }, + {"name": "run_tests", "arguments": {"test": "tests/test_readme.py"}}, + { + "name": "edit_file", + "arguments": { + "path": "README.md", + "old": "Router.dispatch", + "new": "Router.route", + }, + }, + {"name": "read_file", "arguments": {"path": "README.md"}}, + {"name": "run_tests", "arguments": {"test": "tests/test_readme.py"}}, + ) + elif fixture.fragment_id == "coding-retry-policy-tools": + calls = ( + {"name": "repository_search", "arguments": {"query": "retry"}}, + {"name": "repository_search", "arguments": {"query": "retry"}}, + {"name": "record_lookup", "arguments": {"record_id": "issue-219"}}, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/scheduler.py"}, + }, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/scheduler.py"}, + }, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/scheduler.py"}, + }, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/retry.py"}, + }, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/retry.py"}, + }, + {"name": "run_tests", "arguments": {"test": "tests/test_retry.py"}}, + { + "name": "edit_file", + "arguments": { + "path": "src/relaycache/retry.py", + "old": "return attempt + 1", + "new": "return attempt", + }, + }, + { + "name": "read_file", + "arguments": {"path": "src/relaycache/retry.py"}, + }, + {"name": "run_tests", "arguments": {"test": "tests/test_retry.py"}}, + ) + elif fixture.domain == "coding_agent": + raise ValueError(f"fixture {fixture.fragment_id!r} has no scripted coding episode") + else: + identifier = next(value for value in ("order-1001", "warehouse-east") if value in prompt) + calls = ( + {"name": "record_lookup", "arguments": {"record_id": identifier}}, + {"name": "status_lookup", "arguments": {"status_id": identifier}}, + ) + if fixture.domain == "coding_agent": + answer = ( + "I reproduced the failure, updated the affected file, and confirmed the focused test " + "passes on the rerun." + ) + else: + answer = "The local records and policy data support the requested next step." + return tuple({"tool_call": call} for call in calls) + ({"content": answer},) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--condition") + parser.add_argument("--append", action="store_true") + parser.add_argument("--provider", choices=("scripted", "live"), default="scripted") + parser.add_argument("--model") + args = parser.parse_args() + fragments = record( + args.output_dir, + condition=args.condition, + append=args.append, + provider=args.provider, + model=args.model, + ) + print(f"Recorded {len(fragments)} tool-agent fragments in {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/datagen/tool_fixtures.json b/scripts/datagen/tool_fixtures.json new file mode 100644 index 00000000000..43d1320b64a --- /dev/null +++ b/scripts/datagen/tool_fixtures.json @@ -0,0 +1,250 @@ +{ + "customer_support": { + "name": "customer_support", + "documents": [ + { + "id": "returns-policy", + "title": "Returns and refunds", + "text": "Unused merchandise may be returned within 30 calendar days after delivery. Warehouse inspection normally takes three business days, followed by three to seven business days for bank posting." + }, + { + "id": "delivery-guide", + "title": "Delivery status", + "text": "Standard delivery takes four to six business days. Express delivery takes one to two business days. A carrier trace is available after two business days without movement." + } + ], + "records": [ + { + "id": "order-1001", + "customer": "Avery Chen", + "total": 84.5, + "currency": "USD" + }, + { + "id": "order-1002", + "customer": "Sam Rivera", + "total": 129.0, + "currency": "USD" + } + ], + "statuses": [ + { + "id": "order-1001", + "state": "in_transit", + "detail": "Carrier scan unchanged for two business days" + }, + { + "id": "order-1002", + "state": "processing", + "detail": "Preparing for shipment" + } + ] + }, + "data_analyst": { + "name": "data_analyst", + "documents": [ + { + "id": "metric-definitions", + "title": "Governed analytics metrics", + "text": "Net revenue is captured item revenue after discounts, less refunds recognized on the refund date, excluding tax and shipping. Sales bookings are not recognized revenue." + }, + { + "id": "timezone-and-units", + "title": "Timezones and units", + "text": "Daily warehouse metrics use each facility's local date. Monetary records use major currency units unless the source adapter is explicitly marked otherwise." + } + ], + "records": [ + { + "id": "warehouse-east", + "timezone": "America/New_York", + "currency": "USD" + }, + { + "id": "warehouse-west", + "timezone": "America/Los_Angeles", + "currency": "USD" + } + ], + "statuses": [ + { + "id": "warehouse-east", + "state": "ready", + "detail": "Previous local business day is complete" + }, + { + "id": "warehouse-west", + "state": "delayed", + "detail": "Inventory snapshot is awaiting reconciliation" + } + ] + }, + "coding_agent": { + "name": "coding_agent", + "documents": [ + { + "id": "router-api", + "title": "RelayCache routing API", + "text": "Router.route is the current public coroutine. Router.dispatch is a deprecated compatibility alias retained through the next minor release." + }, + { + "id": "retry-ownership", + "title": "Retry ownership", + "text": "The scheduler owns attempt timing and retry.py owns pure retry-budget and delay calculations. Broker adapters own transport and acknowledgement translation." + }, + { + "id": "metrics-conventions", + "title": "Latency metrics conventions", + "text": "Latency histograms must cover the slow tail: buckets end at 5.0 seconds. Emit one histogram per route with route name as a label." + }, + { + "id": "broker-ack-contract", + "title": "Broker acknowledgement contract", + "text": "A delivery is acknowledged exactly once, after handler completion. Adapters must not acknowledge on redelivery of an already-processed message." + }, + { + "id": "release-process", + "title": "Release process", + "text": "Every closed issue that changed behavior gets a CHANGELOG entry under the upcoming version heading before the release branch is cut." + }, + { + "id": "config-reference", + "title": "Configuration reference", + "text": "RELAYCACHE_TIMEOUT accepts seconds as a decimal number. Values are read once at startup by config.load_settings." + } + ], + "records": [ + { + "id": "issue-204", + "title": "README uses deprecated Router.dispatch", + "state": "open" + }, + { + "id": "issue-219", + "title": "Retry backoff starts one attempt early", + "state": "reproduced" + }, + { + "id": "issue-231", + "title": "Latency histogram misses the slow tail", + "state": "open" + }, + { + "id": "issue-242", + "title": "RELAYCACHE_TIMEOUT drops decimal values", + "state": "open" + }, + { + "id": "issue-250", + "title": "Consumers observe duplicate acknowledgements", + "state": "reproduced" + }, + { + "id": "issue-263", + "title": "CHANGELOG missing entries for 1.4", + "state": "open" + } + ], + "statuses": [ + { + "id": "issue-204", + "state": "ready", + "detail": "Documentation-only correction" + }, + { + "id": "issue-219", + "state": "investigating", + "detail": "Focused scheduler test is failing" + }, + { + "id": "issue-231", + "state": "ready", + "detail": "Bucket list fix in metrics.py" + }, + { + "id": "issue-242", + "state": "ready", + "detail": "Parse timeout as float in config.py" + }, + { + "id": "issue-250", + "state": "investigating", + "detail": "Redelivery path acknowledges twice" + }, + { + "id": "issue-263", + "state": "ready", + "detail": "Docs-only release note additions" + } + ], + "repository": { + "search_page_size": 2, + "read_chunk_lines": 2, + "files": [ + { + "path": "README.md", + "content": "# RelayCache\n\nRoute requests with `await Router.dispatch(request)`.\nThe scheduler owns retry timing and broker acknowledgements stay in adapters.\nSee `src/relaycache/router.py` for the public routing API." + }, + { + "path": "src/relaycache/router.py", + "content": "class Router:\n async def route(self, request):\n return await self._send(request)\n\n async def dispatch(self, request):\n return await self.route(request)" + }, + { + "path": "src/relaycache/retry.py", + "content": "def retry_index(attempt: int) -> int:\n if attempt < 0:\n raise ValueError(\"attempt must be non-negative\")\n return attempt + 1\n\ndef retry_delay(attempt: int) -> float:\n return 0.25 * (2 ** retry_index(attempt))" + }, + { + "path": "src/relaycache/scheduler.py", + "content": "from .retry import retry_delay\n\nasync def schedule_retry(job, attempt):\n delay = retry_delay(attempt)\n await job.defer(delay)\n return delay" + }, + { + "path": "src/relaycache/metrics.py", + "content": "LATENCY_BUCKETS = [0.1, 0.5, 1.0]\n\ndef observe_latency(histogram, route, seconds):\n histogram.labels(route=route).observe(seconds)\n" + }, + { + "path": "src/relaycache/config.py", + "content": "import os\n\ndef load_settings():\n timeout = int(os.environ.get(\"RELAYCACHE_TIMEOUT\", \"30\"))\n return {\"timeout\": timeout}\n" + }, + { + "path": "src/relaycache/broker.py", + "content": "class BrokerAdapter:\n def handle(self, message):\n receipt = self._process(message)\n self.ack(message)\n if message.redelivered:\n self.ack(message)\n return receipt\n" + }, + { + "path": "docs/CHANGELOG.md", + "content": "# Changelog\n\n## 1.4 (upcoming)\n\n- Router.route is the public routing coroutine.\n\n## 1.3\n\n- Added broker adapter acknowledgement translation.\n" + } + ], + "transient_failures": { + "read_file:README.md": 1, + "read_file:src/relaycache/scheduler.py": 1 + }, + "tests": [ + { + "name": "tests/test_readme.py", + "path": "README.md", + "contains": "Router.route" + }, + { + "name": "tests/test_retry.py", + "path": "src/relaycache/retry.py", + "contains": "return attempt\n" + }, + { + "name": "tests/test_metrics.py", + "path": "src/relaycache/metrics.py", + "contains": "5.0" + }, + { + "name": "tests/test_config.py", + "path": "src/relaycache/config.py", + "contains": "float(" + }, + { + "name": "tests/test_broker.py", + "path": "src/relaycache/broker.py", + "contains": "return receipt" + } + ] + } + } +}