From d82ada168cfe14aa272ceda8ed1490fa598671be Mon Sep 17 00:00:00 2001 From: Eric Hills <53243273+ebhills@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:52:27 -0500 Subject: [PATCH 1/4] Add extract.ai web search provenance --- .../extract_ai_recipe_examples.ipynb | 96 +++++++- tests/recipes/wrangles/test_extract.py | 64 ++++++ tests/recipes/wrangles/test_main.py | 38 +++ tests/test_openai_extract_ai.py | 180 +++++++++++++++ wrangles/extract.py | 58 ++++- wrangles/openai_responses.py | 133 ++++++++++- wrangles/recipe_wrangles/extract.py | 216 ++++++++++++++---- 7 files changed, 719 insertions(+), 66 deletions(-) diff --git a/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb b/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb index 9a4a01b56..81bb4c465 100644 --- a/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb +++ b/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb @@ -2019,6 +2019,52 @@ "result" ] }, + { + "cell_type": "markdown", + "id": "b18c5f72", + "metadata": {}, + "source": [ + "### Whole-record examples: `name`, `input`, and `output`\n", + "\n", + "Top-level `examples` teach a complete input-to-output mapping. Keep the\n", + "source under `input`, the desired record under `output`, and optionally add\n", + "a descriptive `name`. Sparse example outputs are completed with `null` for\n", + "fields you omit. These are different from `examples` nested under one output\n", + "field, which guide only that field." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9f3a10de", + "metadata": {}, + "outputs": [], + "source": [ + "recipe_record_examples = '''\n", + "wrangles:\n", + " - extract.ai:\n", + " api_key: ${OPENAI_API_KEY}\n", + " output:\n", + " best_match:\n", + " type: string\n", + " description: Best catalog match for the input item\n", + " certainty:\n", + " type: string\n", + " enum: [high, medium, low]\n", + " examples:\n", + " - name: exact manufacturer part match\n", + " input:\n", + " description: ACME valve V-100\n", + " output:\n", + " best_match: ACME V-100 valve\n", + " certainty: high\n", + "'''\n", + "\n", + "df = pd.DataFrame({\"description\": [\"ACME valve model V-100\"]})\n", + "result = recipe.run(recipe_record_examples, variables=variables, dataframe=df)\n", + "result" + ] + }, { "cell_type": "markdown", "id": "e6dc6eb6", @@ -2887,12 +2933,56 @@ "server.shutdown()" ] }, + { + "cell_type": "markdown", + "id": "57f0ca93", + "metadata": {}, + "source": [ + "## 14. Web search and source provenance\n", + "\n", + "Set `web_search: true` to make the native OpenAI Responses web-search tool\n", + "available to the model. The model decides whether a row needs a search.\n", + "Every processed row receives a separate `web_search_sources` column containing\n", + "a deduplicated list of `{title, url}` objects taken from Responses metadata.\n", + "The list is empty when no search/source was used. The reserved column name is\n", + "automatic; do not also define it under `output`.\n", + "\n", + "Web results are cached with the extracted value. Use `cache: false` when the\n", + "recipe must retrieve fresh web information. Web search requires the default\n", + "`responses` protocol and is not available with legacy `chat_completions`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d8641ec7", + "metadata": {}, + "outputs": [], + "source": [ + "recipe_web_search = '''\n", + "wrangles:\n", + " - extract.ai:\n", + " input: company\n", + " api_key: ${OPENAI_API_KEY}\n", + " web_search: true\n", + " cache: false\n", + " output:\n", + " latest_news_title:\n", + " type: string\n", + " description: Title of the company's most recent official news post\n", + "'''\n", + "\n", + "df = pd.DataFrame({\"company\": [\"OpenAI\"]})\n", + "result = recipe.run(recipe_web_search, variables=variables, dataframe=df)\n", + "result[[\"latest_news_title\", \"web_search_sources\"]]" + ] + }, { "cell_type": "markdown", "id": "0f2c87d0", "metadata": {}, "source": [ - "## 14. Quick reference\n", + "## 15. Quick reference\n", "\n", "| What you want | Recipe snippet |\n", "|---|---|\n", @@ -2902,11 +2992,13 @@ "| Constrained values | Add `enum: [...]` to a field |\n", "| Numeric / list result | `type: integer` / `type: array` |\n", "| Style hints | Add `examples: [...]` to a field |\n", + "| Whole-record examples | `examples: [{name: ..., input: ..., output: ...}]` |\n", + "| Web search + sources | `web_search: true` adds `web_search_sources` |\n", "| Reasoning models | `model: gpt-5-mini`, `reasoning: {effort: low}`, `verbosity: low` |\n", "| Parallelism | `threads: ` |\n", "| Resilience | `timeout: `, `retries: ` |\n", "| Legacy endpoint | `url: https://api.openai.com/v1/chat/completions` |\n", - "| WrangleWorks model | `model_id: ` instead of `api_key`/`model` |" + "| Saved WrangleWorks definition | `model_id: ` instead of an `output` schema; `api_key` is still required |" ] } ], diff --git a/tests/recipes/wrangles/test_extract.py b/tests/recipes/wrangles/test_extract.py index ba1b702e6..d6a6f5cd4 100644 --- a/tests/recipes/wrangles/test_extract.py +++ b/tests/recipes/wrangles/test_extract.py @@ -5,6 +5,70 @@ from unittest.mock import patch +class TestExtractAIWebSearch: + @patch("wrangles.recipe_wrangles.extract._extract.ai") + def test_web_search_writes_a_dedicated_sources_column(self, extract_ai): + extract_ai.return_value = [ + { + "manufacturer": "Acme", + "web_search_sources": [{ + "title": "Acme product page", + "url": "https://example.com/acme", + }], + }, + { + "manufacturer": "Contoso", + "web_search_sources": [], + }, + ] + data = pd.DataFrame({"description": ["Acme part", "Contoso part"]}) + recipe = """ + wrangles: + - extract.ai: + input: description + api_key: test-key + web_search: true + output: + manufacturer: + type: string + description: Manufacturer name + """ + + result = wrangles.recipe.run(recipe, dataframe=data) + + assert result["manufacturer"].tolist() == ["Acme", "Contoso"] + assert result["web_search_sources"].tolist() == [ + [{ + "title": "Acme product page", + "url": "https://example.com/acme", + }], + [], + ] + assert extract_ai.call_args.kwargs["web_search"] is True + + @patch("wrangles.recipe_wrangles.extract._extract.ai") + def test_web_search_rejects_an_existing_sources_column(self, extract_ai): + data = pd.DataFrame({ + "description": ["Acme part"], + "web_search_sources": [[]], + }) + recipe = """ + wrangles: + - extract.ai: + input: description + api_key: test-key + web_search: true + output: + manufacturer: + type: string + """ + + with pytest.raises(ValueError, match="is reserved"): + wrangles.recipe.run(recipe, dataframe=data) + + extract_ai.assert_not_called() + + class TestExtractAddress: """ Test extract.address diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index 8e04b070d..e708adc31 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -9339,6 +9339,44 @@ def test_all_wrangle_docstrings_parse_as_yaml(self): assert not failures, 'Wrangle schema docstring YAML parse failures:\n' + '\n'.join(failures) + def test_extract_ai_schema_documents_public_parameters_and_example_shape(self): + import inspect + import yaml + + method = wrangles.recipe._recipe_wrangles.extract.ai + schema = yaml.safe_load(method.__doc__) + properties = schema["properties"] + parameters = { + name + for name in inspect.signature(method).parameters + if name not in {"df", "kwargs"} + } + + assert parameters <= properties.keys() + assert all( + isinstance(definition.get("description"), str) + and definition["description"].strip() + for definition in properties.values() + ) + assert properties["web_search"]["type"] == "boolean" + assert "web_search_sources" in properties["web_search"]["description"] + + examples = properties["examples"] + assert examples["required"] == ["input", "output"] + assert examples["items"]["required"] == ["input", "output"] + assert set(examples["items"]["properties"]) == {"name", "input", "output"} + assert all( + definition["description"].strip() + for definition in examples["items"]["properties"].values() + ) + + field_schema = next(iter(properties["output"]["patternProperties"].values())) + assert all( + definition["description"].strip() + for definition in field_schema["properties"].values() + ) + assert "effort" in properties["reasoning"]["properties"] + def test_extract_codes_schema_matches_microservice_params(self): import yaml diff --git a/tests/test_openai_extract_ai.py b/tests/test_openai_extract_ai.py index 1c9768db6..abdf23eec 100644 --- a/tests/test_openai_extract_ai.py +++ b/tests/test_openai_extract_ai.py @@ -73,6 +73,8 @@ def post(**kwargs): assert payload["text"]["verbosity"] == "low" assert payload["text"]["format"]["strict"] is True assert payload["store"] is False + assert "tools" not in payload + assert "include" not in payload assert calls[0]["timeout"] <= 12 assert "seed" not in payload assert "Ignored legacy OpenAI parameter 'seed'" in caplog.text @@ -82,6 +84,184 @@ def post(**kwargs): assert schema["additionalProperties"] is False +def test_extract_ai_web_search_returns_metadata_sources_and_caches_them(monkeypatch): + calls = [] + body = { + "output": [ + { + "type": "web_search_call", + "action": { + "type": "search", + "sources": [ + {"type": "url", "url": "https://example.com/a"}, + { + "type": "url", + "url": "https://example.com/b", + "title": "Source B", + }, + ], + }, + }, + { + "type": "message", + "content": [{ + "type": "output_text", + "text": '{"output":"Acme"}', + "annotations": [ + { + "type": "url_citation", + "url": "https://example.com/a", + "title": "Source A", + }, + { + "type": "url_citation", + "url": "https://example.com/c", + "title": "Source C", + }, + ], + }], + }, + ] + } + monkeypatch.setattr( + extract._openai_responses._requests, + "post", + lambda **kwargs: calls.append(kwargs) or _Response(body), + ) + arguments = { + "input": "Who makes this product?", + "api_key": "tenant-key", + "output": "Manufacturer name", + "web_search": True, + "threads": 1, + } + + first = extract.ai(**arguments) + second = extract.ai(**arguments) + + assert first == { + "output": "Acme", + "web_search_sources": [ + {"title": "Source A", "url": "https://example.com/a"}, + {"title": "Source B", "url": "https://example.com/b"}, + {"title": "Source C", "url": "https://example.com/c"}, + ], + } + assert second == first + assert len(calls) == 1 + payload = calls[0]["json"] + assert payload["tools"] == [{"type": "web_search"}] + assert payload["include"] == ["web_search_call.action.sources"] + assert payload["tool_choice"] == "auto" + assert "authorized evidence in addition to DATA" in payload["instructions"] + + +def test_extract_ai_web_search_preserves_expert_tool_settings(monkeypatch): + calls = [] + body = { + "output": [{ + "type": "message", + "content": [{ + "type": "output_text", + "text": '{"manufacturer":"Acme"}', + "annotations": [], + }], + }] + } + monkeypatch.setattr( + extract._openai_responses._requests, + "post", + lambda **kwargs: calls.append(kwargs) or _Response(body), + ) + + result = extract.ai( + "Acme part", + "key", + output={"manufacturer": {"type": "string"}}, + web_search=True, + tools=[{"type": "code_interpreter", "container": {"type": "auto"}}], + include=["reasoning.encrypted_content"], + tool_choice="required", + threads=1, + ) + + assert result == {"manufacturer": "Acme", "web_search_sources": []} + payload = calls[0]["json"] + assert payload["tools"] == [ + {"type": "code_interpreter", "container": {"type": "auto"}}, + {"type": "web_search"}, + ] + assert payload["include"] == [ + "reasoning.encrypted_content", + "web_search_call.action.sources", + ] + assert payload["tool_choice"] == "required" + + +def test_extract_ai_web_search_failure_still_returns_empty_sources(monkeypatch): + monkeypatch.setattr( + extract._openai_responses._requests, + "post", + lambda **kwargs: _Response({ + "output": [ + { + "type": "web_search_call", + "action": { + "type": "search", + "sources": [{"url": "https://example.com/failure-source"}], + }, + }, + { + "type": "message", + "content": [{"type": "output_text", "text": "not json"}], + }, + ] + }), + ) + + result = extract.ai( + "Acme part", + "key", + output="Manufacturer name", + web_search=True, + retries=0, + threads=1, + ) + + assert result["web_search_sources"] == [{ + "title": "", + "url": "https://example.com/failure-source", + }] + assert result["output"].startswith("Invalid structured response") + + +def test_extract_ai_web_search_validates_protocol_and_reserved_output(): + with pytest.raises(ValueError, match="only with protocol='responses'"): + extract.ai( + "Acme part", + "key", + output={"manufacturer": {"type": "string"}}, + web_search=True, + protocol="chat_completions", + ) + + with pytest.raises(ValueError, match="web_search must be true or false"): + extract.ai( + "Acme part", + "key", + output={"manufacturer": {"type": "string"}}, + web_search="yes", + ) + + with pytest.raises(ValueError, match="is reserved"): + extract.ai( + "Acme part", + "key", + output={"web search sources": {"type": "string"}}, + web_search=True, + ) + + def test_examples_are_stable_instructions_and_part_of_the_result_cache_key(monkeypatch): calls = [] body = { diff --git a/wrangles/extract.py b/wrangles/extract.py index 68de22ed0..e36267bf4 100644 --- a/wrangles/extract.py +++ b/wrangles/extract.py @@ -69,6 +69,28 @@ def _cacheable_ai_result(result) -> bool: return True +def _enable_responses_web_search(payload: dict) -> None: + """Add native web search without replacing expert Responses settings.""" + tools = payload.setdefault("tools", []) + if not isinstance(tools, list): + raise ValueError("OpenAI Responses 'tools' must be an array.") + if not any( + isinstance(tool, dict) + and tool.get("type") in {"web_search", "web_search_preview"} + for tool in tools + ): + tools.append({"type": "web_search"}) + + included = payload.setdefault("include", []) + if not isinstance(included, list): + raise ValueError("OpenAI Responses 'include' must be an array.") + source_include = "web_search_call.action.sources" + if source_include not in included: + included.append(source_include) + + payload.setdefault("tool_choice", "auto") + + def address( input: _Union[str, list], dataType: str, @@ -126,6 +148,7 @@ def ai( store: bool = None, cache: bool = None, cache_ttl: float = None, + web_search: bool = False, **kwargs ) -> _Union[dict, list]: """ @@ -140,17 +163,17 @@ def ai( :param input: A single value or list of values to extract information from. If a list is provided, \ each element will be analyzed individually and a list of equal length will be returned. - :param api_key: API Key + :param api_key: OpenAI API key. :param output: (Optional) This can be a string prompting the output, a JSON schema definition \ of the output requested or a dict of JSON schema definitions. :param model_id: (Optional) An extract.ai model ID containing a saved definition. Use this or output. \ - If both are provided, output that precedence over the definition from the model_id. + If both are provided, named output fields take precedence over matching saved fields. :param model: (Optional) The model to use for the extraction. :param threads: (Optional) Number of threads to use for parallel processing. :param timeout: (Optional) Timeout in seconds for each API call. :param retries: (Optional) Number of retries to attempt on failure. :param messages: (Optional) Overall prompts to pass additional instructions. - :param examples: (Optional) Holistic examples containing paired input and output values. + :param examples: (Optional) Holistic examples containing separate input and output values, plus an optional name. :param url: (Optional) Override the configured endpoint. :param strict: (Optional) Enable structured output strict mode. Dynamic object schemas \ automatically use non-strict mode and are validated locally. @@ -164,8 +187,11 @@ def ai( :param store: (Optional) Whether OpenAI may store Responses. Defaults to False. :param cache: (Optional) Use the bounded warm-instance result cache. Defaults to True. :param cache_ttl: (Optional) Override the result-cache TTL in seconds for this call. + :param web_search: (Optional) Enable native Responses web search. Each result then includes a + web_search_sources list containing source titles and URLs. Defaults to False. - :return: A scalar or list of extracted information. + :return: Extracted information. When web_search is true, returns a dictionary (or list of + dictionaries) containing web_search_sources, including for single-field output. """ policy = _ai_config.extract_ai() provider = str(provider or policy.get("provider", "openai")).strip().lower() @@ -184,6 +210,10 @@ def ai( else: protocol = policy.get("protocol", "responses") protocol = _normalize_ai_protocol(protocol) + if not isinstance(web_search, bool): + raise ValueError("web_search must be true or false.") + if web_search and protocol != "responses": + raise ValueError("web_search is supported only with protocol='responses'.") if url: if protocol == "responses" and "/chat/completions" in url: @@ -251,6 +281,14 @@ def ai( _needs_remap = compiled.needs_remap root_schema = compiled.root_schema example_guidance = _ai_definition.render_example_guidance(compiled) + if ( + web_search + and _openai_responses.WEB_SEARCH_SOURCES_KEY in compiled.output + ): + raise ValueError( + f"{_openai_responses.WEB_SEARCH_SOURCES_KEY!r} is reserved when " + "web_search is enabled. Choose a different output field name." + ) messages = [ { @@ -277,6 +315,12 @@ def ai( str(message.get("content", "")) for message in messages ) + if web_search: + instructions += "\n\n" + " ".join([ + "Web search is enabled for this call.", + "Information returned by the web search tool is authorized evidence in addition to DATA.", + "Use web search only when it helps answer the requested fields, and return null when neither DATA nor web evidence supports a field.", + ]) payload = { "model": model, @@ -292,6 +336,8 @@ def ai( "store": store, **_openai_responses.sanitize_request_params(kwargs), } + if web_search: + _enable_responses_web_search(payload) configured_reasoning = ( reasoning if reasoning is not None @@ -369,12 +415,12 @@ def ai( ] if input_was_scalar: - if output_generic_key: + if output_generic_key and not web_search: return results[0].get('output', 'Failed') else: return results[0] else: - if output_generic_key: + if output_generic_key and not web_search: return [x.get('output', 'Failed') for x in results] else: return results diff --git a/wrangles/openai_responses.py b/wrangles/openai_responses.py index 60afc8cc0..419b5ffbc 100644 --- a/wrangles/openai_responses.py +++ b/wrangles/openai_responses.py @@ -25,6 +25,7 @@ _LOG = _logging.getLogger(__name__) _LOCK = _threading.Lock() _SUCCESS_STATS = {} +WEB_SEARCH_SOURCES_KEY = "web_search_sources" _JSON_TYPE_MAP = { "string": str, "number": float, @@ -526,11 +527,101 @@ def format_input_data(data: _Any) -> str: return str(data) -def error_result(required_fields: list, message: str) -> dict: - return { +def _uses_web_search(payload: dict) -> bool: + tools = payload.get("tools", []) + if not isinstance(tools, list): + return False + return any( + isinstance(tool, dict) + and tool.get("type") in {"web_search", "web_search_preview"} + for tool in tools + ) + + +def extract_web_search_sources(response_json: dict) -> list: + """Return cited and consulted web URLs in stable response order.""" + annotation_titles = {} + annotations = [] + output = response_json.get("output", []) + if not isinstance(output, list): + output = [] + + for item in output: + if not isinstance(item, dict) or item.get("type") != "message": + continue + content_items = item.get("content", []) + if not isinstance(content_items, list): + continue + for content in content_items: + if not isinstance(content, dict): + continue + content_annotations = content.get("annotations", []) + if not isinstance(content_annotations, list): + continue + for annotation in content_annotations: + if ( + not isinstance(annotation, dict) + or annotation.get("type") != "url_citation" + ): + continue + url = annotation.get("url") + if not isinstance(url, str) or not url.strip(): + continue + url = url.strip() + title = annotation.get("title") + title = title.strip() if isinstance(title, str) else "" + annotations.append((url, title)) + if title and url not in annotation_titles: + annotation_titles[url] = title + + sources = [] + positions = {} + + def add_source(url, title=""): + if not isinstance(url, str) or not url.strip(): + return + url = url.strip() + title = title.strip() if isinstance(title, str) else "" + title = title or annotation_titles.get(url, "") + if url in positions: + existing = sources[positions[url]] + if not existing["title"] and title: + existing["title"] = title + return + positions[url] = len(sources) + sources.append({"title": title, "url": url}) + + for item in output: + if not isinstance(item, dict) or item.get("type") != "web_search_call": + continue + action = item.get("action") + if not isinstance(action, dict): + continue + action_sources = action.get("sources", []) + if isinstance(action_sources, list): + for source in action_sources: + if isinstance(source, dict): + add_source(source.get("url"), source.get("title", "")) + add_source(action.get("url"), action.get("title", "")) + + for url, title in annotations: + add_source(url, title) + + return sources + + +def error_result( + required_fields: list, + message: str, + include_web_search_sources: bool = False, +) -> dict: + result = { field: message for field in required_fields } + if include_web_search_sources: + result[WEB_SEARCH_SOURCES_KEY] = [] + return result def extract_response_text(response_json: dict) -> str: @@ -638,13 +729,24 @@ def call_structured( "content": f"DATA:\n{format_input_data(data)}", } ] + include_web_search_sources = _uses_web_search(request_payload) + + def failure(message: str, response_json: dict = None) -> dict: + result = error_result( + required_fields, + message, + include_web_search_sources=include_web_search_sources, + ) + if include_web_search_sources and isinstance(response_json, dict): + result[WEB_SEARCH_SOURCES_KEY] = extract_web_search_sources(response_json) + return result response = None backoff_time = 1 for attempt in range(retries + 1): remaining = _remaining_seconds(deadline_at) if remaining is not None and remaining <= 0: - return error_result(required_fields, "Deadline Exceeded") + return failure("Deadline Exceeded") request_timeout = timeout if remaining is not None: @@ -662,14 +764,15 @@ def call_structured( elapsed_seconds = _time.time() - started except _requests.exceptions.Timeout: if attempt >= retries: - return error_result(required_fields, "Timed Out") + return failure("Timed Out") except Exception as e: if attempt >= retries: - return error_result(required_fields, str(e)) + return failure(str(e)) if response is not None and response.ok: try: - output_text = extract_response_text(response.json()) + response_json = response.json() + output_text = extract_response_text(response_json) parsed = _json.loads(output_text) if not isinstance(parsed, dict): raise ValueError("Structured response was not a JSON object.") @@ -680,10 +783,18 @@ def call_structured( model=request_payload.get("model"), elapsed_seconds=elapsed_seconds, ) - return validate_structured_output(parsed, schema) + validated = validate_structured_output(parsed, schema) + if include_web_search_sources: + validated[WEB_SEARCH_SOURCES_KEY] = extract_web_search_sources( + response_json + ) + return validated except (_json.JSONDecodeError, _ValidationError, ValueError) as e: if attempt >= retries: - return error_result(required_fields, f"Invalid structured response: {e}") + return failure( + f"Invalid structured response: {e}", + response_json=response_json, + ) else: context = _response_context( response, @@ -700,7 +811,7 @@ def call_structured( raise ValueError("API Key provided is missing or invalid.") if attempt >= retries or not _should_retry(context): _log_api_error(context, final=True) - return error_result(required_fields, _error_message(context)) + return failure(_error_message(context)) _log_api_error(context, final=False) if response is not None and not response.ok: @@ -708,7 +819,7 @@ def call_structured( else: delay = _sleep_for_retry({}, backoff_time, deadline_at) if delay is None: - return error_result(required_fields, "Deadline Exceeded") + return failure("Deadline Exceeded") backoff_time *= 2 - return error_result(required_fields, "Failed") + return failure("Failed") diff --git a/wrangles/recipe_wrangles/extract.py b/wrangles/recipe_wrangles/extract.py index c69650907..c69dcaac2 100644 --- a/wrangles/recipe_wrangles/extract.py +++ b/wrangles/recipe_wrangles/extract.py @@ -26,6 +26,7 @@ "concatenate": "concatenate", "concat": "concatenate", } +_WEB_SEARCH_SOURCES_KEY = "web_search_sources" def _normalize_output_format(output_format, default): @@ -295,11 +296,14 @@ def ai( model_id: str = None, output_format: str = None, char: str = ", ", + web_search: bool = False, **kwargs ): """ type: object - description: Extract data using an AI model. + description: >- + Extract structured data from each input row using an AI model. Define + the desired fields with output, or reuse a saved definition with model_id. additionalProperties: false required: - api_key @@ -314,20 +318,26 @@ def ai( - string - integer - array - description: |- - Name or list of input columns to give to the AI - to use to determine the output. If not specified, all - columns will be used. + description: >- + Input column name, column index, or list of columns supplied together + as DATA for each row. If omitted, all dataframe columns are supplied. + items: + type: [string, integer] output: type: [object, string, array] - description: List and description of the output you want + description: >- + Desired extraction. Use an object keyed by output column name for + structured fields, a string for one prompted value, or an array of + field names/definitions. Each field may use the schema options below. patternProperties: "^[a-zA-Z0-9 _-]+$": type: [object, string] properties: type: type: string - description: The type of data you'd like the model to return. + description: >- + JSON data type required for this field. If omitted, common + scalar types are accepted. Fields allow null by default. enum: - string - number @@ -338,10 +348,14 @@ def ai( - array description: type: string - description: Description of the output you'd like the model to return. + description: >- + Plain-language definition of the value to extract, including + any selection, normalization, unit, or evidence rules. enum: type: array - description: List of possible values for the output. + description: >- + Allowed output values. The model must choose one of these + values; null is also allowed unless nullable is false. default: type: - string @@ -351,7 +365,10 @@ def ai( - "null" - object - array - description: A default value to return. + description: >- + JSON Schema annotation for a preferred default. extract.ai + does not substitute this value when evidence is missing; + describe fallback behavior explicitly or allow null. examples: type: - array @@ -362,27 +379,31 @@ def ai( - boolean - "null" description: >- - Provide typical output values, or paired objects with input - and output keys for field-specific examples. + Field-specific examples. Provide typical output values to + demonstrate style, or objects with separate input and output + keys to show how a source value maps to this field. properties: type: - object - array - string description: >- - Named properties for an object. A list or comma-separated - string creates fixed property names. + Child fields when type is object. Use an object to define a + schema for each child. A list or comma-separated string is a + shortcut that creates fixed child names. additionalProperties: type: - boolean - object description: >- - Set true, or provide a value schema, for a dynamic dictionary. - Dynamic dictionaries automatically use non-strict schema mode - while fixed portions of the output remain constrained. + Controls keys beyond properties when type is object. Set false + for fixed keys, true for arbitrary values, or provide one + schema applied to every dynamic value. Dynamic dictionaries + use non-strict provider mode plus local validation. items: type: object - description: Schema for each item returned by an array. + description: >- + Schema applied to every element when this field's type is array. nullable: type: boolean description: >- @@ -393,100 +414,174 @@ def ai( - array - object description: >- - Holistic input/output example pairs spanning the complete output - record. Sparse outputs are completed with null for omitted fields. + Named, whole-record examples. Each example has a separate input value + or record and the complete expected output record; name is an optional + human-readable label. Use {name: ..., input: ..., output: ...}. Omitted + output fields are completed with null. This differs from examples + nested under one output field, which teach only that field. + required: + - input + - output + properties: + name: + type: string + description: Optional label used to identify this example in the prompt. + input: + description: Source value or record the example should match. + output: + description: Expected result using the field names defined by output. + items: + type: object + required: + - input + - output + properties: + name: + type: string + description: Optional label used to identify this example in the prompt. + input: + description: Source value or record the example should match. + output: + description: Expected result using the field names defined by output. api_key: type: string - description: API Key for the model + description: OpenAI API key used for this wrangle, normally supplied through a recipe variable. model: type: string - description: The name of the AI model to use + description: >- + OpenAI model ID for this call. If omitted, uses the configured + extract.ai default; a saved model definition may supply its own model. threads: type: integer minimum: 1 - description: The number of requests to send in parallel. The configured default is 32. + description: Maximum number of row-level requests sent in parallel. The configured default is 32. timeout: type: number exclusiveMinimum: 0 - description: Per-request timeout in seconds. The configured default is 12. + description: >- + Maximum seconds for one HTTP attempt. The configured default is 12; + deadline can end the overall call sooner. retries: type: integer minimum: 0 description: >- - The number of times to retry if the request fails. - Defaults to 1. Retry delays and request timeouts are bounded - by the total deadline. + Number of additional attempts after a retryable failure. The configured + default is 1. Backoff and request timeouts remain bounded by deadline. url: type: string description: |- - Override the endpoint configured for the selected protocol. - Existing chat/completions URLs are detected for backwards compatibility, - but definitions should set protocol explicitly while being upgraded. + Override the endpoint for the selected protocol. A chat/completions URL + selects the legacy protocol only when protocol is omitted; new recipes + should use the configured Responses endpoint. provider: type: string - description: AI provider. Phase 1 supports OpenAI; this field reserves a stable provider boundary. + description: AI service provider. Currently only OpenAI is supported. enum: - openai protocol: type: string - description: API protocol. Responses is the configured default. + description: >- + OpenAI API protocol. Responses is the configured default and is required + for web_search; chat_completions remains available for legacy definitions. enum: - responses - chat_completions deadline: type: number exclusiveMinimum: 0 - description: Total call budget in seconds, including retries. The configured default is 15. + description: >- + Total seconds allowed for the entire wrangle call, including queued + work, retries, and backoff. The configured default is 15. store: type: boolean - description: Whether OpenAI may store Responses. Defaults to false. + description: Whether OpenAI may store Responses API results. Defaults to false. cache: type: boolean description: >- - Use the bounded warm-instance result cache. Defaults to true. - Set false for a call-level bypass. + Reuse identical successful results from the bounded warm-instance cache. + Defaults to true. Set false when fresh model or web results are required. cache_ttl: type: number exclusiveMinimum: 0 - description: Override the result-cache TTL in seconds for this call. + description: >- + Maximum age in seconds for a cached result used by this call. Applies + to extracted values and web_search_sources together. messages: type: - string - array - description: Optional. Provide additional overall instructions for the AI. + description: >- + Additional overall instruction or list of instructions applied to + every row after the configured extraction prompt and examples. + items: + type: string model_id: type: string - description: Use a saved definition from an extract ai wrangle. + description: >- + ID of a saved extract.ai definition. Use it instead of defining an + output schema. When output is also supplied with model_id in a recipe, + output names the destination column or columns for the saved fields. strict: type: boolean description: >- - Enable structured output strict mode. Default true. Definitions with - dynamic dictionaries automatically use non-strict provider mode and - are validated locally. + Require OpenAI structured-output strict mode. Defaults to true. + Definitions with dynamic dictionary keys automatically switch to + non-strict provider mode and are still validated locally. output_format: type: string - description: Format of the extract output + description: >- + How extracted fields are written. columns writes one dataframe column + per field (default); dictionary keeps one object; concatenate joins + fields into one string using char. enum: - dictionary - columns - concatenate char: type: string - description: Character to use when output_format is concatenate + description: Separator used only when output_format is concatenate. Defaults to comma-space. reasoning: type: object description: >- - Responses API reasoning options. Defaults to effort none for models - that support disabling reasoning; otherwise the provider default applies. + Responses API reasoning controls. Set effort for reasoning-capable + models. The configured default is none when that model supports it; + otherwise the provider default applies. + properties: + effort: + type: string + description: Amount of reasoning work requested from a compatible model. + enum: + - none + - minimal + - low + - medium + - high + - xhigh verbosity: type: string - description: Responses API output verbosity. Defaults to low for models that support low verbosity. + description: >- + Responses API text verbosity for compatible models. Defaults to low + when supported; ignored with a warning for incompatible models. enum: - low - medium - high + web_search: + type: boolean + description: >- + Enable OpenAI Responses web search; the model decides when searching + helps. When true, every row also receives web_search_sources: a + deduplicated list of {title, url} objects in source order, or an empty + list when no source was used. This reserved column is automatic. + Requires protocol responses. Defaults to false. """ output_format_normalized = _normalize_output_format(output_format, "columns") + if not isinstance(web_search, bool): + raise ValueError("web_search must be true or false.") + if web_search and _WEB_SEARCH_SOURCES_KEY in df.columns: + raise ValueError( + f"Column {_WEB_SEARCH_SOURCES_KEY!r} is reserved when web_search is enabled." + ) # If input is provided, extract only those columns # Otherwise, provide the whole dataframe @@ -547,15 +642,35 @@ def ai( # If a schema has been provided, define the target columns if not target_columns and output is not None: target_columns = list(output.keys()) + if web_search and target_columns and _WEB_SEARCH_SOURCES_KEY in target_columns: + raise ValueError( + f"Output column {_WEB_SEARCH_SOURCES_KEY!r} is reserved when web_search is enabled." + ) results = _extract.ai( df_temp.to_dict(orient='records'), api_key=api_key, output=output, model_id=model_id, + web_search=web_search, **kwargs ) + web_search_sources = None + if web_search: + web_search_sources = [] + extraction_results = [] + for result in results: + if not isinstance(result, dict): + web_search_sources.append([]) + extraction_results.append(result) + continue + result = dict(result) + sources = result.pop(_WEB_SEARCH_SOURCES_KEY, []) + web_search_sources.append(sources if isinstance(sources, list) else []) + extraction_results.append(result) + results = extraction_results + try: exploded_df = _pd.json_normalize(results, max_level=0).fillna('').set_index(df.index) @@ -599,6 +714,13 @@ def ai( except: raise RuntimeError("Unable to parse response from AI model") + if web_search_sources is not None: + df[_WEB_SEARCH_SOURCES_KEY] = _pd.Series( + web_search_sources, + index=df.index, + dtype=object, + ) + return df From 08b93ffd64f9a86f6d3a87c3020b3bf02f6a6b82 Mon Sep 17 00:00:00 2001 From: Eric Hills <53243273+ebhills@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:42:25 -0500 Subject: [PATCH 2/4] Clarify extract.ai record and field examples --- .../extract_ai_recipe_examples.ipynb | 17 ++++-- docs/extract_ai_configuration.md | 26 ++++---- tests/recipes/wrangles/test_main.py | 23 ++++++- .../samples/extract ai judge example.wrgl.yml | 2 +- tests/test_ai_definition.py | 38 +++++++++--- tests/test_openai_extract_ai.py | 48 ++++++++++++++- wrangles/ai_definition.py | 47 ++++++++++++--- wrangles/extract.py | 12 +++- wrangles/recipe_wrangles/extract.py | 60 ++++++++++++++++--- 9 files changed, 223 insertions(+), 50 deletions(-) diff --git a/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb b/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb index 81bb4c465..65223d4b0 100644 --- a/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb +++ b/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb @@ -1925,7 +1925,10 @@ "\n", "`examples` on a field nudge the model toward a particular output style or\n", "unit format, without being included in the strict JSON schema sent to\n", - "OpenAI (they're folded into the instructions instead)." + "OpenAI (they're folded into the instructions instead). The existing scalar\n", + "or list form remains supported. A field example can also be a paired object\n", + "with `input` and `output`, plus optional `name` and `notes`; that pair guides\n", + "only the field containing it." ] }, { @@ -2024,11 +2027,12 @@ "id": "b18c5f72", "metadata": {}, "source": [ - "### Whole-record examples: `name`, `input`, and `output`\n", + "### Whole-record examples: `name`, `notes`, `input`, and `output`\n", "\n", - "Top-level `examples` teach a complete input-to-output mapping. Keep the\n", + "Top-level `record_examples` teach a complete input-to-output mapping. Keep the\n", "source under `input`, the desired record under `output`, and optionally add\n", - "a descriptive `name`. Sparse example outputs are completed with `null` for\n", + "a descriptive `name` and explanatory `notes`. Sparse example outputs are\n", + "completed with `null` for\n", "fields you omit. These are different from `examples` nested under one output\n", "field, which guide only that field." ] @@ -2051,8 +2055,9 @@ " certainty:\n", " type: string\n", " enum: [high, medium, low]\n", - " examples:\n", + " record_examples:\n", " - name: exact manufacturer part match\n", + " notes: Use high certainty only when both manufacturer and part number match.\n", " input:\n", " description: ACME valve V-100\n", " output:\n", @@ -2992,7 +2997,7 @@ "| Constrained values | Add `enum: [...]` to a field |\n", "| Numeric / list result | `type: integer` / `type: array` |\n", "| Style hints | Add `examples: [...]` to a field |\n", - "| Whole-record examples | `examples: [{name: ..., input: ..., output: ...}]` |\n", + "| Whole-record examples | `record_examples: [{name: ..., notes: ..., input: ..., output: ...}]` |\n", "| Web search + sources | `web_search: true` adds `web_search_sources` |\n", "| Reasoning models | `model: gpt-5-mini`, `reasoning: {effort: low}`, `verbosity: low` |\n", "| Parallelism | `threads: ` |\n", diff --git a/docs/extract_ai_configuration.md b/docs/extract_ai_configuration.md index 47cb54485..850ab822f 100644 --- a/docs/extract_ai_configuration.md +++ b/docs/extract_ai_configuration.md @@ -45,7 +45,7 @@ cells as sequences; users do not need to quote every object key and value. ## Examples -Definitions support both field-specific and holistic examples. They compile to +Definitions support both field-specific and record examples. They compile to the same stable prompt representation and precede each row's dynamic input. For saved models, the field grid supports: @@ -61,8 +61,9 @@ expected output may be human-friendly JSON/YAML-like syntax. Use explicit list syntax for an array-valued paired output, such as `[Ceramic Tile, Slate]`. An explicit `null` is valid because output fields are nullable by default. -Field-specific examples teach only the named field. Definitions may also -provide holistic examples that pair one input with a multi-field output: +Field-specific examples teach only the named field. Paired field examples may +include optional `name` and `notes` metadata. Definitions may also provide +`record_examples` that pair one input with a multi-field output: ```yaml wrangles: @@ -73,13 +74,16 @@ wrangles: Power Source: type: string examples: - - input: 18V cordless drill + - name: cordless tool + notes: Voltage without a cord indicates a battery. + input: 18V cordless drill output: Battery - Corded Voltage: type: number - examples: + record_examples: - name: corded saw + notes: Use both fields from this complete example. input: 120V corded jig saw output: Power Source: Corded @@ -87,13 +91,15 @@ wrangles: ``` The first `Power Source` item is a paired field example; `Corded` remains -output-only value guidance. Top-level `examples` are holistic. Their output may -be sparse: the compiler inserts `null` for omitted output fields so every -example demonstrates the complete required response shape. Unknown fields and -values that do not match the output schema fail during compilation. +output-only value guidance. Top-level `record_examples` demonstrate the complete +record. Their output may be sparse: the compiler inserts `null` for omitted +output fields so every example demonstrates the complete required response +shape. Unknown fields and values that do not match the output schema fail during +compilation. Optional `name` and `notes` metadata are included in model guidance +at both levels. Saved-model content may likewise include a top-level `Examples` array of -holistic pairs. A dedicated WranglesXL interface for those examples can be +record pairs. A dedicated WranglesXL interface for those examples can be added later without changing the compiler or runtime contract. ## Result cache diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index e708adc31..2fd90f0e5 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -9361,16 +9361,35 @@ def test_extract_ai_schema_documents_public_parameters_and_example_shape(self): assert properties["web_search"]["type"] == "boolean" assert "web_search_sources" in properties["web_search"]["description"] - examples = properties["examples"] + assert "examples" not in properties + examples = properties["record_examples"] + assert examples["title"] == "Record examples" assert examples["required"] == ["input", "output"] assert examples["items"]["required"] == ["input", "output"] - assert set(examples["items"]["properties"]) == {"name", "input", "output"} + assert set(examples["items"]["properties"]) == { + "name", + "notes", + "input", + "output", + } assert all( definition["description"].strip() for definition in examples["items"]["properties"].values() ) field_schema = next(iter(properties["output"]["patternProperties"].values())) + field_examples = field_schema["properties"]["examples"] + assert field_examples["title"] == "Field examples" + assert "backward-compatible" in field_examples["description"] + assert "optional name and notes" in field_examples["description"] + paired_field_example = field_examples["items"]["anyOf"][0] + assert paired_field_example["required"] == ["input", "output"] + assert set(paired_field_example["properties"]) == { + "name", + "notes", + "input", + "output", + } assert all( definition["description"].strip() for definition in field_schema["properties"].values() diff --git a/tests/samples/extract ai judge example.wrgl.yml b/tests/samples/extract ai judge example.wrgl.yml index 1572a253c..3c7359444 100644 --- a/tests/samples/extract ai judge example.wrgl.yml +++ b/tests/samples/extract ai judge example.wrgl.yml @@ -128,7 +128,7 @@ wrangles: description: | True when confidence is Uncertain or Selected Classification is "None (needs review)"; otherwise false. - examples: + record_examples: - name: dynamic two-candidate example input: Input Description: 18V cordless drill with battery and charger diff --git a/tests/test_ai_definition.py b/tests/test_ai_definition.py index 00a63405e..ffc60ebca 100644 --- a/tests/test_ai_definition.py +++ b/tests/test_ai_definition.py @@ -457,13 +457,18 @@ def test_saved_model_compiles_new_field_example_columns_and_legacy_examples(): assert compiled.output["Power_Source"]["examples"] == ["Corded", "Battery"] -def test_recipe_field_pairs_and_holistic_examples_compile_to_stable_guidance(): +def test_recipe_field_pairs_and_record_examples_compile_to_stable_guidance(): compiled = ai_definition.compile_definition( { "Power Source": { "type": "string", "examples": [ - {"input": "18V cordless drill", "output": "Battery"}, + { + "name": "cordless tool", + "notes": "Voltage without a cord indicates a battery.", + "input": "18V cordless drill", + "output": "Battery", + }, "Corded", ], }, @@ -473,6 +478,7 @@ def test_recipe_field_pairs_and_holistic_examples_compile_to_stable_guidance(): examples=[ { "name": "corded saw", + "notes": "Use both fields from this complete record example.", "input": "120V corded jig saw", "output": {"Power Source": "Corded", "Voltage": 120}, }, @@ -486,10 +492,15 @@ def test_recipe_field_pairs_and_holistic_examples_compile_to_stable_guidance(): assert compiled.output["Power_Source"]["examples"] == ["Corded"] assert compiled.field_examples[0].field == "Power_Source" assert compiled.field_examples[0].output == "Battery" + assert compiled.field_examples[0].name == "cordless tool" + assert compiled.field_examples[0].notes == "Voltage without a cord indicates a battery." assert compiled.record_examples[0].output == { "Power_Source": "Corded", "Voltage": 120, } + assert compiled.record_examples[0].notes == ( + "Use both fields from this complete record example." + ) assert compiled.record_examples[1].output == { "Power_Source": None, "Voltage": None, @@ -498,10 +509,13 @@ def test_recipe_field_pairs_and_holistic_examples_compile_to_stable_guidance(): guidance = ai_definition.render_example_guidance(compiled) assert " tuple: f"{path}.examples[{index}].output", "must be provided for a paired field example.", ) - pairs.append({ + pair = { "input": self._parse_json_value(item["input"]), "output": self._parse_json_value(item["output"]), - }) + } + for metadata_key in ("name", "notes"): + metadata_value = item.get(metadata_key) + if metadata_value not in (None, ""): + pair[metadata_key] = str(metadata_value) + pairs.append(pair) if value_examples: schema["examples"] = value_examples @@ -427,11 +435,15 @@ def compile_record_examples( name = raw_example.get("name") if name in (None, ""): name = f"example-{index + 1}" + notes = raw_example.get("notes") + if notes in (None, ""): + notes = None compiled.append( CompiledRecordExample( name=str(name), input=example_input, output=complete_output, + notes=str(notes) if notes is not None else None, ) ) return compiled @@ -810,7 +822,7 @@ def compile_definition( Compile a recipe/Python output and optional saved XL model definition. Direct keyed output overrides fields from the saved model. Saved model - instructions and holistic examples run first, followed by call-level + instructions and record examples run first, followed by call-level messages and examples. """ compiler = _Compiler(source) @@ -918,6 +930,8 @@ def compile_definition( sanitized_output[sanitized_field], f"output.{original_field}.examples[{index}].output", ), + name=pair.get("name"), + notes=pair.get("notes"), ) ) @@ -977,8 +991,17 @@ def render_example_guidance(compiled: CompiledAIDefinition) -> str: if compiled.field_examples: blocks = [] for example in compiled.field_examples: - blocks.append("\n".join([ - f'', + attributes = f" field={_json.dumps(example.field)}" + if example.name: + attributes += f" name={_json.dumps(example.name)}" + lines = [ + f"", + ] + if example.notes: + lines.append( + f"{_json.dumps(example.notes, ensure_ascii=False)}" + ) + lines.extend([ f"{_json.dumps(example.input, ensure_ascii=False, default=str)}", ( "" @@ -986,7 +1009,8 @@ def render_example_guidance(compiled: CompiledAIDefinition) -> str: "" ), "", - ])) + ]) + blocks.append("\n".join(lines)) sections.append( "Each field example below demonstrates only the named output field. " "Do not infer values for other fields from its expected value.\n" @@ -996,8 +1020,14 @@ def render_example_guidance(compiled: CompiledAIDefinition) -> str: if compiled.record_examples: blocks = [] for example in compiled.record_examples: - blocks.append("\n".join([ + lines = [ f'', + ] + if example.notes: + lines.append( + f"{_json.dumps(example.notes, ensure_ascii=False)}" + ) + lines.extend([ f"{_json.dumps(example.input, ensure_ascii=False, default=str)}", ( "" @@ -1005,7 +1035,8 @@ def render_example_guidance(compiled: CompiledAIDefinition) -> str: "" ), "", - ])) + ]) + blocks.append("\n".join(lines)) sections.append( "The record examples below demonstrate the complete expected output shape.\n" + "\n".join(blocks) diff --git a/wrangles/extract.py b/wrangles/extract.py index e36267bf4..a3f0850b4 100644 --- a/wrangles/extract.py +++ b/wrangles/extract.py @@ -137,7 +137,8 @@ def ai( timeout: float = None, retries: int = None, messages: list = None, - examples: list = None, + examples: _Union[dict, list] = None, + record_examples: _Union[dict, list] = None, url: str = None, strict: bool = None, reasoning: dict = None, @@ -173,7 +174,8 @@ def ai( :param timeout: (Optional) Timeout in seconds for each API call. :param retries: (Optional) Number of retries to attempt on failure. :param messages: (Optional) Overall prompts to pass additional instructions. - :param examples: (Optional) Holistic examples containing separate input and output values, plus an optional name. + :param examples: (Optional) Compatibility alias for record_examples. + :param record_examples: (Optional) Whole-record examples containing input and output, with optional name and notes. :param url: (Optional) Override the configured endpoint. :param strict: (Optional) Enable structured output strict mode. Dynamic object schemas \ automatically use non-strict mode and are validated locally. @@ -252,6 +254,10 @@ def ai( if messages is None: messages = [] + if record_examples not in (None, "") and examples not in (None, ""): + raise ValueError("Use record_examples or examples, not both.") + if record_examples in (None, ""): + record_examples = examples # Ensure input is a list input_was_scalar = False @@ -268,7 +274,7 @@ def ai( output, model=model, messages=messages, - examples=examples, + examples=record_examples, strict=strict, saved_model_content=saved_model_content, source=f"saved model {model_id}" if model_id else "recipe/Python output", diff --git a/wrangles/recipe_wrangles/extract.py b/wrangles/recipe_wrangles/extract.py index c69dcaac2..49f56d70e 100644 --- a/wrangles/recipe_wrangles/extract.py +++ b/wrangles/recipe_wrangles/extract.py @@ -294,6 +294,7 @@ def ai( input: list = None, output: _Union[dict, str, list] = None, model_id: str = None, + record_examples: _Union[dict, list] = None, output_format: str = None, char: str = ", ", web_search: bool = False, @@ -370,6 +371,7 @@ def ai( does not substitute this value when evidence is missing; describe fallback behavior explicitly or allow null. examples: + title: Field examples type: - array - object @@ -379,9 +381,40 @@ def ai( - boolean - "null" description: >- - Field-specific examples. Provide typical output values to - demonstrate style, or objects with separate input and output - keys to show how a source value maps to this field. + Field-specific examples. The backward-compatible form is a + scalar or list of typical output values. A paired example may + instead use input and output, with optional name and notes. + Paired examples apply only to this output field; use + record_examples for complete output records. + properties: + name: + type: string + description: Optional label included with this paired field example. + notes: + type: string + description: Optional explanatory guidance included with this paired field example. + input: + description: Source value or record for this paired field example. + output: + description: Expected value for this output field only. + items: + anyOf: + - type: object + required: + - input + - output + properties: + name: + type: string + description: Optional label included with this paired field example. + notes: + type: string + description: Optional explanatory guidance included with this paired field example. + input: + description: Source value or record for this paired field example. + output: + description: Expected value for this output field only. + - description: Backward-compatible output-only example value. properties: type: - object @@ -409,16 +442,18 @@ def ai( description: >- Whether the field may return null. Defaults to true while the field key remains required. Set false to opt out. - examples: + record_examples: + title: Record examples type: - array - object description: >- - Named, whole-record examples. Each example has a separate input value - or record and the complete expected output record; name is an optional - human-readable label. Use {name: ..., input: ..., output: ...}. Omitted - output fields are completed with null. This differs from examples - nested under one output field, which teach only that field. + Whole-record examples. Each example has a separate input value or + record and the complete expected output record. Optional name and + notes provide model-visible context. Use {name: ..., notes: ..., + input: ..., output: ...}. Omitted output fields are completed with + null. This differs from examples nested under one output field, which + teach only that field. required: - input - output @@ -426,6 +461,9 @@ def ai( name: type: string description: Optional label used to identify this example in the prompt. + notes: + type: string + description: Optional explanatory guidance included with this example. input: description: Source value or record the example should match. output: @@ -439,6 +477,9 @@ def ai( name: type: string description: Optional label used to identify this example in the prompt. + notes: + type: string + description: Optional explanatory guidance included with this example. input: description: Source value or record the example should match. output: @@ -652,6 +693,7 @@ def ai( api_key=api_key, output=output, model_id=model_id, + record_examples=record_examples, web_search=web_search, **kwargs ) From 8fa951a2d251f1ca25e6487d6498cb81cba80f12 Mon Sep 17 00:00:00 2001 From: Eric Hills <53243273+ebhills@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:55:17 -0500 Subject: [PATCH 3/4] Rename extract.ai messages to instructions --- .../extract_ai_recipe_examples.ipynb | 1 + docs/extract_ai_configuration.md | 22 ++++++++++++++ tests/recipes/wrangles/test_main.py | 8 +++++ .../samples/extract ai judge example.wrgl.yml | 2 +- tests/test_ai_definition.py | 30 +++++++++++++++++++ tests/test_openai_extract_ai.py | 29 +++++++++++++++--- wrangles/extract.py | 17 +++++++---- wrangles/recipe_wrangles/extract.py | 10 +++++-- 8 files changed, 105 insertions(+), 14 deletions(-) diff --git a/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb b/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb index 65223d4b0..5f5be9de3 100644 --- a/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb +++ b/docs/examples/extract_ai/extract_ai_recipe_examples.ipynb @@ -2995,6 +2995,7 @@ "| Multiple fields | Add more keys under `output` |\n", "| Quick shorthand | `output: ` (uses input column name) |\n", "| Constrained values | Add `enum: [...]` to a field |\n", + "| Overall guidance | `instructions: ...` applies guidance to every input row |\n", "| Numeric / list result | `type: integer` / `type: array` |\n", "| Style hints | Add `examples: [...]` to a field |\n", "| Whole-record examples | `record_examples: [{name: ..., notes: ..., input: ..., output: ...}]` |\n", diff --git a/docs/extract_ai_configuration.md b/docs/extract_ai_configuration.md index 850ab822f..dab71e7b0 100644 --- a/docs/extract_ai_configuration.md +++ b/docs/extract_ai_configuration.md @@ -19,6 +19,28 @@ versioned replacement YAML file to override the complete configuration. Recipes and Python calls can override these settings individually. Saved XL models and recipe outputs are compiled through the same definition compiler. +## Instructions + +Use `instructions` for guidance that applies to every input row: + +```yaml +wrangles: + - extract.ai: + input: Description + api_key: ${OPENAI_API_KEY} + instructions: + - Prefer explicit evidence over inferred evidence. + - Normalize dimensions to inches. + output: + Product Type: + type: string +``` + +Instructions are useful for decision rules, evidence priorities, +normalization requirements, or other behavior that applies to the complete +extraction. The former `messages` parameter remains available as a compatibility +alias but is no longer advertised in the recipe schema. Do not provide both. + ## Nullable output fields Defined output keys remain required so strict Structured Outputs always return diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index 2fd90f0e5..058ad1f2c 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -9361,6 +9361,14 @@ def test_extract_ai_schema_documents_public_parameters_and_example_shape(self): assert properties["web_search"]["type"] == "boolean" assert "web_search_sources" in properties["web_search"]["description"] + assert "messages" not in properties + instructions = properties["instructions"] + assert instructions["title"] == "Instructions" + assert instructions["type"] == ["string", "array"] + assert "every input row" in instructions["description"] + assert "decision rules" in instructions["description"] + assert "after" not in instructions["description"].lower() + assert "examples" not in properties examples = properties["record_examples"] assert examples["title"] == "Record examples" diff --git a/tests/samples/extract ai judge example.wrgl.yml b/tests/samples/extract ai judge example.wrgl.yml index 3c7359444..31193409b 100644 --- a/tests/samples/extract ai judge example.wrgl.yml +++ b/tests/samples/extract ai judge example.wrgl.yml @@ -33,7 +33,7 @@ wrangles: input: Judge Input api_key: ${OPENAI_API_KEY} model: gpt-5.4-mini - messages: | + instructions: | You are an expert classifier for supplier product content. Each input contains: diff --git a/tests/test_ai_definition.py b/tests/test_ai_definition.py index ffc60ebca..fd139a99d 100644 --- a/tests/test_ai_definition.py +++ b/tests/test_ai_definition.py @@ -677,6 +677,7 @@ def call_structured(data, api_key, payload, *args): input: Description api_key: dummy threads: 1 + instructions: Prefer explicit source values. output: Color: type: string @@ -701,8 +702,37 @@ def call_structured(data, api_key, payload, *args): assert result["Voltage"].tolist() == [""] assert " _Union[dict, list]: """ @@ -173,7 +174,10 @@ def ai( :param threads: (Optional) Number of threads to use for parallel processing. :param timeout: (Optional) Timeout in seconds for each API call. :param retries: (Optional) Number of retries to attempt on failure. - :param messages: (Optional) Overall prompts to pass additional instructions. + :param instructions: (Optional) Additional guidance applied to every input row. Use this for + decision rules, evidence priorities, normalization requirements, or other behavior that + applies to the complete extraction. + :param messages: (Optional) Compatibility alias for instructions. :param examples: (Optional) Compatibility alias for record_examples. :param record_examples: (Optional) Whole-record examples containing input and output, with optional name and notes. :param url: (Optional) Override the configured endpoint. @@ -191,7 +195,6 @@ def ai( :param cache_ttl: (Optional) Override the result-cache TTL in seconds for this call. :param web_search: (Optional) Enable native Responses web search. Each result then includes a web_search_sources list containing source titles and URLs. Defaults to False. - :return: Extracted information. When web_search is true, returns a dictionary (or list of dictionaries) containing web_search_sources, including for single-field output. """ @@ -252,8 +255,10 @@ def ai( raise ValueError("reasoning must be an object such as {'effort': 'none'}.") _validate_ai_runtime_settings(threads, timeout, retries, deadline) - if messages is None: - messages = [] + if instructions not in (None, "") and messages not in (None, ""): + raise ValueError("Use instructions or messages, not both.") + if instructions in (None, ""): + instructions = messages if record_examples not in (None, "") and examples not in (None, ""): raise ValueError("Use record_examples or examples, not both.") if record_examples in (None, ""): @@ -273,7 +278,7 @@ def ai( compiled = _ai_definition.compile_definition( output, model=model, - messages=messages, + messages=instructions, examples=record_examples, strict=strict, saved_model_content=saved_model_content, diff --git a/wrangles/recipe_wrangles/extract.py b/wrangles/recipe_wrangles/extract.py index 49f56d70e..a61cb0fd7 100644 --- a/wrangles/recipe_wrangles/extract.py +++ b/wrangles/recipe_wrangles/extract.py @@ -298,6 +298,7 @@ def ai( output_format: str = None, char: str = ", ", web_search: bool = False, + instructions: _Union[str, list] = None, **kwargs ): """ @@ -547,13 +548,15 @@ def ai( description: >- Maximum age in seconds for a cached result used by this call. Applies to extracted values and web_search_sources together. - messages: + instructions: + title: Instructions type: - string - array description: >- - Additional overall instruction or list of instructions applied to - every row after the configured extraction prompt and examples. + Additional guidance applied to every input row. Use this for decision + rules, evidence priorities, normalization requirements, or other + behavior that applies to the complete extraction. items: type: string model_id: @@ -695,6 +698,7 @@ def ai( model_id=model_id, record_examples=record_examples, web_search=web_search, + instructions=instructions, **kwargs ) From dc8d2a8fd919513671088d2aefc8eaf250f7f988 Mon Sep 17 00:00:00 2001 From: Eric Hills <53243273+ebhills@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:48:14 -0500 Subject: [PATCH 4/4] Initialize response.json() Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- wrangles/openai_responses.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wrangles/openai_responses.py b/wrangles/openai_responses.py index 419b5ffbc..7dad8c0df 100644 --- a/wrangles/openai_responses.py +++ b/wrangles/openai_responses.py @@ -771,6 +771,7 @@ def failure(message: str, response_json: dict = None) -> dict: if response is not None and response.ok: try: +response_json = None response_json = response.json() output_text = extract_response_text(response_json) parsed = _json.loads(output_text)