From 03e4365d9a936656195123b3bf654a9d397b05c4 Mon Sep 17 00:00:00 2001 From: Eric Hills <53243273+ebhills@users.noreply.github.com> Date: Sat, 26 Sep 2026 15:14:46 -0500 Subject: [PATCH 1/3] Add Typesafe AI question wrangles --- .gitignore | 1 + docs/ai_configuration.md | 30 +++ docs/ai_questions.md | 242 ++++++++++++++++++ pytest-local.ini | 2 + scripts/test-local.ps1 | 1 + tests/recipes/wrangles/test_ai.py | 367 +++++++++++++++++++++++++++ tests/test_ai.py | 299 ++++++++++++++++++++++ tests/test_ai_config.py | 65 +++++ tests/test_typesafe.py | 269 ++++++++++++++++++++ wrangles/__init__.py | 2 +- wrangles/ai.py | 281 ++++++++++++++++++++ wrangles/ai_cache.py | 30 ++- wrangles/ai_defaults.yml | 67 +++++ wrangles/clients/typesafe.py | 194 ++++++++++++++ wrangles/dataframe.py | 5 + wrangles/recipe.py | 78 +++--- wrangles/recipe_wrangles/__init__.py | 1 + wrangles/recipe_wrangles/ai.py | 182 +++++++++++++ wrangles/recipe_wrangles/main.py | 13 +- 19 files changed, 2067 insertions(+), 62 deletions(-) create mode 100644 docs/ai_questions.md create mode 100644 tests/recipes/wrangles/test_ai.py create mode 100644 tests/test_ai.py create mode 100644 tests/test_typesafe.py create mode 100644 wrangles/ai.py create mode 100644 wrangles/clients/typesafe.py create mode 100644 wrangles/recipe_wrangles/ai.py diff --git a/.gitignore b/.gitignore index a56540f0..ffba90d7 100644 --- a/.gitignore +++ b/.gitignore @@ -124,6 +124,7 @@ env.bak/ venv.bak/ docs/* !docs/ai_configuration.md +!docs/ai_questions.md !docs/extract_ai_user_guide.md !docs/github-app-deployment.md !docs/search_ai_mode.md diff --git a/docs/ai_configuration.md b/docs/ai_configuration.md index bfb93ab5..451e61bd 100644 --- a/docs/ai_configuration.md +++ b/docs/ai_configuration.md @@ -130,12 +130,35 @@ returned dictionary does not alter cached configuration. | Operation | Callers | Configured settings | | --- | --- | --- | +| `ai.choose`, `ai.score`, `ai.true_false`, `ai.questions` | Python and recipe named questions | Provider, model, endpoint, concurrency, timeout, retries, cache | | `extract.ai` | Python and recipe extraction | Model, endpoints, model tuning, concurrency, timeout, retries, strictness, storage, cache, prompt | | `embeddings` | `openai.embeddings` and recipe `create.embeddings` | Provider, model, endpoint, batch size, concurrency, timeout, retries, precision, dimensions, Jina task/normalization/truncation | | `search.retrieve_link_content` | Python, recipe, and Gemini URL-context client | Model, endpoint/API version, concurrency, timeout, retries, temperature/top-p/top-k/token limits/stop sequences | | `generate.ai` | Python and recipe generation | Model, endpoint, reasoning/text tuning, concurrency, timeout, retries, strictness | | `huggingface` | Generic recipe task wrangle | Explicit model, endpoint, timeout, retries, task parameters | +The four named-question operations use Typesafe's `systemone` protocol at +`https://api.typesafe.ai/v1/systemone`, with `jev-1.13.0` as their pinned default +model. Each operation has its own default role; changing an extraction or global +model does not change these operations. Explicit unlisted Typesafe model names +remain available. This adapter supports only `provider: typesafe` and +`protocol: systemone`; adding another catalog provider alone does not implement +an adapter for it. + +Their packaged runtime defaults are 10 concurrent requests, a 30-second timeout, +and one additional attempt after a transient failure. Successful results use a +bounded in-memory cache with a one-hour TTL, at most 512 entries, and a maximum +value size of 65,536 bytes. Duplicate in-flight requests share their result, and +periodic cache logging is disabled. `WRANGLES_AI_CACHE_*` environment controls +apply to these operations independently of the existing +`WRANGLES_EXTRACT_AI_CACHE_*` controls. See [AI questions](ai_questions.md) for +question schemas, examples, output columns, and cache overrides. + +Keep credentials outside the catalog. Supply `api_key` explicitly, use the local +`TYPESAFE_API_KEY` environment variable, or use `api_key: ${TYPESAFE_API_KEY}` in +a hosted recipe with that managed secret. Hosted secrets are supplied as recipe +variables; they are not placed in the worker's environment. + OpenAI embeddings retain `text-embedding-3-small` and their existing dimensions unless explicitly configured otherwise. Jina requires an explicit model or a Jina catalog model assigned the `embeddings` role; the package does not invent a @@ -209,6 +232,13 @@ their original replacement semantics. Existing `model_capabilities` flags remain accepted on that compatibility path and retain packaged capability inheritance. Generation still follows a version-1 extraction model override; other operations use packaged defaults because version 1 did not configure them. +This includes all four Typesafe named-question operations. + +A version-2 replacement remains authoritative: it must define any new operation +you intend to use, its provider, and a default model role (unless you supply the +model explicitly). An older version-2 file does not silently inherit the new +`ai.*` entries from the packaged catalog. Update it from the packaged catalog +before calling these operations. For new files, copy version 2 and edit the provider/model catalog and operation settings. The compatibility API `model_capabilities()` remains available for diff --git a/docs/ai_questions.md b/docs/ai_questions.md new file mode 100644 index 00000000..4f46ef77 --- /dev/null +++ b/docs/ai_questions.md @@ -0,0 +1,242 @@ +# AI questions + +`ai.choose`, `ai.score`, `ai.true_false`, and `ai.questions` answer named questions +about each input record through Typesafe. They use the existing +[AI model catalog](ai_configuration.md), with `provider: typesafe`, +`protocol: systemone`, and the pinned model `jev-1.13.0` as defaults. + +`extract.ai` keeps its current API and behavior. The new `ai` namespace leaves +room for a later compatible `ai.extract` migration; that alias is not introduced +here. + +## Named questions and answers + +Every operation takes a nonempty `questions` mapping. Its keys are the question +names and form the default output column names. Every question has +`instructions`, which can be a string, a JSON object, or a JSON array. +Instructions and supplied descriptions must be nonempty; choose descriptions +may also use `null` when the option label is sufficient. + +| Operation | Criteria | Answer fields and default recipe columns | +| --- | --- | --- | +| `ai.choose` | A mapping of 1–255 nonblank option labels to descriptions. Descriptions can be strings, JSON objects, JSON arrays, or `null`. | `choice` → ``; `confidence` → `_confidence`; `probabilities` → `_probabilities` | +| `ai.score` | An ordered list of 2–10 unique nonblank strings describing the scoring criteria. Structured criteria objects are not accepted. | `score` → ``; `confidence` → `_confidence`; `probabilities` → `_probabilities` | +| `ai.true_false` | Optional `"true"` and/or `"false"` descriptions, supplied as strings, JSON objects, or JSON arrays. Quote these keys in YAML. | `probability_true` → ``; `true_criteria` → `_true_criteria` | +| `ai.questions` | Each question declares `type: choose`, `type: score`, or `type: true_false` and uses the corresponding schema above. | The columns for that question's type. | + +The individual operations allow several questions of the same type. Their +`type` field may be omitted; if supplied, it must match the operation. +`ai.questions` requires each question's `type`, allowing all three types in one +request per input record. + +Choose probabilities use the option labels as keys. Score probabilities use the +exact criterion strings as keys, preserving their order. The score is Typesafe's +native numeric score from `0` to `N-1`, where `N` is the number of criteria; +there is no `scale` parameter or automatic rescaling. +True/false returns a probability, without inventing a Boolean decision, +`probability_false`, or confidence score. Its `true_criteria` field preserves +the question's supplied `"true"` criterion, including structured JSON; if that +criterion is omitted, the value is an empty string. + +## Recipes + +The examples assume a `Description` input column and a managed secret named +`TYPESAFE_API_KEY`. Local recipes can also resolve this placeholder from that +environment variable. Never put the key value in a recipe or the model catalog. + +Choose from named options: + +```yaml +wrangles: + - ai.choose: + input: Description + api_key: ${TYPESAFE_API_KEY} + questions: + Product Class: + instructions: Choose the class supported by the product description. + criteria: + Bearing: A component supporting a rotating shaft. + Belt: A flexible loop that transfers motion. + Other: None of the named classes is supported. +``` + +This adds `Product Class`, `Product Class_confidence`, and +`Product Class_probabilities`. + +Score against an ordered set of criteria: + +```yaml +wrangles: + - ai.score: + input: Description + api_key: ${TYPESAFE_API_KEY} + questions: + Description Quality: + instructions: Assess how specifically this description identifies the product. + criteria: + - The product cannot be identified. + - The general product type can be identified. + - The product type and distinguishing specifications can be identified. +``` + +This adds `Description Quality`, `Description Quality_confidence`, and +`Description Quality_probabilities`. The probability mapping keeps those three +criterion sentences as its keys. + +Evaluate a true/false question: + +```yaml +wrangles: + - ai.true_false: + input: Description + api_key: ${TYPESAFE_API_KEY} + questions: + Stainless Steel: + instructions: Is the product explicitly described as stainless steel? + criteria: + "true": The description explicitly identifies stainless steel. + "false": The description identifies a different material or omits the material. +``` + +This adds `Stainless Steel` containing `probability_true` and +`Stainless Steel_true_criteria`. A downstream rule can apply the threshold +appropriate for the workflow. + +Combine question types in one operation: + +```yaml +wrangles: + - ai.questions: + input: [Manufacturer, Part Number, Description] + api_key: ${TYPESAFE_API_KEY} + questions: + Product Class: + type: choose + instructions: Select the product class using all available fields. + criteria: + Bearing: A component supporting a rotating shaft. + Belt: A flexible loop that transfers motion. + Other: None of the named classes is supported. + Description Quality: + type: score + instructions: Assess how specifically the description identifies the product. + criteria: + - The product cannot be identified. + - The general product type can be identified. + - The product type and distinguishing specifications can be identified. + Stainless Steel: + type: true_false + instructions: Is the product explicitly described as stainless steel? + criteria: + "true": The product information explicitly identifies stainless steel. + "false": The product information identifies another material or omits material. +``` + +### Output names + +Omitting a question's `output`, setting it to `null`, or using an empty or +whitespace-only string selects the default columns. To rename columns, supply +the complete ordered list inside that question: + +```yaml +questions: + Product Class: + instructions: Choose the product class. + criteria: + Bearing: A component supporting a rotating shaft. + Other: Any other product. + output: [Class, Class Confidence, Class Probabilities] +``` + +Choose and score require exactly three output names; true/false requires exactly +two, in the order shown in the answer table. Every name must be a nonblank +string. Partial lists, an empty list, and a single nonempty string are invalid. +Output names must also be unique across all questions in the operation. +The operation's `questions` mapping controls the outputs; there is no single +operation-level `output` selection. + +## Python + +The public functions share these parameters: + +```python +from wrangles import ai + +ai.choose(data, questions, api_key=None, *, model=None, provider=None, + protocol=None, threads=None, timeout=None, retries=None, + cache=None, cache_ttl=None) +ai.score(data, questions, api_key=None, *, model=None, provider=None, + protocol=None, threads=None, timeout=None, retries=None, + cache=None, cache_ttl=None) +ai.true_false(data, questions, api_key=None, *, model=None, provider=None, + protocol=None, threads=None, timeout=None, retries=None, + cache=None, cache_ttl=None) +ai.questions(data, questions, api_key=None, *, model=None, provider=None, + protocol=None, threads=None, timeout=None, retries=None, + cache=None, cache_ttl=None) +``` + +A string or dictionary input returns a dictionary of named answers. A list of +strings or dictionaries returns one named-answer dictionary per input record, +in input order. Question definitions use the same schema as the recipe examples. + +```python +from wrangles import ai + +# The client uses TYPESAFE_API_KEY from the environment when api_key is omitted. +answers = ai.choose( + {"Description": "Stainless steel ball bearing"}, + questions={ + "Product Class": { + "instructions": "Choose the product class.", + "criteria": {"Bearing": "Supports a rotating shaft.", "Other": None}, + }, + }, +) +selected_class = answers["Product Class"]["choice"] +probabilities = answers["Product Class"]["probabilities"] +``` + +Python returns the answer fields listed in the table rather than flattened +column names. A question's `output` names are recipe column controls. + +## Runtime settings and compatibility + +Runtime options override the operation's catalog settings: `threads` controls +concurrent requests, `timeout` is in seconds, and `retries` counts additional +attempts after the initial request. The defaults are 10 threads, 30 seconds, +and one retry. Use `retries: 0` to disable retries. Only transient failures are +retried; invalid credentials, question definitions, and malformed results fail. +Errors do not become fabricated classification answers. + +Successful results are cached in memory by request identity. The packaged +defaults enable a one-hour TTL, up to 512 entries, and up to 65,536 bytes per +cached value. Concurrent duplicate requests share one in-flight request. +Use `cache: false` or `cache_ttl` in seconds to override the corresponding +catalog values. Process-level environment controls take precedence: + +- `WRANGLES_AI_CACHE_ENABLED` +- `WRANGLES_AI_CACHE_TTL_SECONDS` +- `WRANGLES_AI_CACHE_MAX_ENTRIES` +- `WRANGLES_AI_CACHE_MAX_VALUE_BYTES` +- `WRANGLES_AI_CACHE_SINGLE_FLIGHT` +- `WRANGLES_AI_CACHE_LOG_EVERY` + +These controls do not change `extract.ai`'s existing +`WRANGLES_EXTRACT_AI_CACHE_*` settings. Provider, endpoint, model, credentials, +input, and question definitions are included in the cache identity, preventing +results from being shared across different requests or credentials. Secrets +and raw prompts are not written into cache keys or logs. Cache logging is +disabled by default for these operations. + +A version-1 AI configuration uses packaged defaults for the new operations. +A version-2 replacement file must include their provider and operation entries; +missing entries are not silently merged. Only Typesafe's `systemone` adapter is +implemented for these functions, even if another provider appears in the catalog. +Explicit Typesafe model names may be supplied; the provider determines whether +that model is available. + +See the [Typesafe introduction](https://docs.typesafe.ai/introduction) for the +provider's API documentation. Offline tests validate request and response +contracts; live accuracy, account limits, and deployment behavior require a +separate live-service check. diff --git a/pytest-local.ini b/pytest-local.ini index eefb97b1..f24f4704 100644 --- a/pytest-local.ini +++ b/pytest-local.ini @@ -1,6 +1,7 @@ [pytest] testpaths = tests/test_ai_cache.py + tests/test_ai.py tests/test_ai_config.py tests/test_ai_caller_config.py tests/test_ai_definition.py @@ -13,6 +14,7 @@ testpaths = tests/test_search-ai_mode.py tests/test_search_ai_extraction.py tests/test_standardize_clean_markup.py + tests/test_typesafe.py "tests/test_back_ processes.py" tests/recipes tests/connectors/test_access.py diff --git a/scripts/test-local.ps1 b/scripts/test-local.ps1 index 367ef9f7..f8bcb27f 100644 --- a/scripts/test-local.ps1 +++ b/scripts/test-local.ps1 @@ -25,6 +25,7 @@ foreach ($name in @( "HUGGINGFACE_TOKEN", "OPENAI_API_KEY", "SERPAPI_API_KEY", + "TYPESAFE_API_KEY", "WRANGLES_PASSWORD", "WRANGLES_USER", "WRANGLES_AI_CONFIG" diff --git a/tests/recipes/wrangles/test_ai.py b/tests/recipes/wrangles/test_ai.py new file mode 100644 index 00000000..9ed740b4 --- /dev/null +++ b/tests/recipes/wrangles/test_ai.py @@ -0,0 +1,367 @@ +"""Offline recipe integration for named Typesafe questions.""" + +import copy +import inspect +import os + +import jsonschema +import pandas as pd +import pytest +import yaml + +import wrangles +from wrangles.recipe_wrangles import ai as recipe_ai +from wrangles.clients import typesafe + + +@pytest.fixture(autouse=True) +def no_live_auth(monkeypatch): + monkeypatch.setattr(wrangles.recipe._auth, "get_applied_permission_group", lambda: None) + + +def question(kind, **overrides): + result = {"instructions": "Evaluate the supplied product."} + if kind == "choose": + result["criteria"] = {"bearing": "A bearing", "belt": "A belt"} + elif kind == "score": + result["criteria"] = ["Low", "Medium", "High"] + else: + result["criteria"] = {"true": "Suitable outdoors", "false": "Indoor use only"} + return {**result, **overrides} + + +@pytest.fixture +def core_calls(monkeypatch): + """Mock only the network-facing core entry points, retaining real validation.""" + calls = [] + + def operation(kind): + def execute(data, questions, **settings): + calls.append({"kind": kind, "data": data, "questions": questions, "settings": settings}) + prepared = wrangles.ai._prepare_questions(questions, kind=kind) + results = [] + for row in data: + response = {} + for label, definition in prepared.items(): + if definition["type"] == "choose": + response[label] = { + "choice": "bearing", "confidence": 0.8, + "probabilities": {"bearing": 0.8, "belt": 0.2}, + } + elif definition["type"] == "score": + response[label] = { + "score": 1.42, "confidence": 0.7, + "probabilities": {"Low": 0.0, "Medium": 0.58, "High": 0.42}, + } + else: + response[label] = { + "probability_true": 0.75, + "true_criteria": definition.get("criteria", {}).get("true", ""), + } + results.append(response) + return results + return execute + + for name in ("choose", "score", "true_false", "questions"): + monkeypatch.setattr(wrangles.ai, name, operation(None if name == "questions" else name)) + return calls + + +def run(kind, definitions, dataframe=None, **options): + if dataframe is None: + dataframe = pd.DataFrame({"Description": ["first", "second"], "Ignore": [1, 2]}) + return wrangles.recipe.run( + {"wrangles": [{f"ai.{kind}": {"input": "Description", "questions": definitions, **options}}]}, + dataframe=dataframe, + ) + + +@pytest.mark.parametrize("kind", ["choose", "score", "true_false"]) +def test_homogeneous_named_questions_share_row_records_and_defaults(kind, core_calls): + source = pd.DataFrame({"Description": ["first", "second"], "Ignore": [1, 2]}, index=[8, 3]) + result = run(kind, {"First": question(kind), "Second": question(kind)}, dataframe=source, + api_key="test-key", threads=3, cache=False) + + assert result.index.tolist() == [8, 3] + assert len(core_calls) == 1 + assert core_calls[0]["data"] == [{"Description": "first"}, {"Description": "second"}] + assert core_calls[0]["settings"]["threads"] == 3 + assert core_calls[0]["settings"]["cache"] is False + if kind == "true_false": + assert result.columns.tolist() == ["Description", "Ignore", "First", "First_true_criteria", + "Second", "Second_true_criteria"] + assert result["First"].tolist() == [0.75, 0.75] + assert result["First_true_criteria"].tolist() == ["Suitable outdoors"] * 2 + else: + assert result.columns.tolist() == ["Description", "Ignore", "First", "First_confidence", + "First_probabilities", "Second", "Second_confidence", + "Second_probabilities"] + assert result["First"].tolist() == (["bearing"] * 2 if kind == "choose" else [1.42] * 2) + assert isinstance(result["First_probabilities"].iloc[0], dict) + + +@pytest.mark.parametrize("blank", [None, "", " "]) +def test_whole_blank_output_uses_default_columns(blank, core_calls): + result = run("choose", {"Category": question("choose", output=blank)}) + assert result["Category"].tolist() == ["bearing", "bearing"] + assert "Category_probabilities" in result + + +def test_mixed_questions_project_explicit_columns_and_native_values(core_calls): + result = run("questions", { + "Category": question("choose", type="choose", output=["Category Name", "Certainty", "Distribution"]), + "Severity": question("score", type="score"), + "Outdoor": question("true_false", type="true_false", output=["Outdoor Probability", "Criterion"]), + }) + + assert core_calls[0]["kind"] is None + assert list(core_calls[0]["questions"]) == ["Category", "Severity", "Outdoor"] + assert result["Category Name"].tolist() == ["bearing"] * 2 + assert result["Certainty"].tolist() == [0.8] * 2 + assert result["Severity"].tolist() == [1.42] * 2 + assert result["Severity_probabilities"].iloc[0] == {"Low": 0.0, "Medium": 0.58, "High": 0.42} + assert result["Outdoor Probability"].tolist() == [0.75] * 2 + assert result["Criterion"].tolist() == ["Suitable outdoors"] * 2 + assert "Outdoor" not in result + + +def test_true_false_without_criteria_has_blank_criteria_output(core_calls): + result = run("true_false", {"Outdoor": {"instructions": "Can it be used outdoors?"}}) + assert result["Outdoor"].tolist() == [0.75, 0.75] + assert result["Outdoor_true_criteria"].tolist() == ["", ""] + + +def test_structured_true_criteria_stays_structured(core_calls): + criteria = {"requirements": ["water resistant", "UV resistant"]} + result = run("true_false", {"Outdoor": question("true_false", criteria={"true": criteria})}) + assert result["Outdoor_true_criteria"].tolist() == [criteria, criteria] + + +def test_where_preserves_untouched_rows_and_overwrites_nested_destinations(core_calls): + source = pd.DataFrame({"Description": ["first", "second"], "Category": ["old1", "old2"]}, index=[5, 9]) + result = run("questions", { + "Category": question("choose", type="choose"), + "Outdoor": question("true_false", type="true_false"), + }, dataframe=source, where="Description = 'second'") + + assert core_calls[0]["data"] == [{"Description": "second"}] + assert result.index.tolist() == [5, 9] + assert result["Category"].tolist() == ["old1", "bearing"] + assert result["Category_confidence"].tolist() == ["", 0.8] + assert result["Outdoor"].tolist() == ["", 0.75] + assert result["Outdoor_true_criteria"].tolist() == ["", "Suitable outdoors"] + + +def test_where_no_rows_creates_all_nested_destinations_without_calling_core(core_calls): + source = pd.DataFrame({"Description": ["first"], "Category": ["kept"]}) + result = run("questions", { + "Category": question("choose", type="choose"), + "Outdoor": question("true_false", type="true_false", output=["Probability", "Rule"]), + }, dataframe=source, where="Description = 'missing'") + + assert core_calls == [] + assert result.columns.tolist() == ["Description", "Category", "Category_confidence", + "Category_probabilities", "Probability", "Rule"] + assert result["Category"].tolist() == ["kept"] + assert result["Probability"].tolist() == [""] + + +def test_empty_dataframe_has_declared_schema_without_calling_core(core_calls): + result = run("score", {"Severity": question("score")}, dataframe=pd.DataFrame({"Description": []})) + assert core_calls == [] + assert result.empty + assert result.columns.tolist() == ["Description", "Severity", "Severity_confidence", "Severity_probabilities"] + + +@pytest.mark.parametrize("where", ["Description = 'second'", "Description = 'missing'"]) +def test_nested_output_names_are_literal_even_when_they_contain_wildcards(where, core_calls): + source = pd.DataFrame({"Description": ["first", "second"], "OutdoorElse": ["kept1", "kept2"]}) + result = run("true_false", {"Outdoor*": question("true_false")}, dataframe=source, where=where) + assert "Outdoor*" in result + assert "Outdoor*_true_criteria" in result + assert result["Outdoor*"].tolist() == (["", 0.75] if core_calls else ["", ""]) + assert result["OutdoorElse"].tolist() == ["kept1", "kept2"] + + +def test_duplicate_destinations_rejected_even_when_where_matches_no_rows(core_calls): + with pytest.raises(ValueError, match="(?i)(duplicate|unique|collision|already)"): + run("choose", {"First": question("choose", output=["Same", "Certainty1", "Probs1"]), + "Second": question("choose", output=["Same", "Certainty2", "Probs2"])}, + where="Description = 'missing'") + assert core_calls == [] + + +def test_concurrent_collects_nested_outputs_and_ordinary_wrangle_output(core_calls): + source = pd.DataFrame({"Description": ["first", "second"], "Severity": [99, 99]}) + result = wrangles.recipe.run({"wrangles": [{"concurrent": {"wrangles": [ + {"ai.choose": {"input": "Description", "questions": {"Category": question("choose")}}}, + {"ai.questions": {"input": "Description", "questions": { + "Severity": question("score", type="score"), + "Outdoor": question("true_false", type="true_false"), + }}}, + {"copy": {"input": "Description", "output": "Copied"}}, + ]}}]}, dataframe=source) + + assert len(core_calls) == 2 + assert result["Category"].tolist() == ["bearing"] * 2 + assert result["Severity"].tolist() == [1.42] * 2 + assert result["Outdoor"].tolist() == [0.75] * 2 + assert result["Copied"].tolist() == ["first", "second"] + assert set(result) == {"Description", "Severity", "Category", "Category_confidence", + "Category_probabilities", "Severity_confidence", "Severity_probabilities", + "Outdoor", "Outdoor_true_criteria", "Copied"} + + +def test_dataframe_accessor_uses_recipe_wrapper_without_mutating_source(core_calls): + source = wrangles.DataFrame({"Description": ["first"]}) + result = source.wrangles.ai.score(input="Description", questions={"Severity": question("score")}) + assert result["Severity"].tolist() == [1.42] + assert source.columns.tolist() == ["Description"] + + +def test_wrapper_rejects_response_count_mismatch(monkeypatch): + monkeypatch.setattr(wrangles.ai, "choose", lambda *args, **kwargs: []) + with pytest.raises(RuntimeError, match="response count"): + run("choose", {"Category": question("choose")}) + + +@pytest.mark.parametrize("kind", ["choose", "score", "true_false", "questions"]) +def test_schema_exposes_public_parameters_and_accepts_named_questions(kind): + function = getattr(recipe_ai, kind) + schema = yaml.safe_load(function.__doc__) + jsonschema.Draft202012Validator.check_schema(schema) + assert set(inspect.signature(function).parameters) - {"df"} == set(schema["properties"]) + definition = question("choose" if kind == "questions" else kind) + if kind == "questions": + definition["type"] = "choose" + jsonschema.validate({"questions": {"Question": definition}}, schema) + definition["output"] = ["wrong length"] + with pytest.raises(jsonschema.ValidationError): + jsonschema.validate({"questions": {"Question": definition}}, schema) + + +def test_declared_output_helper_preserves_ordinary_wrangle_shapes(): + helper = wrangles.recipe._declared_output_columns + assert helper("copy", {"output": "Copied"}) == ["Copied"] + assert helper("split.dictionary", {"output": [{"choice": "Category"}, "confidence"]}) == ["Category", "confidence"] + assert helper("extract.ai", {"output": {"Field": {"type": "string"}}}) == ["Field"] + assert helper("custom.dynamic", {}) is None + + +@pytest.fixture +def http_transport(monkeypatch): + """Exercise the complete runtime, intercepting only the provider HTTP call.""" + for name in list(os.environ): + if name.startswith(("WRANGLES_AI_CACHE_", "WRANGLES_EXTRACT_AI_CACHE_")): + monkeypatch.delenv(name) + monkeypatch.delenv("WRANGLES_AI_CONFIG", raising=False) + monkeypatch.setenv("TYPESAFE_API_KEY", "synthetic-recipe-integration-key") + wrangles.ai_config.clear_cache() + wrangles.ai_cache.clear() + calls = [] + + class Response: + status_code = 200 + headers = {} + + def __init__(self, payload): + self.payload = payload + + def json(self): + return self.payload + + def close(self): + pass + + def post(url, **kwargs): + calls.append({"url": url, **copy.deepcopy(kwargs)}) + request = kwargs["json"] + first_row = request["state"]["Description"] == "first" + answers = {} + for label, definition in reversed(list(request["questions"].items())): + kind = definition["type"] + if kind == "choice": + answers[label] = { + "type": "choice", "choice": "bearing" if first_row else "belt", + "confidence": 0.8 if first_row else 0.9, + "probabilities": {"belt": 0.2 if first_row else 0.85, + "bearing": 0.8 if first_row else 0.15}, + } + elif kind == "score": + answers[label] = { + "type": "score", "score": 1.42 if first_row else 0.2, + "confidence": 0.7, + "legend": {"2": "High", "0": "Low", "1": "Medium"}, + "probabilities": {"2": 0.42 if first_row else 0.0, + "0": 0.0 if first_row else 0.8, + "1": 0.58 if first_row else 0.2}, + } + else: + answers[label] = {"type": "noul", "noul": 0.0 if first_row else 0.6} + return Response({"model": request["model"], "answers": answers, + "usage": {"input_tokens": 24, "output_tokens": 12}}) + + monkeypatch.setattr(typesafe._requests, "post", post) + yield calls + wrangles.ai_config.clear_cache() + wrangles.ai_cache.clear() + + +def test_mixed_recipe_through_http_preserves_native_columns_and_provider_payload(http_transport): + definitions = { + "Category": question("choose", type="choose", output=["Picked", "Certainty", "Distribution"]), + "Severity": question("score", type="score"), + "Outdoor": {"type": "true_false", "instructions": "Can it be used outdoors?"}, + } + source = pd.DataFrame({"Description": ["first", "second"], "Excluded": [1, 2]}, index=[12, 4]) + result = run("questions", definitions, dataframe=source, cache=False, threads=1) + + assert result.index.tolist() == [12, 4] + assert result["Picked"].tolist() == ["bearing", "belt"] + assert result["Certainty"].tolist() == [0.8, 0.9] + assert result["Distribution"].tolist() == [{"bearing": 0.8, "belt": 0.2}, + {"bearing": 0.15, "belt": 0.85}] + assert result["Severity"].tolist() == [1.42, 0.2] + assert result["Severity_probabilities"].tolist() == [ + {"Low": 0.0, "Medium": 0.58, "High": 0.42}, + {"Low": 0.8, "Medium": 0.2, "High": 0.0}, + ] + assert result["Outdoor"].tolist() == [0.0, 0.6] + assert result["Outdoor_true_criteria"].tolist() == ["", ""] + assert len(http_transport) == 2 + assert [call["json"]["state"] for call in http_transport] == [ + {"Description": "first"}, {"Description": "second"}, + ] + for call in http_transport: + assert call["url"] == "https://api.typesafe.ai/v1/systemone" + assert call["headers"]["Authorization"] == "Bearer synthetic-recipe-integration-key" + assert set(call["json"]) == {"state", "model", "questions"} + assert call["json"]["questions"] == { + "Category": {"type": "choice", "instructions": definitions["Category"]["instructions"], + "criteria": {"bearing": "A bearing", "belt": "A belt"}}, + "Severity": {"type": "score", "instructions": definitions["Severity"]["instructions"], + "criteria": ["Low", "Medium", "High"]}, + "Outdoor": {"type": "noul", "instructions": "Can it be used outdoors?"}, + } + + +def test_homogeneous_recipe_through_http_batches_both_named_questions_per_row(http_transport): + definitions = { + "Category": question("choose"), + "Alternate": question("choose", output=["Alternative", "Alternative Confidence", "Alternative Probabilities"]), + } + result = run("choose", definitions, cache=False, threads=1) + + assert len(http_transport) == 2 + assert result["Category"].tolist() == ["bearing", "belt"] + assert result["Alternative"].tolist() == ["bearing", "belt"] + assert result["Category_probabilities"].tolist() == result["Alternative Probabilities"].tolist() + assert result["Category_confidence"].tolist() == [0.8, 0.9] + for call in http_transport: + assert set(call["json"]) == {"state", "model", "questions"} + assert set(call["json"]["questions"]) == {"Category", "Alternate"} + for definition in call["json"]["questions"].values(): + assert definition == { + "type": "choice", "instructions": "Evaluate the supplied product.", + "criteria": {"bearing": "A bearing", "belt": "A belt"}, + } diff --git a/tests/test_ai.py b/tests/test_ai.py new file mode 100644 index 00000000..ccac78dc --- /dev/null +++ b/tests/test_ai.py @@ -0,0 +1,299 @@ +"""Offline contracts for the shared AI questions runtime.""" +import copy +import os +import traceback + +import pytest +import yaml + +from wrangles import ai, ai_cache, ai_config +from wrangles.clients import typesafe + + +SCORE = {"severity": {"instructions": "How severe?", "criteria": ["Cosmetic", "Workaround", "Blocking"]}} +CHOICE = {"team": {"instructions": "Which team?", "criteria": {"Support": None, "Billing": "Payments"}}} +NOUL = {"repeat": {"instructions": "Has this happened before?"}} + + +@pytest.fixture(autouse=True) +def isolate(monkeypatch): + for name in list(os.environ): + if name.startswith(("WRANGLES_AI_CACHE_", "WRANGLES_EXTRACT_AI_CACHE_")): + monkeypatch.delenv(name) + monkeypatch.delenv("TYPESAFE_API_KEY", raising=False) + monkeypatch.delenv("WRANGLES_AI_CONFIG", raising=False) + ai_config.clear_cache() + ai_cache.clear() + yield + ai_config.clear_cache() + ai_cache.clear() + + +@pytest.fixture +def transport(monkeypatch): + calls = [] + + def call(**kwargs): + calls.append(copy.deepcopy(kwargs)) + answers = {} + # Deliberately return a different question ordering than the request. + for label, question in reversed(list(kwargs["questions"].items())): + kind = question["type"] + if kind == "score": + answers[label] = {"type": kind, "score": 1.42, "confidence": 0.37, + "probabilities": {"2": 0.42, "0": 0.0, "1": 0.58}, + "legend": dict(enumerate(question["criteria"]))} + elif kind == "choice": + options = list(question["criteria"]) + answers[label] = {"type": kind, "choice": options[0], "confidence": 0.0, + "probabilities": {option: 1.0 / len(options) for option in options}} + else: + answers[label] = {"type": kind, "noul": 0.0 if kwargs["state"] == "zero" else 0.94} + return {"model": kwargs["model"], "answers": answers, "usage": {"input_tokens": 10}} + + monkeypatch.setattr(typesafe, "call_systemone", call) + return calls + + +def save_config(config, monkeypatch, tmp_path): + path = tmp_path / "catalog.yml" + path.write_text(yaml.safe_dump(config), encoding="utf-8") + monkeypatch.setenv("WRANGLES_AI_CONFIG", str(path)) + ai_config.clear_cache() + + +def test_named_mixed_questions_share_one_request_and_preserve_native_values(transport): + questions = { + **{key: {**value, "type": "score"} for key, value in SCORE.items()}, + **{key: {**value, "type": "choose"} for key, value in CHOICE.items()}, + **{key: {**value, "type": "true_false"} for key, value in NOUL.items()}, + } + questions["severity"]["output"] = ["Result", "Certainty", "Distribution"] + original = copy.deepcopy(questions) + result = ai.questions({"Description": "broken"}, questions, api_key="synthetic-key") + assert len(transport) == 1 + assert list(result) == ["severity", "team", "repeat"] + assert result["severity"] == { + "score": 1.42, "confidence": 0.37, + "probabilities": {"Cosmetic": 0.0, "Workaround": 0.58, "Blocking": 0.42}, + } + assert result["team"] == {"choice": "Support", "confidence": 0.0, + "probabilities": {"Support": 0.5, "Billing": 0.5}} + assert result["repeat"] == {"probability_true": 0.94, "true_criteria": ""} + assert questions == original + assert transport[0]["questions"]["team"]["type"] == "choice" + assert transport[0]["questions"]["repeat"]["type"] == "noul" + assert all("output" not in q for q in transport[0]["questions"].values()) + + +@pytest.mark.parametrize("method,definitions", [("choose", CHOICE), ("score", SCORE), ("true_false", NOUL)]) +def test_individual_wrangles_support_multiple_questions_in_one_request(transport, method, definitions): + questions = {"first": copy.deepcopy(next(iter(definitions.values()))), + "second": copy.deepcopy(next(iter(definitions.values())))} + result = getattr(ai, method)("state", questions, api_key="synthetic-key") + assert list(result) == ["first", "second"] + assert len(transport) == 1 + assert len(transport[0]["questions"]) == 2 + + +def test_batch_shape_order_duplicate_suppression_and_independent_results(transport): + result = ai.true_false(["zero", "positive", "zero"], NOUL, api_key="synthetic-key") + assert [row["repeat"]["probability_true"] for row in result] == [0.0, 0.94, 0.0] + assert len(transport) == 2 + result[0]["repeat"]["probability_true"] = 1 + assert result[2]["repeat"]["probability_true"] == 0 + assert ai.true_false("zero", NOUL, api_key="synthetic-key")["repeat"]["probability_true"] == 0 + assert len(transport) == 2 + + +def test_true_criterion_is_supplied_context_and_not_a_generated_explanation(transport): + definitions = {"repeat": {**NOUL["repeat"], "criteria": { + "true": {"description": "Prior contact", "examples": ["Called twice"]}, + "false": "First contact", + }}} + result = ai.true_false("positive", definitions, api_key="synthetic-key") + assert result["repeat"] == {"probability_true": 0.94, "true_criteria": definitions["repeat"]["criteria"]["true"]} + result["repeat"]["true_criteria"]["examples"].append("mutation") + assert definitions["repeat"]["criteria"]["true"]["examples"] == ["Called twice"] + + +@pytest.mark.parametrize("blank", [None, "", " "]) +def test_default_names_and_empty_output_forms(blank): + assert ai._output_columns({"severity": {**SCORE["severity"], "output": blank}}, "score") == [ + "severity", "severity_confidence", "severity_probabilities"] + assert ai._output_columns({"repeat": {**NOUL["repeat"], "output": blank}}, "true_false") == [ + "repeat", "repeat_true_criteria"] + + +def test_explicit_output_names_and_projection_validation(): + assert ai._output_columns({"severity": {**SCORE["severity"], "output": ["A", "B", "C"]}}, "score") == ["A", "B", "C"] + assert ai._output_columns({"repeat": {**NOUL["repeat"], "output": ["A", "B"]}}, "true_false") == ["A", "B"] + + +@pytest.mark.parametrize("output", [["one"], ["a", "b"], ["a", "b", ""], ["a", "a", "c"], [], "name", ["a", "b", "c", "d"]]) +def test_invalid_positional_outputs_fail_before_calls(transport, output): + with pytest.raises(ValueError): + ai.score("state", {"severity": {**SCORE["severity"], "output": output}}) + assert transport == [] + + +def test_generated_and_explicit_names_must_not_collide(): + definitions = {"q": NOUL["repeat"], "q_true_criteria": NOUL["repeat"]} + with pytest.raises(ValueError, match="unique"): + ai._prepare_questions(definitions, "true_false") + + +@pytest.mark.parametrize("definitions,kind", [ + ({}, "score"), + ({"q": {"instructions": "score", "criteria": ["same", "same"]}}, "score"), + ({"q": {"instructions": "score", "criteria": ["only"]}}, "score"), + ({"q": {"instructions": "score", "criteria": list(map(str, range(11)))}}, "score"), + ({"q": {"instructions": "score", "criteria": ["low", {"high": "very"}]}}, "score"), + ({"q": {"instructions": "choose", "criteria": {}}}, "choose"), + ({"q": {"instructions": "choose", "criteria": {str(i): None for i in range(256)}}}, "choose"), + ({"q": {"instructions": "choose", "criteria": {True: None}}}, "choose"), + ({"q": {"instructions": "noul", "criteria": {True: "yes"}}}, "true_false"), + ({"q": {"instructions": "noul", "criteria": {"other": "yes"}}}, "true_false"), + ({"q": {"instructions": "noul", "criteria": None}}, "true_false"), + ({"q": {"instructions": "noul", "type": "score"}}, "true_false"), + ({"q": {"instructions": "question"}}, None), + ({"q": {"instructions": "question", "type": "noul"}}, None), + ({"q": {"instructions": "question", "type": [], "criteria": []}}, None), + ({"": {"instructions": "question"}}, "true_false"), + ({"q": {"instructions": " "}}, "true_false"), + ({"q": {"instructions": "question", "scale": 100}}, "true_false"), + ({"q": {"instructions": {"value": float("nan")}}}, "true_false"), +]) +def test_invalid_questions_fail_before_credentials_or_provider(transport, definitions, kind): + function = getattr(ai, kind or "questions") + with pytest.raises(ValueError): + function("state", definitions) + assert transport == [] + + +def test_json_descriptions_and_boundary_criteria_are_preserved(): + definitions = {"q": {"instructions": {"question": "Choose", "context": [1, None]}, + "criteria": {str(i): None for i in range(255)}}} + prepared = ai._prepare_questions(definitions, "choose") + assert len(prepared["q"]["criteria"]) == 255 + assert prepared["q"]["instructions"] == definitions["q"]["instructions"] + assert len(ai._prepare_questions({"q": {"instructions": "Score", "criteria": list(map(str, range(10)))}}, "score")["q"]["criteria"]) == 10 + + +@pytest.mark.parametrize("row", [42, None, {"value": float("nan")}, {1: "non-string-key"}, {"value": object()}]) +def test_all_rows_validated_before_sending_any(transport, row): + with pytest.raises(ValueError): + ai.true_false(["valid first", row], NOUL, api_key="synthetic-key") + assert transport == [] + + +def test_empty_batch_needs_no_credentials_or_catalog(transport, monkeypatch): + monkeypatch.setattr(ai_config, "resolve", lambda *a, **kw: pytest.fail("No work needs no config")) + assert ai.score([], SCORE) == [] + assert transport == [] + + +def test_config_resolves_once_with_explicit_zero_and_false_overrides(transport, monkeypatch, tmp_path): + config = ai_config.load() + config["operations"]["ai.true_false"]["defaults"].update(default_concurrency=2, request_timeout_seconds=7.5, retries=3) + config["providers"]["typesafe"]["endpoints"]["systemone"] = "https://typesafe.example/evaluate" + save_config(config, monkeypatch, tmp_path) + original = ai_config.resolve + resolutions = [] + def resolve(*args, **kwargs): + resolutions.append((args, kwargs)) + return original(*args, **kwargs) + monkeypatch.setattr(ai_config, "resolve", resolve) + ai.true_false(["a", "a"], NOUL, api_key="synthetic-key", model="custom-jev", retries=0, cache=False, timeout=2.5, threads=1) + assert len(resolutions) == 1 + assert len(transport) == 2 + assert all(call["model"] == "custom-jev" and call["retries"] == 0 and call["timeout"] == 2.5 + and call["url"] == "https://typesafe.example/evaluate" for call in transport) + + +def test_local_environment_key_and_explicit_key_precedence(transport, monkeypatch): + monkeypatch.setenv("TYPESAFE_API_KEY", "environment-test-key") + ai.true_false("one", NOUL) + ai.true_false("two", NOUL, api_key="explicit-test-key") + assert [call["api_key"] for call in transport] == ["environment-test-key", "explicit-test-key"] + with pytest.raises(ValueError, match="api_key"): + ai.true_false("three", NOUL, api_key="") + + +@pytest.mark.parametrize("overrides", [ + {"threads": 0}, {"threads": True}, {"timeout": 0}, {"timeout": float("nan")}, + {"timeout": 10 ** 1000}, {"cache_ttl": 10 ** 1000}, + {"retries": -1}, {"retries": True}, {"cache": "yes"}, {"cache_ttl": float("inf")}, + {"provider": "openai"}, {"protocol": "responses"}, +]) +def test_invalid_runtime_settings_fail_before_requests(transport, overrides): + with pytest.raises(ValueError): + ai.true_false("state", NOUL, api_key="synthetic-key", **overrides) + assert transport == [] + + +def test_unusable_endpoint_error_does_not_disclose_credentials(transport, monkeypatch, tmp_path): + config = ai_config.load() + config["providers"]["typesafe"]["endpoints"]["systemone"] = "https://user:private-value@example.com" + save_config(config, monkeypatch, tmp_path) + with pytest.raises(ValueError) as exc: + ai.true_false("private-state", NOUL, api_key="private-key") + assert "private" not in str(exc.value) + assert transport == [] + + +def test_typesafe_cache_controls_are_independent_of_extraction(transport, monkeypatch): + monkeypatch.setenv("WRANGLES_EXTRACT_AI_CACHE_ENABLED", "false") + for _ in range(2): + ai.true_false("state", NOUL, api_key="synthetic-key") + assert len(transport) == 1 + monkeypatch.setenv("WRANGLES_AI_CACHE_ENABLED", "false") + for _ in range(2): + ai.true_false("state", NOUL, api_key="synthetic-key") + assert len(transport) == 3 + monkeypatch.setenv("WRANGLES_EXTRACT_AI_CACHE_ENABLED", "true") + assert ai_cache.resolve_policy({}).enabled is True + + +def test_invalid_cache_environment_value_is_not_exposed_in_traceback(transport, monkeypatch): + monkeypatch.setenv("WRANGLES_AI_CACHE_TTL_SECONDS", "private-invalid-value") + with pytest.raises(ValueError) as exc: + ai.true_false("state", NOUL, api_key="synthetic-key") + assert "private-invalid-value" not in "".join(traceback.format_exception(exc.value)) + assert transport == [] + + +def test_cache_identity_includes_key_model_endpoint_state_and_questions(transport, monkeypatch, tmp_path): + ai.true_false("a", NOUL, api_key="key-a") + ai.true_false("b", NOUL, api_key="key-a") + ai.true_false("a", NOUL, api_key="key-b") + ai.true_false("a", NOUL, api_key="key-a", model="other-model") + ai.true_false("a", {"repeat": {"instructions": "A different question?"}}, api_key="key-a") + config = ai_config.load() + config["providers"]["typesafe"]["endpoints"]["systemone"] = "https://other.example/systemone" + save_config(config, monkeypatch, tmp_path) + ai.true_false("a", NOUL, api_key="key-a") + # Output names only affect local projection, so renaming may reuse the answer. + ai.true_false("a", {"repeat": {**NOUL["repeat"], "output": ["new", "new_criteria"]}}, api_key="key-a") + assert len(transport) == 6 + + +def test_failures_are_not_cached(monkeypatch): + attempts = [] + def fail(**kwargs): + attempts.append(1) + raise RuntimeError("synthetic transport failure") + monkeypatch.setattr(typesafe, "call_systemone", fail) + for _ in range(2): + with pytest.raises(RuntimeError, match="synthetic"): + ai.true_false("state", NOUL, api_key="synthetic-key") + assert len(attempts) == 2 + assert ai_cache.stats()["entries"] == 0 + + +def test_deprecated_model_warning_is_once_per_invocation(transport, monkeypatch, tmp_path, caplog): + config = ai_config.load() + config["providers"]["typesafe"]["models"]["jev-1.13.0"]["status"] = "deprecated" + save_config(config, monkeypatch, tmp_path) + ai.true_false(["a", "b", "c"], NOUL, api_key="synthetic-key") + assert sum("deprecated status" in record.message for record in caplog.records) == 1 diff --git a/tests/test_ai_config.py b/tests/test_ai_config.py index dcb0e1a5..f7a586ec 100644 --- a/tests/test_ai_config.py +++ b/tests/test_ai_config.py @@ -59,6 +59,71 @@ def test_catalog_metadata_is_separate_from_request_defaults(): assert "documentation" not in policy +@pytest.mark.parametrize("operation", ["ai.choose", "ai.score", "ai.true_false", "ai.questions"]) +def test_typesafe_operations_resolve_independent_packaged_defaults(operation): + policy = ai_config.resolve(operation) + assert policy["provider"] == "typesafe" + assert policy["protocol"] == "systemone" + assert policy["model"] == "jev-1.13.0" + assert policy["endpoints"] == {"systemone": "https://api.typesafe.ai/v1/systemone"} + assert policy["default_concurrency"] == 10 + assert policy["request_timeout_seconds"] == 30 + assert policy["retries"] == 1 + assert policy["cache"] == { + "enabled": True, "ttl_seconds": 3600, "max_entries": 512, + "max_value_bytes": 65536, "single_flight": True, "log_every": 0, + } + assert "reasoning" not in policy + assert "text" not in policy + for role in ("global", "test"): + with pytest.raises(ValueError, match="cannot select"): + ai_config.resolve(operation, role=role) + + +@pytest.mark.parametrize("operation", ["ai.choose", "ai.score", "ai.true_false", "ai.questions"]) +def test_typesafe_explicit_unlisted_models_preserve_provider_settings(operation): + policy = ai_config.resolve(operation, provider=" TYPESAFE ", model="jev-custom") + assert policy["model"] == "jev-custom" + assert policy["provider"] == "typesafe" + assert policy["protocol"] == "systemone" + assert policy["endpoints"]["systemone"] == "https://api.typesafe.ai/v1/systemone" + with pytest.raises(ValueError, match="No AI provider"): + ai_config.resolve(operation, provider="unconfigured", model="jev-custom") + with pytest.raises(ValueError, match="specify model explicitly"): + ai_config.resolve(operation, provider="openai") + + +def test_typesafe_custom_model_defaults_are_provider_local(monkeypatch, tmp_path): + config = ai_config.load() + config["providers"]["typesafe"]["models"]["jev-1.13.0"]["default_for"].remove("ai.score") + config["providers"]["typesafe"]["models"]["jev-custom-score"] = { + "status": "active", "default_for": ["ai.score"], + } + config["providers"]["typesafe"]["endpoints"]["systemone"] = "https://custom.example/systemone" + use_config(config, monkeypatch, tmp_path) + assert ai_config.resolve("ai.score")["model"] == "jev-custom-score" + assert ai_config.resolve("ai.choose")["model"] == "jev-1.13.0" + assert ai_config.resolve("extract.ai")["model"] == "gpt-6-luna" + assert ai_config.resolve("ai.score")["endpoints"]["systemone"] == "https://custom.example/systemone" + + +@pytest.mark.parametrize("operation", ["ai.choose", "ai.score", "ai.true_false", "ai.questions"]) +def test_version_1_overrides_keep_packaged_typesafe_operations(monkeypatch, tmp_path, operation): + packaged = ai_config.resolve(operation) + use_config({"version": 1, "extract_ai": {"model": "legacy-custom"}}, monkeypatch, tmp_path) + assert ai_config.resolve(operation) == packaged + + +@pytest.mark.parametrize("operation", ["ai.choose", "ai.score", "ai.true_false", "ai.questions"]) +def test_version_2_does_not_merge_missing_typesafe_operations(monkeypatch, tmp_path, operation): + config = ai_config.load() + del config["operations"][operation] + config["providers"]["typesafe"]["models"]["jev-1.13.0"]["default_for"].remove(operation) + use_config(config, monkeypatch, tmp_path) + with pytest.raises(ValueError, match="No AI operation"): + ai_config.resolve(operation) + + def test_generic_task_operation_requires_explicit_model(): with pytest.raises(ValueError, match="requires an explicit model"): ai_config.resolve("huggingface") diff --git a/tests/test_typesafe.py b/tests/test_typesafe.py new file mode 100644 index 00000000..895b5c44 --- /dev/null +++ b/tests/test_typesafe.py @@ -0,0 +1,269 @@ +"""Credential-free TypeSafe wire contracts; all HTTP calls are mocked.""" + +import copy +import traceback +from datetime import datetime, timedelta, timezone +from email.utils import format_datetime + +import pytest +import requests + +from wrangles.clients import typesafe + + +QUESTIONS = { + "route": {"type": "choice", "instructions": "Select a team", "criteria": {"billing": None, "sales": "Quotes"}}, + "urgency": {"type": "score", "instructions": "Rate urgency", "criteria": ["Normal", "Urgent", "Critical"]}, + "duplicate": {"type": "noul", "instructions": "Is this a duplicate?"}, +} +RESPONSE = { + "model": "jev-1.13.0", + "answers": { + "route": {"type": "choice", "choice": "billing", "confidence": 0.8, "probabilities": {"billing": 0.9, "sales": 0.1}}, + "urgency": {"type": "score", "score": 1.05, "confidence": 0.92, "legend": {"0": "Normal", "1": "Urgent", "2": "Critical"}, "probabilities": {"0": 0, "1": 0.95, "2": 0.05}}, + "duplicate": {"type": "noul", "noul": 0.7}, + }, + "usage": {"input_tokens": 42, "output_tokens": 8}, +} + + +class Response: + def __init__(self, payload=None, status=200, headers=None, error=None): + self.payload = copy.deepcopy(RESPONSE if payload is None else payload) + self.status_code = status + self.headers = headers or {} + self.error = error + self.closed = False + + def json(self): + if self.error is not None: + raise self.error + return self.payload + + def close(self): + self.closed = True + + +@pytest.fixture +def transport(monkeypatch): + calls, sleeps, results = [], [], [] + + def post(url, **kwargs): + calls.append({"url": url, **copy.deepcopy(kwargs)}) + result = results.pop(0) + if isinstance(result, Exception): + raise result + return result + + monkeypatch.setattr(typesafe._requests, "post", post) + monkeypatch.setattr(typesafe._time, "sleep", sleeps.append) + return calls, sleeps, results + + +def call(**overrides): + arguments = dict( + state={"description": "Synthetic invoice question"}, questions=copy.deepcopy(QUESTIONS), + model="jev-1.13.0", api_key="synthetic-private-key", url="https://typesafe.example/v1/systemone", + timeout=7.5, retries=1, + ) + return typesafe.call_systemone(**(arguments | overrides)) + + +def test_typed_response_and_request_contract_are_preserved(transport): + calls, sleeps, results = transport + response = Response() + results.append(response) + state, questions = {"description": "Synthetic"}, copy.deepcopy(QUESTIONS) + result = call(state=state, questions=questions) + assert result is response.payload + assert result == RESPONSE + assert calls == [{ + "url": "https://typesafe.example/v1/systemone", "timeout": 7.5, "allow_redirects": False, + "headers": {"Authorization": "Bearer synthetic-private-key", "Content-Type": "application/json"}, + "json": {"state": state, "model": "jev-1.13.0", "questions": QUESTIONS}, + }] + assert questions == QUESTIONS + assert sleeps == [] + assert response.closed + + +@pytest.mark.parametrize("status", [429, 500, 502, 503, 529]) +@pytest.mark.parametrize("retries", [0, 1, 2]) +def test_transient_status_retry_budget(status, retries, transport): + calls, sleeps, results = transport + responses = [Response(status=status) for _ in range(retries + 1)] + results.extend(responses) + with pytest.raises(typesafe.TypesafeError) as caught: + call(retries=retries) + assert caught.value.status_code == status + assert caught.value.attempts == retries + 1 + assert len(calls) == retries + 1 + assert sleeps == [2 ** attempt for attempt in range(retries)] + assert all(response.closed for response in responses) + + +@pytest.mark.parametrize("status", [301, 307, 400, 401, 403, 404, 422]) +def test_permanent_status_and_redirects_fail_without_retry(status, transport): + calls, sleeps, results = transport + response = Response(status=status) + results.append(response) + with pytest.raises(typesafe.TypesafeError) as caught: + call(retries=3) + assert caught.value.status_code == status + assert len(calls) == 1 + assert sleeps == [] + assert response.closed + + +@pytest.mark.parametrize("error_type", [requests.Timeout, requests.ConnectionError, requests.exceptions.ChunkedEncodingError, requests.exceptions.ContentDecodingError]) +def test_transient_transport_error_retries_and_recovers(error_type, transport): + calls, sleeps, results = transport + response = Response() + results.extend([error_type("SECRET FROM TRANSPORT"), response]) + assert call() == RESPONSE + assert len(calls) == 2 + assert [item["timeout"] for item in calls] == [7.5, 7.5] + assert sleeps == [1] + assert response.closed + + +@pytest.mark.parametrize("error_type", [requests.Timeout, requests.ConnectionError, requests.exceptions.InvalidURL, ValueError]) +def test_transport_failure_cannot_expose_exception_contents(error_type, transport): + calls, sleeps, results = transport + results.extend([error_type("DO-NOT-EXPOSE-secret-and-request-body") for _ in range(3)]) + with pytest.raises(typesafe.TypesafeError) as caught: + call(retries=2) + rendered = "".join(traceback.format_exception(caught.type, caught.value, caught.tb)) + assert "DO-NOT-EXPOSE" not in rendered + assert caught.value.__cause__ is None + assert caught.value.__suppress_context__ + expected = 3 if error_type in (requests.Timeout, requests.ConnectionError) else 1 + assert len(calls) == expected + assert len(sleeps) == expected - 1 + + +@pytest.mark.parametrize("header,expected", [("2.5", 2.5), ("0", 0), ("900", 60), ("-1", 1), ("NaN", 1), ("inf", 1), ("invalid", 1)]) +def test_retry_after_is_honored_and_bounded(header, expected, transport): + calls, sleeps, results = transport + results.extend([Response(status=429, headers={"Retry-After": header}), Response()]) + assert call() == RESPONSE + assert sleeps == [expected] + + +def test_retry_after_http_date_and_backoff_are_bounded(): + future = format_datetime(datetime.now(timezone.utc) + timedelta(seconds=25)) + assert 23 <= typesafe._retry_delay(future, 0) <= 25 + past = format_datetime(datetime.now(timezone.utc) - timedelta(seconds=25)) + assert typesafe._retry_delay(past, 0) == 0 + assert typesafe._retry_delay(None, 100000) == 60 + + +def test_invalid_json_fails_without_retry_or_body_leak(transport): + calls, sleeps, results = transport + response = Response(error=ValueError("DO-NOT-EXPOSE-body")) + results.append(response) + with pytest.raises(typesafe.TypesafeResponseError) as caught: + call(retries=2) + assert "DO-NOT-EXPOSE" not in "".join(traceback.format_exception(caught.type, caught.value, caught.tb)) + assert len(calls) == 1 + assert sleeps == [] + assert response.closed + + +@pytest.mark.parametrize("modify", [ + lambda p: p.update(model=""), + lambda p: p.update(model=None), + lambda p: p.pop("usage"), + lambda p: p["usage"].update(input_tokens=-1), + lambda p: p["usage"].update(output_tokens=True), + lambda p: p["usage"].update(input_tokens=1.5), + lambda p: p["answers"].pop("duplicate"), + lambda p: p["answers"].update(extra={"type": "noul", "noul": 0.5}), + lambda p: p["answers"].update(duplicate=[]), + lambda p: p["answers"]["duplicate"].update(type="score"), + lambda p: p["answers"]["duplicate"].update(noul=None), + lambda p: p["answers"]["duplicate"].update(noul=True), + lambda p: p["answers"]["duplicate"].update(noul=float("nan")), + lambda p: p["answers"]["duplicate"].update(noul=10 ** 1000), + lambda p: p["answers"]["duplicate"].update(noul=1.01), + lambda p: p["answers"]["route"].update(choice="DO-NOT-EXPOSE-unrequested"), + lambda p: p["answers"]["route"].update(choice=[]), + lambda p: p["answers"]["route"]["probabilities"].pop("sales"), + lambda p: p["answers"]["route"]["probabilities"].update(other=0), + lambda p: p["answers"]["route"]["probabilities"].update(sales=0.5), + lambda p: p["answers"]["route"]["probabilities"].update(sales=-0.1), + lambda p: p["answers"]["route"]["probabilities"].update(sales=float("inf")), + lambda p: p["answers"]["route"]["probabilities"].update(sales=True), + lambda p: p["answers"]["route"].update(confidence="0.9"), + lambda p: p["answers"]["route"].update(confidence=float("nan")), + lambda p: p["answers"]["urgency"].update(score=3), + lambda p: p["answers"]["urgency"].update(score=-0.1), + lambda p: p["answers"]["urgency"].update(score=float("inf")), + lambda p: p["answers"]["urgency"].update(score=True), + lambda p: p["answers"]["urgency"]["probabilities"].update({"3": 0}), + lambda p: p["answers"]["urgency"]["legend"].pop("2"), + lambda p: p["answers"]["urgency"]["legend"].update({"1": "DO-NOT-EXPOSE-wrong-rubric"}), +]) +def test_invalid_typed_answers_fail_without_retry_or_values_in_error(modify, transport): + calls, sleeps, results = transport + payload = copy.deepcopy(RESPONSE) + modify(payload) + response = Response(payload) + results.append(response) + with pytest.raises(typesafe.TypesafeResponseError) as caught: + call(retries=2) + assert "DO-NOT-EXPOSE" not in str(caught.value) + assert len(calls) == 1 + assert sleeps == [] + assert response.closed + + +@pytest.mark.parametrize("usage", [{}, {"input_tokens": None, "output_tokens": 0}]) +def test_optional_usage_is_preserved_without_manufactured_counts(usage, transport): + _, _, results = transport + payload = copy.deepcopy(RESPONSE) + payload["usage"] = usage + results.append(Response(payload)) + assert call()["usage"] == usage + + +def test_rounded_probabilities_are_preserved_not_renormalized(transport): + _, _, results = transport + payload = copy.deepcopy(RESPONSE) + payload["answers"]["route"]["probabilities"] = {"billing": 0.9, "sales": 0.099} + results.append(Response(payload)) + assert call() == payload + + +@pytest.mark.parametrize("payload", [[], "not an object", False]) +def test_non_object_response_is_rejected_and_closed(payload, transport): + _, _, results = transport + response = Response(payload) + results.append(response) + with pytest.raises(typesafe.TypesafeResponseError, match="expected an object"): + call() + assert response.closed + + +def test_failed_status_does_not_read_error_body_and_closes_response(transport): + _, _, results = transport + response = Response(status=401, error=AssertionError("error body must not be read")) + results.append(response) + with pytest.raises(typesafe.TypesafeError, match="HTTP 401"): + call() + assert response.closed + + +def test_error_response_cleanup_does_not_override_sanitized_error(transport): + _, _, results = transport + response = Response(status=401) + + def close(): + raise OSError("DO-NOT-EXPOSE-close-details") + + response.close = close + results.append(response) + with pytest.raises(typesafe.TypesafeError) as caught: + call() + assert "HTTP 401" in str(caught.value) + assert "DO-NOT-EXPOSE" not in str(caught.value) diff --git a/wrangles/__init__.py b/wrangles/__init__.py index 21cf489f..d72bd76b 100644 --- a/wrangles/__init__.py +++ b/wrangles/__init__.py @@ -18,6 +18,7 @@ from .classify import classify from . import extract +from . import ai from .lookup import lookup from .translate import translate from .standardize import standardize @@ -34,4 +35,3 @@ from . import compare from . import generate - diff --git a/wrangles/ai.py b/wrangles/ai.py new file mode 100644 index 00000000..09c62fd9 --- /dev/null +++ b/wrangles/ai.py @@ -0,0 +1,281 @@ +"""Named AI questions over text or structured records. + +The individual operations imply a question type; ``questions`` mixes types. +Each input record is evaluated in one request. Python results are dictionaries +of named answers; recipe wrappers project their fields into columns. +""" + +import copy as _copy +import json as _json +import math as _math +import os as _os +from urllib.parse import urlsplit as _urlsplit + +from . import ai_cache as _cache +from . import ai_config as _config +from .clients import typesafe as _typesafe + + +_KINDS = {"choose": "choice", "score": "score", "true_false": "noul"} +_FIELDS = { + "choose": ("choice", "confidence", "probabilities"), + "score": ("score", "confidence", "probabilities"), + "true_false": ("probability_true", "true_criteria"), +} + + +def _json_value(value, location): + """Validate JSON without quoting private input values in an exception.""" + def check(item): + if item is None or type(item) in (str, bool, int): + return + if type(item) is float and _math.isfinite(item): + return + if isinstance(item, list): + for child in item: + check(child) + return + if isinstance(item, dict) and all(isinstance(key, str) for key in item): + for child in item.values(): + check(child) + return + raise ValueError + + try: + check(value) + # Also rejects circular containers before copying or sending requests. + _json.dumps(value, allow_nan=False) + except (TypeError, ValueError, OverflowError, RecursionError): + raise ValueError(f"{location} must contain finite JSON-compatible values and string object keys.") from None + + +def _description(value, location, *, nullable=False): + if value is None and nullable: + return + if not isinstance(value, (str, dict, list)): + raise ValueError(f"{location} must be text, an object, or an array.") + if not value or (isinstance(value, str) and not value.strip()): + raise ValueError(f"{location} must not be empty.") + _json_value(value, location) + + +def _prepare_questions(questions, kind=None): + """Validate question definitions and resolve their positional output names.""" + if kind is not None and kind not in _KINDS: + raise ValueError("Unknown AI question type.") + if not isinstance(questions, dict) or not questions: + raise ValueError("questions must be a non-empty mapping of labels to definitions.") + prepared = {} + destinations = set() + for label, definition in questions.items(): + if not isinstance(label, str) or not label.strip(): + raise ValueError("Question labels must be non-empty strings.") + if not isinstance(definition, dict): + raise ValueError("Each question must be an object.") + if set(definition) - {"type", "instructions", "criteria", "output"}: + raise ValueError("Question definitions only accept type, instructions, criteria, and output.") + question_kind = definition.get("type", kind) + if not isinstance(question_kind, str) or question_kind not in _KINDS: + raise ValueError("Question type must be choose, score, or true_false.") + if kind is not None and question_kind != kind: + raise ValueError(f"ai.{kind} only accepts {kind} questions.") + instructions = definition.get("instructions") + _description(instructions, "Question instructions") + criteria = definition.get("criteria") + if question_kind == "choose": + if not isinstance(criteria, dict) or not 1 <= len(criteria) <= 255: + raise ValueError("Choose criteria must map between 1 and 255 option labels to descriptions.") + for option, description in criteria.items(): + if not isinstance(option, str) or not option.strip(): + raise ValueError("Choose option labels must be non-empty strings.") + _description(description, "Choose criterion", nullable=True) + elif question_kind == "score": + if (not isinstance(criteria, list) or not 2 <= len(criteria) <= 10 + or any(not isinstance(item, str) or not item.strip() for item in criteria)): + raise ValueError("Score criteria must be an ordered list of 2 to 10 non-empty string descriptions.") + if len(set(criteria)) != len(criteria): + raise ValueError("Score criterion descriptions must be unique probability labels.") + elif "criteria" in definition: + if not isinstance(criteria, dict) or set(criteria) - {"true", "false"}: + raise ValueError('True/false criteria must be an object with quoted "true"/"false" keys.') + for description in criteria.values(): + _description(description, "True/false criterion") + + fields = _FIELDS[question_kind] + names = definition.get("output") + if names is None or (isinstance(names, str) and not names.strip()): + names = [label, *(f"{label}_{field}" for field in fields[1:])] + elif (not isinstance(names, list) or len(names) != len(fields) + or any(not isinstance(name, str) or not name.strip() for name in names)): + raise ValueError(f"{question_kind} output must be blank or a list of {len(fields)} non-empty column names.") + if len(set(names)) != len(names) or destinations.intersection(names): + raise ValueError("AI output column names must be unique across all questions.") + destinations.update(names) + normalized = {"type": question_kind, "instructions": _copy.deepcopy(instructions), + "output": list(names)} + if "criteria" in definition: + normalized["criteria"] = _copy.deepcopy(criteria) + prepared[label] = normalized + return prepared + + +def _output_columns(questions, kind=None): + """Collect destinations before execution, including recipes selecting no rows.""" + return [name for question in _prepare_questions(questions, kind).values() + for name in question["output"]] + + +def _wire_questions(prepared): + return { + label: {"type": _KINDS[question["type"]], "instructions": question["instructions"], + **({"criteria": question["criteria"]} if "criteria" in question else {})} + for label, question in prepared.items() + } + + +def _answers(response, prepared): + result = {} + for label, question in prepared.items(): + answer = response["answers"][label] + kind = question["type"] + if kind == "true_false": + result[label] = { + "probability_true": answer["noul"], + "true_criteria": _copy.deepcopy(question.get("criteria", {}).get("true", "")), + } + else: + value_field = _FIELDS[kind][0] + probabilities = answer["probabilities"] + if kind == "score": + probabilities = {description: probabilities[str(index)] + for index, description in enumerate(question["criteria"])} + else: + probabilities = {option: probabilities[option] for option in question["criteria"]} + result[label] = {value_field: answer[value_field], "confidence": answer["confidence"], + "probabilities": probabilities} + return result + + +def _positive_number(value, name): + try: + valid = type(value) in (int, float) and _math.isfinite(value) and value > 0 + except OverflowError: + valid = False + if not valid: + raise ValueError(f"{name} must be a positive finite number.") + return value + + +def _run(data, questions, kind, *, api_key, model, provider, protocol, threads, + timeout, retries, cache, cache_ttl): + prepared = _prepare_questions(questions, kind) + rows = data if isinstance(data, list) else [data] + for row in rows: + if not isinstance(row, (str, dict, list)): + raise ValueError("Each AI input must be text, a JSON object, or an array.") + _json_value(row, "AI input") + # All input/question validation precedes credentials and any provider request. + if not rows: + return [] + + operation = f"ai.{kind or 'questions'}" + settings = _config.resolve(operation, model=model, provider=provider, protocol=protocol) + if settings["provider"] != "typesafe" or settings["protocol"] != "systemone": + raise ValueError("AI question wrangles currently require provider typesafe and protocol systemone.") + workers = threads if threads is not None else settings.get("default_concurrency") + if type(workers) is not int or workers < 1: + raise ValueError("threads/default_concurrency must be a positive integer.") + request_timeout = _positive_number( + timeout if timeout is not None else settings.get("request_timeout_seconds"), "timeout") + attempts = retries if retries is not None else settings.get("retries") + if type(attempts) is not int or attempts < 0: + raise ValueError("retries must be a non-negative integer.") + url = settings.get("endpoints", {}).get("systemone") + try: + parsed_url = _urlsplit(url) if isinstance(url, str) else None + valid_url = (parsed_url is not None and parsed_url.scheme == "https" and parsed_url.hostname + and not parsed_url.username and not parsed_url.password and not parsed_url.fragment) + except ValueError: + valid_url = False + if not valid_url: + raise ValueError("The Typesafe systemone endpoint must be an HTTPS URL without embedded credentials or fragments.") + + secret = api_key if api_key is not None else _os.getenv("TYPESAFE_API_KEY") + if not isinstance(secret, str) or not secret.strip(): + raise ValueError("Set api_key or TYPESAFE_API_KEY for AI question wrangles.") + secret = secret.strip() + if any(ord(char) < 33 or ord(char) > 126 for char in secret): + raise ValueError("The Typesafe API key must contain only printable ASCII without whitespace.") + if cache_ttl is not None: + _positive_number(cache_ttl, "cache_ttl") + policy = _cache.resolve_policy(settings.get("cache", {}), enabled=cache, ttl_seconds=cache_ttl, + env_prefix="WRANGLES_AI_CACHE") + _positive_number(policy.ttl_seconds, "cache_ttl") + _config.warn_if_deprecated(settings["model"], provider=settings["provider"]) + wire = _wire_questions(prepared) + static_request = {"endpoint": url, "model": settings["model"], "questions": wire} + + def key_for(row): + return _cache.make_key(namespace=operation, provider=settings["provider"], + protocol=settings["protocol"], tenant_secret=secret, + static_request=static_request, data=row) + + def compute(row): + return _typesafe.call_systemone(state=row, questions=wire, model=settings["model"], + api_key=secret, url=url, timeout=request_timeout, retries=attempts) + + responses = _cache.execute_batch(rows, key_for=key_for, compute=compute, cacheable=lambda result: True, + max_workers=workers, policy=policy, preflight_first=True) + results = [_answers(response, prepared) for response in responses] + return results if isinstance(data, list) else results[0] + + +def choose(data, questions, api_key=None, *, model=None, provider=None, protocol=None, + threads=None, timeout=None, retries=None, cache=None, cache_ttl=None): + """Answer named choice questions; return choice, confidence, and probabilities. + + ``data`` is text/a record, or a list of input records. Each question supplies + instructions and a mapping of option labels to descriptions (or None). + A scalar input returns a named-answer dictionary; a list returns an ordered + list of those dictionaries. Model and runtime defaults come from the catalog. + """ + return _run(data, questions, "choose", api_key=api_key, model=model, provider=provider, + protocol=protocol, threads=threads, timeout=timeout, retries=retries, + cache=cache, cache_ttl=cache_ttl) + + +def score(data, questions, api_key=None, *, model=None, provider=None, protocol=None, + threads=None, timeout=None, retries=None, cache=None, cache_ttl=None): + """Answer named score questions using ordered, unique criterion descriptions. + + Scores retain the provider's native 0..N-1 scale. Confidence and complete + probabilities keyed by description accompany each score. Input/output batch + shape and configuration follow :func:`choose`. + """ + return _run(data, questions, "score", api_key=api_key, model=model, provider=provider, + protocol=protocol, threads=threads, timeout=timeout, retries=retries, + cache=cache, cache_ttl=cache_ttl) + + +def true_false(data, questions, api_key=None, *, model=None, provider=None, protocol=None, + threads=None, timeout=None, retries=None, cache=None, cache_ttl=None): + """Return probability_true and the supplied true_criteria for named questions. + + Optional criteria use string keys "true" and "false". No Boolean conversion + is performed. Input/output batch shape and configuration follow :func:`choose`. + """ + return _run(data, questions, "true_false", api_key=api_key, model=model, provider=provider, + protocol=protocol, threads=threads, timeout=timeout, retries=retries, + cache=cache, cache_ttl=cache_ttl) + + +def questions(data, questions, api_key=None, *, model=None, provider=None, protocol=None, + threads=None, timeout=None, retries=None, cache=None, cache_ttl=None): + """Evaluate mixed named questions in one request per input record. + + Each definition declares type choose, score, or true_false along with its + instructions and criteria. Other behavior follows the individual operations. + """ + return _run(data, questions, None, api_key=api_key, model=model, provider=provider, + protocol=protocol, threads=threads, timeout=timeout, retries=retries, + cache=cache, cache_ttl=cache_ttl) diff --git a/wrangles/ai_cache.py b/wrangles/ai_cache.py index 7ee787cc..a74e634a 100644 --- a/wrangles/ai_cache.py +++ b/wrangles/ai_cache.py @@ -41,6 +41,7 @@ class CachePolicy: max_value_bytes: int single_flight: bool log_every: int + event_name: str = "extract_ai_result_cache" class _Flight: @@ -68,8 +69,8 @@ def _env_number(name: str, default, converter): return default try: return converter(value) - except (TypeError, ValueError) as exc: - raise ValueError(f"{name} must be a valid {converter.__name__}.") from exc + except (TypeError, ValueError): + raise ValueError(f"{name} must be a valid {converter.__name__}.") from None def resolve_policy( @@ -77,8 +78,9 @@ def resolve_policy( *, enabled: bool = None, ttl_seconds: float = None, + env_prefix: str = "WRANGLES_EXTRACT_AI_CACHE", ) -> CachePolicy: - """Resolve config, per-call overrides, and operational environment switches.""" + """Resolve cache settings, retaining extraction's existing switches by default.""" config = config or {} configured_enabled = config.get("enabled", True) configured_single_flight = config.get("single_flight", True) @@ -92,7 +94,7 @@ def resolve_policy( else configured_enabled ) resolved_enabled = _env_bool( - "WRANGLES_EXTRACT_AI_CACHE_ENABLED", + f"{env_prefix}_ENABLED", resolved_enabled, ) resolved_ttl = ( @@ -101,26 +103,26 @@ def resolve_policy( else config.get("ttl_seconds", 3600) ) resolved_ttl = _env_number( - "WRANGLES_EXTRACT_AI_CACHE_TTL_SECONDS", + f"{env_prefix}_TTL_SECONDS", resolved_ttl, float, ) max_entries = _env_number( - "WRANGLES_EXTRACT_AI_CACHE_MAX_ENTRIES", + f"{env_prefix}_MAX_ENTRIES", config.get("max_entries", 512), int, ) max_value_bytes = _env_number( - "WRANGLES_EXTRACT_AI_CACHE_MAX_VALUE_BYTES", + f"{env_prefix}_MAX_VALUE_BYTES", config.get("max_value_bytes", 65536), int, ) single_flight = _env_bool( - "WRANGLES_EXTRACT_AI_CACHE_SINGLE_FLIGHT", + f"{env_prefix}_SINGLE_FLIGHT", configured_single_flight, ) log_every = _env_number( - "WRANGLES_EXTRACT_AI_CACHE_LOG_EVERY", + f"{env_prefix}_LOG_EVERY", config.get("log_every", 100), int, ) @@ -134,15 +136,15 @@ def resolve_policy( ): raise ValueError("cache_ttl must be a positive number of seconds.") if not isinstance(max_entries, int) or isinstance(max_entries, bool) or max_entries < 0: - raise ValueError("WRANGLES_EXTRACT_AI_CACHE_MAX_ENTRIES must be non-negative.") + raise ValueError(f"{env_prefix}_MAX_ENTRIES must be non-negative.") if ( not isinstance(max_value_bytes, int) or isinstance(max_value_bytes, bool) or max_value_bytes < 0 ): - raise ValueError("WRANGLES_EXTRACT_AI_CACHE_MAX_VALUE_BYTES must be non-negative.") + raise ValueError(f"{env_prefix}_MAX_VALUE_BYTES must be non-negative.") if not isinstance(log_every, int) or isinstance(log_every, bool) or log_every < 0: - raise ValueError("WRANGLES_EXTRACT_AI_CACHE_LOG_EVERY must be non-negative.") + raise ValueError(f"{env_prefix}_LOG_EVERY must be non-negative.") return CachePolicy( enabled=resolved_enabled and max_entries > 0 and max_value_bytes > 0, @@ -151,6 +153,8 @@ def resolve_policy( max_value_bytes=max_value_bytes, single_flight=single_flight, log_every=log_every, + event_name=("extract_ai_result_cache" if env_prefix == "WRANGLES_EXTRACT_AI_CACHE" + else "ai_result_cache"), ) @@ -232,7 +236,7 @@ def _maybe_log(policy: CachePolicy) -> None: if operations == 0 or operations % policy.log_every: return payload = { - "event": "extract_ai_result_cache", + "event": policy.event_name, **_STATS, "entries": len(_CACHE), "inflight": len(_INFLIGHT), diff --git a/wrangles/ai_defaults.yml b/wrangles/ai_defaults.yml index 18a32dab..f02cd397 100644 --- a/wrangles/ai_defaults.yml +++ b/wrangles/ai_defaults.yml @@ -105,6 +105,17 @@ providers: supported_values: task: [retrieval.query, retrieval.passage, text-matching, clustering, classification] + typesafe: + endpoints: + systemone: https://api.typesafe.ai/v1/systemone + documentation: + api: https://docs.typesafe.ai/introduction + models: + jev-1.13.0: + status: active + applications: [classification, scoring, true_false] + default_for: [ai.choose, ai.score, ai.true_false, ai.questions] + huggingface: endpoints: hf_inference: https://router.huggingface.co/hf-inference/models @@ -118,6 +129,62 @@ providers: models: {} operations: + ai.choose: + provider: typesafe + protocol: systemone + defaults: + default_concurrency: 10 + request_timeout_seconds: 30 + retries: 1 + cache: + enabled: true + ttl_seconds: 3600 + max_entries: 512 + max_value_bytes: 65536 + single_flight: true + log_every: 0 + ai.score: + provider: typesafe + protocol: systemone + defaults: + default_concurrency: 10 + request_timeout_seconds: 30 + retries: 1 + cache: + enabled: true + ttl_seconds: 3600 + max_entries: 512 + max_value_bytes: 65536 + single_flight: true + log_every: 0 + ai.true_false: + provider: typesafe + protocol: systemone + defaults: + default_concurrency: 10 + request_timeout_seconds: 30 + retries: 1 + cache: + enabled: true + ttl_seconds: 3600 + max_entries: 512 + max_value_bytes: 65536 + single_flight: true + log_every: 0 + ai.questions: + provider: typesafe + protocol: systemone + defaults: + default_concurrency: 10 + request_timeout_seconds: 30 + retries: 1 + cache: + enabled: true + ttl_seconds: 3600 + max_entries: 512 + max_value_bytes: 65536 + single_flight: true + log_every: 0 extract.ai: provider: openai protocol: responses diff --git a/wrangles/clients/typesafe.py b/wrangles/clients/typesafe.py new file mode 100644 index 00000000..25ffdfd7 --- /dev/null +++ b/wrangles/clients/typesafe.py @@ -0,0 +1,194 @@ +"""Small HTTP adapter for TypeSafe's typed System One responses. + +Public wrangles own input normalization, configuration, and row concurrency. +This adapter validates one response and never includes request or response +content in errors. See https://docs.typesafe.ai/api for the wire contract. +""" + +import math as _math +import time as _time +from datetime import datetime as _datetime, timezone as _timezone +from email.utils import parsedate_to_datetime as _parsedate_to_datetime + +import requests as _requests + + +class TypesafeError(RuntimeError): + """A sanitized provider failure, without request, response, or credentials.""" + + def __init__(self, message, *, status_code=None, attempts=None): + super().__init__(message) + self.status_code = status_code + self.attempts = attempts + + +class TypesafeResponseError(TypesafeError): + """The provider returned an invalid typed response; retrying is unsafe.""" + + +def _invalid(reason): + raise TypesafeResponseError(f"Invalid Typesafe response: {reason}.") from None + + +def _number(value, low, high): + # Compare integers directly: coercing an unusually large JSON integer to + # float inside isfinite() can overflow instead of rejecting the response. + return ( + type(value) is int and low <= value <= high + or type(value) is float and _math.isfinite(value) and low <= value <= high + ) + + +def _probabilities(answer, expected): + probabilities = answer.get("probabilities") + if not isinstance(probabilities, dict) or set(probabilities) != expected: + _invalid("probability keys do not match the requested criteria") + if any(not _number(value, 0, 1) for value in probabilities.values()): + _invalid("probabilities must be finite numbers between zero and one") + # The documented distribution sums approximately to one. Allow rounding, + # but neither renormalize provider values nor accept incomplete mass. + if not _math.isclose(_math.fsum(probabilities.values()), 1, rel_tol=0, abs_tol=0.01): + _invalid("probabilities do not sum approximately to one") + if not _number(answer.get("confidence"), 0, 1): + _invalid("confidence must be a finite number between zero and one") + + +def _validate_response(payload, questions): + if not isinstance(payload, dict): + _invalid("expected an object") + if not isinstance(payload.get("model"), str) or not payload["model"].strip(): + _invalid("model must be a non-empty string") + usage = payload.get("usage") + if not isinstance(usage, dict): + _invalid("usage must be an object") + # The current official SDK permits missing/null usage counts. Preserve + # those values instead of inventing zero counts for unreported usage. + for key in ("input_tokens", "output_tokens"): + value = usage.get(key) + if value is not None and (type(value) is not int or value < 0): + _invalid("reported token counts must be non-negative integers") + answers = payload.get("answers") + if not isinstance(answers, dict) or set(answers) != set(questions): + _invalid("answer IDs do not match the requested questions") + for question_id, question in questions.items(): + answer = answers[question_id] + kind = question["type"] + if not isinstance(answer, dict) or answer.get("type") != kind: + _invalid("answer type does not match its question") + if kind == "noul": + if not _number(answer.get("noul"), 0, 1): + _invalid("noul must be a finite number between zero and one") + elif kind == "choice": + expected = set(question["criteria"]) + if not isinstance(answer.get("choice"), str) or answer["choice"] not in expected: + _invalid("choice is not one of the requested options") + _probabilities(answer, expected) + elif kind == "score": + criteria = question["criteria"] + expected = {str(index) for index in range(len(criteria))} + if not _number(answer.get("score"), 0, len(criteria) - 1): + _invalid("score is outside the requested rubric") + legend = answer.get("legend") + if not isinstance(legend, dict) or set(legend) != expected: + _invalid("score legend does not match the requested levels") + if any(legend[str(index)] != description for index, description in enumerate(criteria)): + _invalid("score legend descriptions do not match the requested rubric") + _probabilities(answer, expected) + else: + _invalid("unsupported question type") + return payload + + +def _retry_delay(header, attempt): + """Honor numeric/HTTP-date Retry-After, bounded to sixty seconds.""" + fallback = min(2 ** min(attempt, 6), 60) + if not isinstance(header, str) or not header.strip(): + return fallback + try: + delay = float(header) + except ValueError: + try: + retry_at = _parsedate_to_datetime(header) + if retry_at.tzinfo is None: + retry_at = retry_at.replace(tzinfo=_timezone.utc) + delay = max(0, (retry_at - _datetime.now(_timezone.utc)).total_seconds()) + except (ValueError, TypeError, OverflowError): + return fallback + if not _math.isfinite(delay) or delay < 0: + return fallback + return min(delay, 60) + + +def _close_response(response): + try: + response.close() + except (OSError, _requests.exceptions.RequestException): + # A failed cleanup must not replace a sanitized provider error with + # an exception that may contain connection details or credentials. + pass + + +def call_systemone(*, state, questions, model, api_key, url, timeout, retries): + """Evaluate one state and return the validated, unmodified response. + + ``retries`` counts additional attempts. ``timeout`` applies to each HTTP + attempt; it is not a deadline for the complete operation. HTTP 429 and 5xx, + timeouts, and connection failures retry; invalid answers and other statuses + fail immediately. Redirects are disabled so credentials stay at the chosen + endpoint. All raised provider errors omit bodies and low-level causes. + """ + if type(retries) is not int or retries < 0: + raise ValueError("retries must be a non-negative integer.") + if type(timeout) not in (int, float) or not _math.isfinite(timeout) or timeout <= 0: + raise ValueError("timeout must be a positive finite number.") + payload = {"state": state, "model": model, "questions": questions} + headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} + transient_errors = ( + _requests.exceptions.Timeout, + _requests.exceptions.ConnectionError, + _requests.exceptions.ChunkedEncodingError, + _requests.exceptions.ContentDecodingError, + ) + for attempt in range(retries + 1): + response = None + delay = _retry_delay(None, attempt) + try: + try: + response = _requests.post( + url, headers=headers, json=payload, timeout=timeout, + allow_redirects=False, + ) + except (_requests.exceptions.RequestException, ValueError) as error: + # RequestException may contain a URL, headers, or body. Never + # retain it in the surfaced error's message or displayed chain. + error_response = getattr(error, "response", None) + if error_response is not None: + _close_response(error_response) + if isinstance(error, transient_errors) and attempt < retries: + pass + else: + raise TypesafeError( + "Typesafe request failed during transport.", attempts=attempt + 1, + ) from None + if response is not None: + status = response.status_code + if 200 <= status < 300: + try: + result = response.json() + except (ValueError, _requests.exceptions.RequestException): + raise TypesafeResponseError( + "Invalid Typesafe response: expected valid JSON.", + status_code=status, attempts=attempt + 1, + ) from None + return _validate_response(result, questions) + if (status == 429 or 500 <= status < 600) and attempt < retries: + delay = _retry_delay(response.headers.get("Retry-After"), attempt) + else: + raise TypesafeError( + f"Typesafe request failed with HTTP {status}.", + status_code=status, attempts=attempt + 1, + ) from None + finally: + if response is not None: + _close_response(response) + _time.sleep(delay) diff --git a/wrangles/dataframe.py b/wrangles/dataframe.py index 0e181a3d..0acacefc 100644 --- a/wrangles/dataframe.py +++ b/wrangles/dataframe.py @@ -89,6 +89,11 @@ def method(self, *args, **kwargs): else: setattr(self, name, make_method(name).__get__(self)) + @property + def ai(self): + """Ask named AI questions about each row.""" + return _wrangles_accessor(self._df, _recipe_wrangles.ai) + @property def compare(self): """ diff --git a/wrangles/recipe.py b/wrangles/recipe.py index f67a55d8..2a4d9503 100644 --- a/wrangles/recipe.py +++ b/wrangles/recipe.py @@ -559,6 +559,31 @@ def _read_data( else: return results +_NAMED_AI_WRANGLES = {"ai.choose", "ai.score", "ai.true_false", "ai.questions"} + + +def _declared_output_columns(wrangle: str, params: dict): + """Return declared destinations, including outputs nested in AI questions. + + None means the wrangle does not declare outputs. Keep ordinary wrangles' + existing output interpretation while sharing AI discovery with concurrent. + """ + if wrangle in _NAMED_AI_WRANGLES: + kind = wrangle.split(".", 1)[1] + return _recipe_wrangles.ai._output_columns( + params.get("questions"), kind=None if kind == "questions" else kind, + ) + if "output" not in params: + return None + output = params["output"] + if isinstance(output, list): + return [name for column in output + for name in (column.values() if isinstance(column, dict) else [column])] + if isinstance(output, dict): + return list(output.keys()) + return [output] + + def _execute_wrangles( df: _pandas.DataFrame, wrangles_list: list, @@ -672,28 +697,14 @@ def _execute_wrangles( # at least one row matches - otherwise batching with where can # produce inconsistent columns between batches if len(df) == 0: - if 'output' in params: - if isinstance(params['output'], list): - output_columns = [ - list(col.values()) if isinstance(col, dict) else [col] - for col in params['output'] - ] - output_columns = [ - item - for sublist in output_columns - for item in sublist - ] - elif isinstance(params['output'], dict): - output_columns = list(params['output'].keys()) - else: - output_columns = [params['output']] - + output_columns = _declared_output_columns(wrangle, params) + if output_columns is not None: for col in output_columns: # Wildcard outputs (e.g. 'Col*') are expanded into # concrete column names based on the actual data, # which can't be determined with no rows to work # with - skip adding those, only add named columns - if '*' in str(col): + if wrangle not in _NAMED_AI_WRANGLES and '*' in str(col): continue if col not in df_original.columns: df_original[col] = '' @@ -895,32 +906,11 @@ def _execute_wrangles( if 'where' in original_params and wrangle not in _where_overwrite_output: # Wrangle explictly defined the output - if 'output' in params.keys(): - # Get the columns that should have been added - if isinstance(params['output'], list): - # Wrangle output was a list - # this may be a list of columns or - # a list of dictionaries with renamed outputs - output_columns = [ - list(col.values()) if isinstance(col, dict) else [col] - for col in params['output'] - ] - # Spread to a 1D list - output_columns = [ - item - for sublist in output_columns - for item in sublist - ] - elif isinstance(params['output'], dict): - # Wrangle output was a dictionary, - # the keys should be the columns that were added - output_columns = list(params['output'].keys()) - else: - # Scalar value - output_columns = [params['output']] - + output_columns = _declared_output_columns(wrangle, params) + if output_columns is not None: # Expand the columns if using any wildcards - output_columns = _wildcard_expansion(df.columns, output_columns) + if wrangle not in _NAMED_AI_WRANGLES: + output_columns = _wildcard_expansion(df.columns, output_columns) df = df[output_columns] @@ -976,7 +966,9 @@ def _execute_wrangles( df = df.fillna('0') if wrangle != 'log': # Determine what columns were actually produced for logging - if 'output' in params: + if wrangle in _NAMED_AI_WRANGLES: + output_columns = _declared_output_columns(wrangle, params) + elif 'output' in params: if isinstance(params['output'], list): # Handle mixed list types (strings and dicts) output_columns = [] diff --git a/wrangles/recipe_wrangles/__init__.py b/wrangles/recipe_wrangles/__init__.py index 46f1e5e2..14632e8f 100644 --- a/wrangles/recipe_wrangles/__init__.py +++ b/wrangles/recipe_wrangles/__init__.py @@ -22,6 +22,7 @@ from . import convert from . import create from . import extract +from . import ai from . import format from . import merge from . import select diff --git a/wrangles/recipe_wrangles/ai.py b/wrangles/recipe_wrangles/ai.py new file mode 100644 index 00000000..81ac5188 --- /dev/null +++ b/wrangles/recipe_wrangles/ai.py @@ -0,0 +1,182 @@ +"""Dataframe adapters for named AI questions.""" + +import pandas as _pd +import yaml as _yaml + +from .. import ai as _ai + + +_FIELDS = _ai._FIELDS + + +def _output_columns(questions, kind=None): + """Resolve nested destinations before filtering or concurrent execution.""" + return _ai._output_columns(questions, kind=kind) + + +def _run(df, questions, kind, input=None, **settings): + prepared = _ai._prepare_questions(questions, kind=kind) + if df.empty: + for question in prepared.values(): + for column in question["output"]: + if column not in df.columns: + df[column] = _pd.Series(index=df.index, dtype=object) + return df + + source = df if input is None else df[input if isinstance(input, list) else [input]] + operation = getattr(_ai, kind or "questions") + results = operation(source.to_dict(orient="records"), questions=questions, **settings) + if not isinstance(results, list) or len(results) != len(df): + raise RuntimeError("AI response count does not match the input row count.") + for label, question in prepared.items(): + for column, field in zip(question["output"], _FIELDS[question["type"]]): + df[column] = [result[label][field] for result in results] + return df + + +def choose(df: _pd.DataFrame, questions: dict, input=None, api_key=None, + model=None, provider=None, protocol=None, threads=None, timeout=None, + retries=None, cache=None, cache_ttl=None) -> _pd.DataFrame: + return _run(df, questions, "choose", input=input, api_key=api_key, + model=model, provider=provider, protocol=protocol, threads=threads, + timeout=timeout, retries=retries, cache=cache, cache_ttl=cache_ttl) + + +def score(df: _pd.DataFrame, questions: dict, input=None, api_key=None, + model=None, provider=None, protocol=None, threads=None, timeout=None, + retries=None, cache=None, cache_ttl=None) -> _pd.DataFrame: + return _run(df, questions, "score", input=input, api_key=api_key, + model=model, provider=provider, protocol=protocol, threads=threads, + timeout=timeout, retries=retries, cache=cache, cache_ttl=cache_ttl) + + +def true_false(df: _pd.DataFrame, questions: dict, input=None, api_key=None, + model=None, provider=None, protocol=None, threads=None, timeout=None, + retries=None, cache=None, cache_ttl=None) -> _pd.DataFrame: + return _run(df, questions, "true_false", input=input, api_key=api_key, + model=model, provider=provider, protocol=protocol, threads=threads, + timeout=timeout, retries=retries, cache=cache, cache_ttl=cache_ttl) + + +def questions(df: _pd.DataFrame, questions: dict, input=None, api_key=None, + model=None, provider=None, protocol=None, threads=None, timeout=None, + retries=None, cache=None, cache_ttl=None) -> _pd.DataFrame: + return _run(df, questions, None, input=input, api_key=api_key, + model=model, provider=provider, protocol=protocol, threads=threads, + timeout=timeout, retries=retries, cache=cache, cache_ttl=cache_ttl) + + +def _description_schema(nullable=False): + variants = [ + {"type": "string", "pattern": r"\S"}, + {"type": "object", "minProperties": 1}, + {"type": "array", "minItems": 1}, + ] + if nullable: + variants.append({"type": "null"}) + return {"anyOf": variants} + + +def _question_schema(kind): + fields = _FIELDS[kind] + properties = { + "type": {"type": "string", "enum": [kind], + "description": "Question type; optional for a wrangle with one question type."}, + "instructions": {**_description_schema(), + "description": "Question or structured instructions applied to this row's input."}, + "output": { + "description": ( + "Destination columns in order: " + ", ".join(fields) + ". " + "Omit, use null, or use a blank string for defaults. " + "Defaults are the question label followed by _confidence and _probabilities " + "for choose/score, or the question label and _true_criteria for true_false. " + "An explicit list must name every output column." + ), + "oneOf": [ + {"type": "null"}, + {"type": "string", "pattern": r"^\s*$"}, + {"type": "array", "minItems": len(fields), "maxItems": len(fields), + "uniqueItems": True, "items": {"type": "string", "pattern": r"\S"}}, + ], + }, + } + required = ["instructions"] + if kind == "score": + properties["criteria"] = { + "type": "array", "minItems": 2, "maxItems": 10, + "uniqueItems": True, "items": {"type": "string", "pattern": r"\S"}, + "description": ( + "Distinct criterion descriptions used as probability keys. " + "For score, order defines the native zero-based score positions; no custom scale." + ), + } + required.append("criteria") + elif kind == "choose": + properties["criteria"] = { + "type": "object", "minProperties": 1, "maxProperties": 255, + "propertyNames": {"pattern": r"\S"}, + "additionalProperties": _description_schema(nullable=True), + "description": "Choice labels mapped to their descriptions or structured criteria.", + } + required.append("criteria") + else: + properties["criteria"] = { + "type": "object", "propertyNames": {"enum": ["true", "false"]}, + "additionalProperties": _description_schema(), + "description": ( + "Optional criteria keyed by the strings true and false; quote these YAML keys. " + "The true criterion is retained in the true_criteria output." + ), + } + return {"type": "object", "additionalProperties": False, + "required": required, "properties": properties} + + +def _schema(kind): + if kind is None: + question_schemas = [] + for question_type in _FIELDS: + question_schema = _question_schema(question_type) + question_schema["required"].append("type") + question_schemas.append(question_schema) + question_definition = {"oneOf": question_schemas} + else: + question_definition = _question_schema(kind) + return { + "type": "object", "additionalProperties": False, "required": ["questions"], + "description": ( + "Ask named " + (kind.replace("_", "/") if kind else "mixed") + + " questions in one provider request per row. " + "Question labels determine default output names. Scores are native numeric scores; " + "true_false returns the probability of true, without thresholding or a Boolean." + ), + "properties": { + "input": {"type": ["string", "integer", "array"], + "items": {"type": ["string", "integer"]}, + "description": "Input column(s); omit to supply all columns as a row record."}, + "questions": {"type": "object", "minProperties": 1, + "propertyNames": {"pattern": r"\S"}, + "additionalProperties": question_definition, + "description": "Questions keyed by their distinct labels."}, + "api_key": {"type": "string", "description": "Typesafe API key; defaults to TYPESAFE_API_KEY."}, + "model": {"type": "string", "description": "Model ID; defaults to this operation's AI catalog selection."}, + "provider": {"type": "string", "enum": ["typesafe"], + "description": "Provider resolved through the AI catalog; currently typesafe."}, + "protocol": {"type": "string", "enum": ["systemone"], + "description": "Provider protocol selected through the AI catalog."}, + "threads": {"type": "integer", "minimum": 1, + "description": "Maximum row requests in parallel; defaults to the AI catalog."}, + "timeout": {"type": "number", "exclusiveMinimum": 0, + "description": "Per-attempt request timeout in seconds; defaults to the AI catalog."}, + "retries": {"type": "integer", "minimum": 0, + "description": "Additional attempts for retryable failures; defaults to the AI catalog."}, + "cache": {"type": "boolean", "description": "Reuse identical successful requests through the shared AI cache."}, + "cache_ttl": {"type": "number", "exclusiveMinimum": 0, + "description": "Override the cache result lifetime in seconds."}, + }, + } + + +for _name, _kind in (("choose", "choose"), ("score", "score"), + ("true_false", "true_false"), ("questions", None)): + globals()[_name].__doc__ = _yaml.safe_dump(_schema(_kind), sort_keys=False) diff --git a/wrangles/recipe_wrangles/main.py b/wrangles/recipe_wrangles/main.py index b01806e6..d468eb80 100644 --- a/wrangles/recipe_wrangles/main.py +++ b/wrangles/recipe_wrangles/main.py @@ -488,10 +488,13 @@ def concurrent( futures = [] futures_output_map = {} for wrangle_definition in wrangles: - if ( - not isinstance(wrangle_definition, dict) or - "output" not in list(wrangle_definition.values())[0] - ): + if not isinstance(wrangle_definition, dict) or len(wrangle_definition) != 1: + raise ValueError('Using concurrent requires that each wrangle specify output column(s).') + wrangle_name, wrangle_params = next(iter(wrangle_definition.items())) + output_columns = _wrangles.recipe._declared_output_columns( + wrangle_name, wrangle_params or {}, + ) + if output_columns is None: raise ValueError('Using concurrent requires that each wrangle specify output column(s).') if use_multiprocessing: @@ -514,7 +517,7 @@ def concurrent( futures.append(future) # Add output columns to reference on completion - futures_output_map[future] = list(wrangle_definition.values())[0]["output"] + futures_output_map[future] = output_columns # Wait for all futures to complete for future in _futures.as_completed(futures): From d5fbaa849664c6f6a63693e740ecb57b1c5187e1 Mon Sep 17 00:00:00 2001 From: Eric Hills <53243273+ebhills@users.noreply.github.com> Date: Sat, 26 Sep 2026 15:48:44 -0500 Subject: [PATCH 2/3] Expand AI wrangle guide with complete input and output examples --- docs/ai_questions.md | 384 +++++++++++++++++++++++++++---------------- 1 file changed, 241 insertions(+), 143 deletions(-) diff --git a/docs/ai_questions.md b/docs/ai_questions.md index 4f46ef77..58c6f4d5 100644 --- a/docs/ai_questions.md +++ b/docs/ai_questions.md @@ -1,213 +1,316 @@ # AI questions `ai.choose`, `ai.score`, `ai.true_false`, and `ai.questions` answer named questions -about each input record through Typesafe. They use the existing -[AI model catalog](ai_configuration.md), with `provider: typesafe`, -`protocol: systemone`, and the pinned model `jev-1.13.0` as defaults. +about each input record through [Typesafe](https://docs.typesafe.ai/introduction). +They use the existing [AI model catalog](ai_configuration.md), with +`provider: typesafe`, `protocol: systemone`, and the pinned model `jev-1.13.0` +as defaults. `extract.ai` keeps its current API and behavior. -`extract.ai` keeps its current API and behavior. The new `ai` namespace leaves -room for a later compatible `ai.extract` migration; that alias is not introduced -here. +## Input and output schema -## Named questions and answers +Every wrangle takes a nonempty `questions` mapping. Each key is a nonblank +question label, used to construct the default output column names. Each +question requires `instructions`: nonempty text, a JSON object, or a JSON array. +Objects must have string keys, and all supplied values must be JSON-compatible. -Every operation takes a nonempty `questions` mapping. Its keys are the question -names and form the default output column names. Every question has -`instructions`, which can be a string, a JSON object, or a JSON array. -Instructions and supplied descriptions must be nonempty; choose descriptions -may also use `null` when the option label is sufficient. - -| Operation | Criteria | Answer fields and default recipe columns | +| Wrangle | Criteria | Answer fields and default output columns | | --- | --- | --- | -| `ai.choose` | A mapping of 1–255 nonblank option labels to descriptions. Descriptions can be strings, JSON objects, JSON arrays, or `null`. | `choice` → ``; `confidence` → `_confidence`; `probabilities` → `_probabilities` | -| `ai.score` | An ordered list of 2–10 unique nonblank strings describing the scoring criteria. Structured criteria objects are not accepted. | `score` → ``; `confidence` → `_confidence`; `probabilities` → `_probabilities` | -| `ai.true_false` | Optional `"true"` and/or `"false"` descriptions, supplied as strings, JSON objects, or JSON arrays. Quote these keys in YAML. | `probability_true` → ``; `true_criteria` → `_true_criteria` | -| `ai.questions` | Each question declares `type: choose`, `type: score`, or `type: true_false` and uses the corresponding schema above. | The columns for that question's type. | +| `ai.choose` | A mapping of 1–255 nonblank option labels to descriptions. Descriptions can be nonempty strings, JSON objects, JSON arrays, or `null`. | `choice` → `