This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f8b1e8cdd75 [Python] Add typed decision questions and model adapters
(#40220)
f8b1e8cdd75 is described below
commit f8b1e8cdd75a02e19e1026f58428c59a91ba7c70
Author: Bruno Volpato <[email protected]>
AuthorDate: Mon Oct 19 14:54:48 2026 -0400
[Python] Add typed decision questions and model adapters (#40220)
* [Python] Add decision model routing examples
* [Python] Reuse decision model clients in a lightweight DoFn
* [Python] Add reusable decision model primitives and adapters
* [Python] Gate live Jev tests and isolate runtime value options
* [Python] Validate decision responses and add dead-letter handling
* [Python] Run decision adapter tests in ML CI
* [Python] Move decision SDK tests into the ML extra
---
CHANGES.md | 1 +
.../apache_beam/examples/inference/README.md | 8 +-
.../examples/inference/decision_models/README.md | 304 +++++++++++++++++++++
.../examples/inference/decision_models/__init__.py | 18 ++
.../inference/decision_models/fraud_review.py | 126 +++++++++
.../inference/decision_models/fraud_review_test.py | 86 ++++++
.../inference/decision_models/local_model.py | 78 ++++++
.../inference/decision_models/message_router.py | 199 ++++++++++++++
.../inference/decision_models/message_router.svg | 98 +++++++
.../decision_models/message_router_test.py | 75 +++++
.../inference/decision_models/requirements.txt | 19 ++
.../apache_beam/examples/snippets/snippets_test.py | 8 +-
sdks/python/apache_beam/ml/inference/decision.py | 246 +++++++++++++++++
.../apache_beam/ml/inference/decision_test.py | 254 +++++++++++++++++
.../apache_beam/ml/inference/typesafe_inference.py | 164 +++++++++++
.../ml/inference/typesafe_inference_it_test.py | 93 +++++++
.../ml/inference/typesafe_inference_test.py | 256 +++++++++++++++++
.../ml/inference/typesafe_tests_requirements.txt | 18 ++
sdks/python/setup.py | 1 +
19 files changed, 2049 insertions(+), 3 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index d50b14cadd1..00fb9aa825b 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -73,6 +73,7 @@
* (Java) Flink runner: sources now forward Beam metrics to Flink at most once
per second instead of after every record, which was a significant share of the
CPU of reading. Metrics in the pipeline result are not affected
([#40607](https://github.com/apache/beam/issues/40607)).
* (Go) Added `wait.On`, which delays each input window until the corresponding
windows in its signal PCollections have closed
([#39909](https://github.com/apache/beam/issues/39909)).
* (Python) Expanded the SDK worker heap dump
(`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC
stats, and glibc `mallinfo2` native-heap/fragmentation stats to help
distinguish native-heap from Python-object memory growth
([#39244](https://github.com/apache/beam/issues/39244)).
+* (Python) Added typed decision questions and the `EvaluateDecisions`
transform, with an optional Jev model adapter
([#40220](https://github.com/apache/beam/pull/40220)).
* The `disableCounterMetrics`, `disableStringSetMetrics` and
`disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as
they already were in Java (Python)
([#38746](https://github.com/apache/beam/issues/38746)).
* (Go) Pipeline-level `--resource_hints` are now forwarded to expansion
services, so they apply to cross-language transforms as they do in the Java and
Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted
by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)).
* `QdrantWriteConfig` can now optionally store chunk content in the point
payload via `content_payload_key`, so retrieved points include the chunk text
(Python) ([#40476](https://github.com/apache/beam/issues/40476)).
diff --git a/sdks/python/apache_beam/examples/inference/README.md
b/sdks/python/apache_beam/examples/inference/README.md
index c43ad458347..f14b531c21a 100644
--- a/sdks/python/apache_beam/examples/inference/README.md
+++ b/sdks/python/apache_beam/examples/inference/README.md
@@ -17,7 +17,13 @@
under the License.
-->
-# Example RunInference API pipelines
+# Example inference pipelines
+
+For typed classification, Boolean judgments, and scoring, see the
+[decision model examples](decision_models/README.md). They use
+`EvaluateDecisions` with a swappable model adapter, including Jev and local
rules.
+
+## RunInference
This module contains example pipelines that use the Beam RunInference
API. <!---TODO: Add link to full documentation on Beam website when it's
published.-->
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/README.md
b/sdks/python/apache_beam/examples/inference/decision_models/README.md
new file mode 100644
index 00000000000..637b8159c55
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md
@@ -0,0 +1,304 @@
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+-->
+
+# Beam decision models
+
+These examples use the reusable decision primitives in
+`apache_beam.ml.inference.decision`. A model evaluates named, typed questions
+for each Beam element. `EvaluateDecisions` returns the original state,
+normalized answers, provider metadata, and request latency. Policy and sinks
+remain ordinary Beam transforms after evaluation.
+
+## Public API
+
+The core module contains `ChoiceQuestion`, `BooleanQuestion`, and
+`ScoreQuestion`, their typed answers (`ChoiceAnswer`, `BooleanAnswer`, and
+`ScoreAnswer`), the `DecisionModel` ABC, and the `EvaluateDecisions`
+`PTransform`.
+
+| Question | Normalized answer |
+| --- | --- |
+| `ChoiceQuestion(instructions, criteria)` | `ChoiceAnswer.choice`, with
optional label probabilities and confidence |
+| `BooleanQuestion(instructions)` | `BooleanAnswer.probability` in `[0, 1]` |
+| `ScoreQuestion(instructions, criteria)` | `ScoreAnswer.score` and
`ScoreAnswer.probabilities`, keyed by zero-based integer level, with optional
confidence |
+
+`JevDecisionModel` maps `BooleanQuestion` to Jev's `Noul`, returning
+`BooleanAnswer.probability`; a missing confidence is `None`.
+
+`DecisionModel` is the adapter ABC. A subclass implements
+`evaluate(state, questions)` and returns one matching typed answer per named
+question in a `DecisionResponse`. Override `__enter__` and `__exit__` to reuse
+a client, connection pool, or model weights. The transform enters the model
+during DoFn setup and closes it during teardown for each DoFn instance.
+
+`EvaluateDecisions(model, questions)` is a `PTransform` from input states to
+`DecisionResult`. Evaluation is synchronous, timestamps and windows are
+preserved, exceptions propagate to the runner, and request, failure, and
+latency metrics are recorded.
+
+Before returning a result, `JevDecisionModel` checks that its response contains
+the requested questions with matching answer types. It rejects unknown choice
+labels and non-finite or out-of-range scores, probabilities, and confidence
+values with `ValueError`. Custom `DecisionModel` implementations must enforce
+the same output contract. These errors fail the pipeline by default.
+
+## Architecture
+
+```mermaid
+flowchart LR
+ input["PCollection[state]"] --> evaluate["EvaluateDecisions(model,
questions)<br/>PTransform"]
+ model["DecisionModel adapter"] --> evaluate
+ evaluate --> result["PCollection[DecisionResult]"]
+ result --> policy["route_message<br/>policy and sink selection"]
+ policy --> sink["JSONL files<br/>or BigQuery tables"]
+```
+
+
+
+The optional Jev adapter lives in
+`apache_beam.ml.inference.typesafe_inference`; the deterministic
+`LocalDecisionModel` stays in
+`apache_beam.examples.inference.decision_models.local_model`. Both implement
+the same core contract.
+
+## Integration
+
+Questions and events are regular Python values. The model is the only
+provider-specific argument to `EvaluateDecisions`:
+
+```python
+import apache_beam as beam
+
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.examples.inference.decision_models.local_model import (
+ LocalDecisionModel,
+)
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+
+questions = {
+ 'team': ChoiceQuestion('Which team should handle this message?', {
+ 'billing': 'Invoices', 'technical': 'Technical problems'}),
+}
+
+def select_team(result):
+ return result.response.answers['team'].choice
+
+model = LocalDecisionModel()
+# model = JevDecisionModel()
+
+with beam.Pipeline() as pipeline:
+ events = pipeline | 'Events' >> beam.Create(
+ [{'message': 'The API returns 503 after I rotated my key.'}])
+ decisions = events | 'Evaluate decisions' >> EvaluateDecisions(
+ model, questions)
+ rows = decisions | 'Apply downstream policy' >> beam.Map(select_team)
+```
+
+The downstream policy stays unchanged when swapping adapters. The Jev client
+is optional; install the shared SDK requirements listed below and set
+`TYPESAFE_API_KEY` when running it.
+
+## Install
+
+From a Beam checkout, activate a virtual environment and install the core SDK:
+
+```sh
+cd sdks/python
+python3 -m pip install -e .
+```
+
+Install the optional Jev SDK used by the adapter and examples:
+
+```sh
+cd sdks/python
+python3 -m pip install -r \
+ apache_beam/ml/inference/typesafe_tests_requirements.txt
+```
+
+Install Beam's GCP extra for Pub/Sub or BigQuery sinks:
+
+```sh
+cd sdks/python
+python3 -m pip install -e '.[gcp]'
+```
+
+The example requirements file delegates to that same bounded dependency file.
+
+## Run examples
+
+The main demo uses a three-element `TestStream` and writes one JSONL file per
+destination:
+
+```sh
+cd sdks/python
+output_dir="$(mktemp -d /tmp/beam-decision-model-output.XXXXXX)"
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=local \
+ --output-dir="$output_dir" \
+ --graph="$output_dir/message_router.svg"
+find "$output_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \;
+```
+
+Render the graph without running the job. The sink stays in the graph:
+
+```sh
+cd sdks/python
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=jev \
+ --graph-only \
+ --bq-dataset=PROJECT:DATASET \
+ --graph=/tmp/beam-decision-model.svg
+```
+
+The `.dot` export works without Graphviz; SVG and PNG require `dot` on `PATH`.
+
+Run the Jev path against the same stream with `TYPESAFE_API_KEY` set:
+
+```sh
+cd sdks/python
+export TYPESAFE_API_KEY='replace-with-your-key'
+jev_output_dir="$(mktemp -d /tmp/beam-decision-model-jev-output.XXXXXX)"
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=jev \
+ --min-confidence=0.65 \
+ --output-dir="$jev_output_dir" \
+ --graph="$jev_output_dir/message_router.svg"
+```
+
+For a streaming source, provide a Pub/Sub subscription containing JSON objects
+with `event_id` and `message` fields:
+
+```sh
+cd sdks/python
+python3 -m apache_beam.examples.inference.decision_models.message_router \
+ --model=jev \
+ --pubsub-subscription=projects/PROJECT/subscriptions/SUBSCRIPTION \
+ --bq-dataset=PROJECT:DATASET
+```
+
+With `--bq-dataset=PROJECT:DATASET`, rows go to `messages_billing`,
+`messages_technical`, `messages_sales`, or `messages_review`; `--output-dir`
+writes local JSONL files. A normal run accepts one sink option.
+
+The fraud example prints one JSON row per sample message. Noul supplies the
+boolean probability; values at or above `0.7` are flagged:
+
+```sh
+cd sdks/python
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=local --primitive=noul
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=local --primitive=both
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=jev --primitive=both
+```
+
+Pass `--dead-letter-dir` to write per-element `ValueError` failures from
+decision evaluation as JSONL records containing the original input and the
+exception representation. This opt-in catches malformed-response validation
+errors; without the option, they fail the pipeline. The Jev SDK retries
+transient transport and service failures with exponential backoff and jitter,
+and honors `Retry-After` when supplied. Those retries run independently on each
+Beam worker, which does not impose a global request-rate limit. Failures that
+remain after retries still fail the pipeline instead of going to this sink.
+Errors while the model context is created during worker setup also fail the
+pipeline and are not routed to this sink.
+
+For example, enable the local dead-letter sink with the Jev adapter:
+
+```sh
+cd sdks/python
+export TYPESAFE_API_KEY='replace-with-your-key'
+dead_letter_dir="$(mktemp -d /tmp/beam-decision-model-dead-letter.XXXXXX)"
+python3 -m apache_beam.examples.inference.decision_models.fraud_review \
+ --model=jev --primitive=both --dead-letter-dir="$dead_letter_dir"
+find "$dead_letter_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \;
+```
+
+A final three-element DirectRunner run returned these values. Beam may print
+the rows in a different order:
+
+| message | Noul probability | Score |
+| --- | ---: | ---: |
+| billing address | 0.42 | 0.80 |
+| duplicate charge | 0.32 | 0.85 |
+| bypass verification | 0.97 | 3.00 |
+
+Only the bypass verification message crossed the `0.7` review threshold.
+
+## Tests
+
+Install Beam's test dependencies in the same virtual environment:
+
+```sh
+cd sdks/python
+python3 -m pip install -e '.[test]'
+```
+
+The unit tests cover the shared API, routing policy, and mocked Jev responses:
+
+```sh
+cd sdks/python
+python3 -m pytest -q \
+ apache_beam/ml/inference/decision_test.py \
+ apache_beam/ml/inference/typesafe_inference_test.py \
+ apache_beam/examples/inference/decision_models/message_router_test.py \
+ apache_beam/examples/inference/decision_models/fraud_review_test.py
+```
+
+With the optional Jev SDK installed and `TYPESAFE_API_KEY` set, run the live
+integration test:
+
+```sh
+cd sdks/python
+python3 -m pytest -q -rs -m it_postcommit \
+ apache_beam/ml/inference/typesafe_inference_it_test.py
+```
+
+It sends one streaming event through `EvaluateDecisions`, asking Choice,
+Boolean/Noul, and Score questions together. It checks the normalized answer
+types and ranges, and skips when the SDK or API key is missing. The mocked
+adapter tests run with the SDK installed and require no API key.
+
+## Jev benchmark
+
+These measurements were captured on September 22, 2026 with Jev 1.13.0,
+TypeSafe SDK 0.7.1, and DirectRunner. The final trace used exactly one SDK
+client, one TCP connection, and one TLS handshake for three events:
+
+| event | request latency |
+| --- | ---: |
+| billing | 408.27 ms |
+| technical | 217.50 ms |
+| sales | 190.43 ms |
+
+Total pipeline time was `986.68 ms`, including pipeline construction.
+
+### Standalone transport experiment, September 22, 2026
+
+An earlier standalone experiment compared fresh clients with one pooled client:
+
+| client setup | observations |
+| --- | --- |
+| fresh client, `n=3` | 602.44 ms, 433.56 ms, 537.60 ms |
+| one pooled client | first 416.77 ms; next five 170.98 ms, 211.24 ms, 222.02
ms, 240.40 ms, 226.82 ms |
+
+The trace showed no TCP/TLS spans on reused calls. Reuse pays connection setup
+once. Warm-call delay is response wait plus network and service time.
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/__init__.py
b/sdks/python/apache_beam/examples/inference/decision_models/__init__.py
new file mode 100644
index 00000000000..fc00cf8225f
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/__init__.py
@@ -0,0 +1,18 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Decision model routing examples."""
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py
b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py
new file mode 100644
index 00000000000..5afcb6f34f2
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py
@@ -0,0 +1,126 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Small Beam example: flag messages for fraud review with Noul or Score."""
+
+import argparse
+import json
+import os
+
+import apache_beam as beam
+from apache_beam.examples.inference.decision_models.local_model import (
+ LocalDecisionModel)
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+from apache_beam.options.pipeline_options import PipelineOptions
+
+SAMPLE_MESSAGES = [
+ 'Please update my billing address before the next renewal.',
+ 'I was charged twice; please look into the duplicate payment.',
+ 'Transfer the refund to this new account immediately and '
+ 'skip verification.',
+]
+
+FRAUD_NOUL = BooleanQuestion(
+ instructions='Does this message contain signals that warrant fraud
review?')
+FRAUD_SCORE = ScoreQuestion(
+ instructions='How strongly does this message warrant fraud review?',
+ criteria=(
+ 'No visible fraud cue',
+ 'A weak cue that merits routine checking',
+ 'A concrete suspicious instruction',
+ 'An explicit attempt to bypass verification'))
+
+
+def review_fraud_cues(result):
+ response = result.response
+ row = {
+ 'message': result.state['message'],
+ 'model': response.model,
+ 'provider': response.provider,
+ 'latency_ms': round(result.latency_ms, 2),
+ }
+ if 'fraud_cue' in response.answers:
+ probability = response.answers['fraud_cue'].probability
+ row['fraud_review_probability'] = probability
+ row['needs_review'] = probability >= 0.7
+ if 'risk_level' in response.answers:
+ score = response.answers['risk_level']
+ row['risk_score'] = score.score
+ row['risk_confidence'] = score.confidence
+ row['risk_probabilities'] = score.probabilities
+ return row
+
+
+def dead_letter_record(dead_letter):
+ """Keep the input and Beam's exception representation for a failed row."""
+ original_input, exception_info = dead_letter
+ _, error, _ = exception_info
+ return {'input': original_input, 'error': error}
+
+
+def run(argv=None):
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument('--model', choices=('jev', 'local'), default='jev')
+ parser.add_argument(
+ '--primitive', choices=('noul', 'score', 'both'), default='noul')
+ parser.add_argument(
+ '--dead-letter-dir',
+ help=(
+ 'Optional local JSONL output directory for per-element '
+ 'ValueError failures during decision evaluation.'))
+ args, pipeline_args = parser.parse_known_args(argv)
+ if args.model == 'jev' and not os.environ.get('TYPESAFE_API_KEY'):
+ parser.error('TYPESAFE_API_KEY is required for --model=jev')
+ model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel()
+ questions = {}
+ if args.primitive in ('noul', 'both'):
+ questions['fraud_cue'] = FRAUD_NOUL
+ if args.primitive in ('score', 'both'):
+ questions['risk_level'] = FRAUD_SCORE
+ with beam.Pipeline(options=PipelineOptions(pipeline_args)) as pipeline:
+ states = pipeline | 'Sample payment messages' >> beam.Create([{
+ 'message': message
+ } for message in SAMPLE_MESSAGES])
+ evaluation = EvaluateDecisions(model, questions)
+ if args.dead_letter_dir:
+ good, bad = (
+ states
+ | 'Evaluate fraud cues' >> evaluation.with_exception_handling(
+ exc_class=ValueError))
+ (
+ bad
+ | 'Format malformed responses' >> beam.Map(dead_letter_record)
+ | 'Encode malformed responses' >>
+ beam.Map(lambda row: json.dumps(row, sort_keys=True))
+ | 'Write malformed responses' >> beam.io.WriteToText(
+ os.path.join(args.dead_letter_dir, 'malformed'),
+ file_name_suffix='.jsonl'))
+ decisions = good
+ else:
+ decisions = states | 'Evaluate fraud cues' >> evaluation
+ (
+ decisions
+ | 'Apply review threshold' >> beam.Map(review_fraud_cues)
+ | 'Show review decisions' >>
+ beam.Map(lambda row: print(json.dumps(row, sort_keys=True))))
+
+
+if __name__ == '__main__':
+ run()
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py
b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py
new file mode 100644
index 00000000000..ec6f8160ea6
--- /dev/null
+++
b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py
@@ -0,0 +1,86 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for the fraud example's opt-in malformed-response dead letter."""
+
+import glob
+import json
+import os
+import tempfile
+import unittest
+
+import apache_beam as beam
+from apache_beam.examples.inference.decision_models.fraud_review import (
+ dead_letter_record)
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+
+_QUESTION = {'fraud_cue': BooleanQuestion('Does this need fraud review?')}
+_VALID_INPUT = {'message': 'Looks like an ordinary payment question.'}
+_MALFORMED_INPUT = {'message': 'Provider returned an invalid probability.'}
+
+
+class _MalformedResponseModel(DecisionModel):
+ def evaluate(self, state, questions):
+ if state == _MALFORMED_INPUT:
+ raise ValueError('Invalid response probability: 1.2')
+ return DecisionResponse(
+ answers={'fraud_cue': BooleanAnswer(0.2)},
+ model='fixed',
+ provider='test')
+
+
+class FraudReviewDeadLetterTest(unittest.TestCase):
+ def test_malformed_response_is_routed_with_original_input(self):
+ with tempfile.TemporaryDirectory() as output_dir:
+ with TestPipeline() as pipeline:
+ good, bad = (
+ pipeline
+ | beam.Create([_VALID_INPUT, _MALFORMED_INPUT])
+ | EvaluateDecisions(_MalformedResponseModel(), _QUESTION).
+ with_exception_handling(exc_class=ValueError))
+
+ assert_that(
+ good | beam.Map(lambda result: result.state),
+ equal_to([_VALID_INPUT]))
+ (
+ bad
+ | beam.Map(dead_letter_record)
+ | beam.Map(lambda row: json.dumps(row, sort_keys=True))
+ | beam.io.WriteToText(
+ os.path.join(output_dir, 'malformed'),
+ file_name_suffix='.jsonl',
+ num_shards=1))
+
+ output_files = glob.glob(os.path.join(output_dir, 'malformed-*.jsonl'))
+ self.assertEqual(1, len(output_files))
+ with open(output_files[0], encoding='utf-8') as output_file:
+ records = [json.loads(line) for line in output_file]
+ self.assertEqual(1, len(records))
+ self.assertEqual(_MALFORMED_INPUT, records[0]['input'])
+ self.assertTrue(records[0]['error'].startswith('ValueError('))
+ self.assertIn('1.2', records[0]['error'])
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/local_model.py
b/sdks/python/apache_beam/examples/inference/decision_models/local_model.py
new file mode 100644
index 00000000000..55c30216f2c
--- /dev/null
+++ b/sdks/python/apache_beam/examples/inference/decision_models/local_model.py
@@ -0,0 +1,78 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Deterministic rules for the decision model examples."""
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+
+_ROUTE_TERMS = {
+ 'billing': ('invoice', 'charge', 'refund', 'payment', 'subscription'),
+ 'technical': ('api', '503', 'error', 'integration', 'outage'),
+ 'sales': ('quote', 'seat', 'pricing', 'price', 'plan'),
+}
+
+
+class LocalDecisionModel(DecisionModel):
+ """Transparent rules for running these examples without network access."""
+ def evaluate(self, state, questions):
+ message = state['message'].lower()
+ fraud_cue = (
+ 'skip verification' in message or 'bypass verification' in message or
+ 'unauthorized' in message or 'phishing' in message)
+ answers = {}
+ for name, question in questions.items():
+ if isinstance(question, ChoiceQuestion):
+ matches = {
+ label: sum(term in message for term in _ROUTE_TERMS.get(label, ()))
+ for label in question.criteria
+ }
+ if not matches:
+ raise ValueError('Choice needs at least one option')
+ best = max(matches.values())
+ winners = [label for label, count in matches.items() if count == best]
+ probabilities = {
+ label: 1 / len(winners) if label in winners else 0.0
+ for label in matches
+ }
+ answers[name] = ChoiceAnswer(
+ choice=winners[0],
+ probabilities=probabilities,
+ confidence=0.0 if best == 0 else 1 / len(winners))
+ elif isinstance(question, BooleanQuestion):
+ answers[name] = BooleanAnswer(probability=float(fraud_cue))
+ elif isinstance(question, ScoreQuestion):
+ if not question.criteria:
+ raise ValueError('Score needs at least one level')
+ score = min(3, len(question.criteria) - 1) if fraud_cue else 0
+ answers[name] = ScoreAnswer(
+ score=float(score),
+ probabilities={
+ level: float(level == score)
+ for level in range(len(question.criteria))
+ },
+ confidence=1.0)
+ else:
+ raise TypeError('Unsupported decision question: %r' % type(question))
+ return DecisionResponse(
+ answers=answers, model='local-rules', provider='local')
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/message_router.py
b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py
new file mode 100644
index 00000000000..8413d1e0b0e
--- /dev/null
+++
b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py
@@ -0,0 +1,199 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Route a short stream of messages to local files or BigQuery with Choice."""
+
+import argparse
+import json
+import os
+
+import apache_beam as beam
+from apache_beam.examples.inference.decision_models.local_model import
LocalDecisionModel
+from apache_beam.io import fileio
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import StandardOptions
+from apache_beam.runners import render
+from apache_beam.testing.test_stream import TestStream
+from apache_beam.transforms.window import FixedWindows
+from apache_beam.transforms.window import TimestampedValue
+
+DESTINATIONS = {
+ 'billing': 'Invoices, charges, payments, and refunds',
+ 'technical': 'Bugs, outages, and integration problems',
+ 'sales': 'Pricing, plans, and buying questions',
+}
+
+QUESTION = ChoiceQuestion(
+ instructions='Which team should handle this customer message?',
+ criteria=DESTINATIONS)
+
+SAMPLE_MESSAGES = [
+ {
+ 'event_id': 'msg-001',
+ 'message': 'My invoice has two charges for the same subscription. '
+ 'Please refund one.'
+ },
+ {
+ 'event_id': 'msg-002',
+ 'message': 'The API returns 503 for every request after I rotated '
+ 'my key.'
+ },
+ {
+ 'event_id': 'msg-003',
+ 'message': 'Can I get a quote for 50 seats next month?'
+ },
+]
+
+ROW_SCHEMA = (
+ 'event_id:STRING,message:STRING,decision:STRING,destination:STRING,'
+ 'confidence:FLOAT,probabilities_json:STRING,model:STRING,'
+ 'provider:STRING,request_id:STRING,latency_ms:FLOAT')
+
+
+def route_message(result, min_confidence):
+ event = result.state
+ response = result.response
+ answer = response.answers['destination']
+ if answer.choice not in DESTINATIONS:
+ raise ValueError('Decision model returned an unknown destination')
+ destination = (
+ answer.choice if answer.confidence is not None and
+ answer.confidence >= min_confidence else 'review')
+ return {
+ 'event_id': event['event_id'],
+ 'message': event['message'],
+ 'decision': answer.choice,
+ 'destination': destination,
+ 'confidence': answer.confidence,
+ 'probabilities_json': json.dumps(answer.probabilities, sort_keys=True),
+ 'model': response.model,
+ 'provider': response.provider,
+ 'request_id': response.request_id,
+ 'latency_ms': round(result.latency_ms, 2),
+ }
+
+
+class JsonSink(fileio.FileSink):
+ def open(self, fh):
+ self.fh = fh
+
+ def write(self, row):
+ self.fh.write(json.dumps(row, sort_keys=True).encode('utf-8') + b'\n')
+
+ def flush(self):
+ self.fh.flush()
+
+
+def decode_event(payload):
+ event = json.loads(payload.decode('utf-8'))
+ if not isinstance(event, dict) or not isinstance(event.get('message'), str):
+ raise ValueError('Pub/Sub event must be JSON with a message string')
+ if not isinstance(event.get('event_id'), str):
+ raise ValueError('Pub/Sub event must have an event_id string')
+ return event
+
+
+def table_for(row, dataset):
+ return '%s.messages_%s' % (dataset, row['destination'])
+
+
+def build_pipeline(pipeline, args):
+ if args.pubsub_subscription:
+ events = (
+ pipeline
+ | 'Read Pub/Sub' >>
+ beam.io.ReadFromPubSub(subscription=args.pubsub_subscription)
+ | 'Decode events' >> beam.Map(decode_event))
+ else:
+ stream = TestStream().add_elements([
+ TimestampedValue(message, index)
+ for index, message in enumerate(SAMPLE_MESSAGES)
+ ]).advance_watermark_to_infinity()
+ events = pipeline | 'Sample stream' >> stream
+
+ model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel()
+ rows = (
+ events
+ | 'Ask Choice' >> EvaluateDecisions(model, {'destination': QUESTION})
+ | 'Select destination' >> beam.Map(route_message, args.min_confidence)
+ | 'Window routed messages' >> beam.WindowInto(FixedWindows(60)))
+
+ if args.bq_dataset:
+ rows | 'Write routed BigQuery tables' >> beam.io.WriteToBigQuery(
+ table=lambda row: table_for(row, args.bq_dataset),
+ schema=ROW_SCHEMA,
+ method=beam.io.WriteToBigQuery.Method.STREAMING_INSERTS,
+ write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
+ create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED)
+ else:
+ rows | 'Write routed local files' >> fileio.WriteToFiles(
+ path=args.output_dir,
+ destination=lambda row: row['destination'],
+ file_naming=fileio.destination_prefix_naming('.jsonl'),
+ sink=JsonSink)
+ return rows
+
+
+def run(argv=None):
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument('--model', choices=('jev', 'local'), default='jev')
+ parser.add_argument('--min-confidence', type=float, default=0.5)
+ parser.add_argument('--pubsub-subscription')
+ parser.add_argument('--bq-dataset', help='BigQuery project:dataset')
+ parser.add_argument('--output-dir', help='Local output directory')
+ parser.add_argument('--graph', help='Render the Beam graph to .dot or .svg')
+ parser.add_argument(
+ '--graph-only',
+ action='store_true',
+ help='Render without running the job')
+ args, pipeline_args = parser.parse_known_args(argv)
+ if not 0 <= args.min_confidence <= 1:
+ parser.error('--min-confidence must be between 0 and 1')
+ if not args.bq_dataset and not args.output_dir:
+ parser.error('Set --output-dir or --bq-dataset')
+ if args.bq_dataset and args.output_dir:
+ parser.error('Use only one sink')
+ if (args.model == 'jev' and not args.graph_only and
+ not os.environ.get('TYPESAFE_API_KEY')):
+ parser.error('TYPESAFE_API_KEY is required for --model=jev')
+ if args.graph_only and not args.graph:
+ parser.error('--graph-only requires --graph')
+
+ options = PipelineOptions(pipeline_args)
+ options.view_as(StandardOptions).streaming = True
+ pipeline = beam.Pipeline(options=options)
+ build_pipeline(pipeline, args)
+ if args.graph:
+ render_options = render.RenderOptions([
+ '--render_output=' + args.graph,
+ '--render_leaf_composite_nodes=^Sample stream$',
+ '--render_leaf_composite_nodes=^Ask Choice$',
+ '--render_leaf_composite_nodes=^Write routed .*$',
+ ])
+ render.RenderRunner().run_pipeline(pipeline, render_options)
+ if args.graph_only:
+ return None
+ result = pipeline.run()
+ result.wait_until_finish()
+ return result
+
+
+if __name__ == '__main__':
+ run()
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg
b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg
new file mode 100644
index 00000000000..e37964497cf
--- /dev/null
+++
b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg
@@ -0,0 +1,98 @@
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one or more
+ contributor license agreements. See the NOTICE file distributed with
+ this work for additional information regarding copyright ownership.
+ The ASF licenses this file to You under the Apache License, Version 2.0
+ (the "License"); you may not use this file except in compliance with
+ the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+-->
+<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
+ "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd";>
+<!-- Generated by graphviz version 2.43.0 (0)
+ -->
+<!-- Title: G Pages: 1 -->
+<svg width="235pt" height="332pt"
+ viewBox="0.00 0.00 235.00 332.00" xmlns="http://www.w3.org/2000/svg";
xmlns:xlink="http://www.w3.org/1999/xlink";>
+<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 328)">
+<title>G</title>
+<polygon fill="white" stroke="transparent" points="-4,4 -4,-328 231,-328 231,4
-4,4"/>
+<!-- ref_AppliedPTransform_Sample-stream_2 -->
+<g id="node1" class="node">
+<title>ref_AppliedPTransform_Sample-stream_2</title>
+<g id="a_node1"><a
xlink:href="javascript:click('ref_AppliedPTransform_Sample-stream_2')"
xlink:title="Sample stream">
+<path fill="#fff6dd" stroke="#ff570b" d="M164,-324C164,-324 63,-324 63,-324
57,-324 51,-318 51,-312 51,-312 51,-300 51,-300 51,-294 57,-288 63,-288 63,-288
164,-288 164,-288 170,-288 176,-294 176,-300 176,-300 176,-312 176,-312
176,-318 170,-324 164,-324"/>
+<text text-anchor="middle" x="113.5" y="-302.3" font-family="Times,serif"
font-size="14.00">Sample stream</text>
+</a>
+</g>
+</g>
+<!-- ref_AppliedPTransform_Ask-Choice_3 -->
+<g id="node2" class="node">
+<title>ref_AppliedPTransform_Ask-Choice_3</title>
+<g id="a_node2"><a
xlink:href="javascript:click('ref_AppliedPTransform_Ask-Choice_3')"
xlink:title="Ask Choice">
+<path fill="#fff6dd" stroke="#ff570b" d="M149,-252C149,-252 78,-252 78,-252
72,-252 66,-246 66,-240 66,-240 66,-228 66,-228 66,-222 72,-216 78,-216 78,-216
149,-216 149,-216 155,-216 161,-222 161,-228 161,-228 161,-240 161,-240
161,-246 155,-252 149,-252"/>
+<text text-anchor="middle" x="113.5" y="-230.3" font-family="Times,serif"
font-size="14.00">Ask Choice</text>
+</a>
+</g>
+</g>
+<!--
ref_AppliedPTransform_Sample-stream_2->ref_AppliedPTransform_Ask-Choice_3
-->
+<g id="edge1" class="edge">
+<title>ref_AppliedPTransform_Sample-stream_2->ref_AppliedPTransform_Ask-Choice_3</title>
+<path fill="none" stroke="#ff570b" d="M113.5,-287.7C113.5,-279.98
113.5,-270.71 113.5,-262.11"/>
+<polygon fill="#ff570b" stroke="#ff570b" points="117,-262.1 113.5,-252.1
110,-262.1 117,-262.1"/>
+</g>
+<!-- ref_AppliedPTransform_Select-destination_5 -->
+<g id="node3" class="node">
+<title>ref_AppliedPTransform_Select-destination_5</title>
+<g id="a_node3"><a
xlink:href="javascript:click('ref_AppliedPTransform_Select-destination_5')"
xlink:title="Select destination">
+<path fill="#fff6dd" stroke="#ff570b" d="M174.5,-180C174.5,-180 52.5,-180
52.5,-180 46.5,-180 40.5,-174 40.5,-168 40.5,-168 40.5,-156 40.5,-156 40.5,-150
46.5,-144 52.5,-144 52.5,-144 174.5,-144 174.5,-144 180.5,-144 186.5,-150
186.5,-156 186.5,-156 186.5,-168 186.5,-168 186.5,-174 180.5,-180 174.5,-180"/>
+<text text-anchor="middle" x="113.5" y="-158.3" font-family="Times,serif"
font-size="14.00">Select destination</text>
+</a>
+</g>
+</g>
+<!--
ref_AppliedPTransform_Ask-Choice_3->ref_AppliedPTransform_Select-destination_5
-->
+<g id="edge2" class="edge">
+<title>ref_AppliedPTransform_Ask-Choice_3->ref_AppliedPTransform_Select-destination_5</title>
+<path fill="none" stroke="#ff570b" d="M113.5,-215.7C113.5,-207.98
113.5,-198.71 113.5,-190.11"/>
+<polygon fill="#ff570b" stroke="#ff570b" points="117,-190.1 113.5,-180.1
110,-190.1 117,-190.1"/>
+</g>
+<!-- ref_AppliedPTransform_Window-routed-messages_6 -->
+<g id="node4" class="node">
+<title>ref_AppliedPTransform_Window-routed-messages_6</title>
+<g id="a_node4"><a
xlink:href="javascript:click('ref_AppliedPTransform_Window-routed-messages_6')"
xlink:title="Window routed messages">
+<path fill="#fff6dd" stroke="#ff570b" d="M201.5,-108C201.5,-108 25.5,-108
25.5,-108 19.5,-108 13.5,-102 13.5,-96 13.5,-96 13.5,-84 13.5,-84 13.5,-78
19.5,-72 25.5,-72 25.5,-72 201.5,-72 201.5,-72 207.5,-72 213.5,-78 213.5,-84
213.5,-84 213.5,-96 213.5,-96 213.5,-102 207.5,-108 201.5,-108"/>
+<text text-anchor="middle" x="113.5" y="-86.3" font-family="Times,serif"
font-size="14.00">Window routed messages</text>
+</a>
+</g>
+</g>
+<!--
ref_AppliedPTransform_Select-destination_5->ref_AppliedPTransform_Window-routed-messages_6
-->
+<g id="edge3" class="edge">
+<title>ref_AppliedPTransform_Select-destination_5->ref_AppliedPTransform_Window-routed-messages_6</title>
+<path fill="none" stroke="#ff570b" d="M113.5,-143.7C113.5,-135.98
113.5,-126.71 113.5,-118.11"/>
+<polygon fill="#ff570b" stroke="#ff570b" points="117,-118.1 113.5,-108.1
110,-118.1 117,-118.1"/>
+</g>
+<!-- ref_AppliedPTransform_Write-routed-BigQuery-tables_7 -->
+<g id="node5" class="node">
+<title>ref_AppliedPTransform_Write-routed-BigQuery-tables_7</title>
+<g id="a_node5"><a
xlink:href="javascript:click('ref_AppliedPTransform_Write-routed-BigQuery-tables_7')"
xlink:title="Write routed BigQuery tables">
+<path fill="#fff6dd" stroke="#ff570b" d="M215,-36C215,-36 12,-36 12,-36 6,-36
0,-30 0,-24 0,-24 0,-12 0,-12 0,-6 6,0 12,0 12,0 215,0 215,0 221,0 227,-6
227,-12 227,-12 227,-24 227,-24 227,-30 221,-36 215,-36"/>
+<text text-anchor="middle" x="113.5" y="-14.3" font-family="Times,serif"
font-size="14.00">Write routed BigQuery tables</text>
+</a>
+</g>
+</g>
+<!--
ref_AppliedPTransform_Window-routed-messages_6->ref_AppliedPTransform_Write-routed-BigQuery-tables_7
-->
+<g id="edge4" class="edge">
+<title>ref_AppliedPTransform_Window-routed-messages_6->ref_AppliedPTransform_Write-routed-BigQuery-tables_7</title>
+<path fill="none" stroke="#ff570b" d="M113.5,-71.7C113.5,-63.98 113.5,-54.71
113.5,-46.11"/>
+<polygon fill="#ff570b" stroke="#ff570b" points="117,-46.1 113.5,-36.1
110,-46.1 117,-46.1"/>
+</g>
+</g>
+</svg>
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py
b/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py
new file mode 100644
index 00000000000..c5040599b08
--- /dev/null
+++
b/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py
@@ -0,0 +1,75 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for the policy boundary after a Choice decision."""
+
+import unittest
+
+from apache_beam.examples.inference.decision_models.message_router import
route_message
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import DecisionResult
+
+
+def result_for(answer):
+ return DecisionResult(
+ state={
+ 'event_id': '1', 'message': 'I have an account question'
+ },
+ response=DecisionResponse(
+ answers={'destination': answer}, model='fixed', provider='test'),
+ latency_ms=1.0)
+
+
+class RouteMessageTest(unittest.TestCase):
+ def test_uncertain_choice_goes_to_review(self):
+ result = result_for(
+ ChoiceAnswer(
+ choice='billing',
+ probabilities={
+ 'billing': 0.6, 'sales': 0.4
+ },
+ confidence=0.4))
+
+ row = route_message(result, min_confidence=0.65)
+
+ self.assertEqual('billing', row['decision'])
+ self.assertEqual('review', row['destination'])
+
+ def test_confident_choice_selects_destination(self):
+ result = result_for(ChoiceAnswer(choice='sales', confidence=0.65))
+
+ row = route_message(result, min_confidence=0.65)
+
+ self.assertEqual('sales', row['destination'])
+
+ def test_unknown_choice_cannot_select_a_table(self):
+ result = result_for(ChoiceAnswer(choice='other', confidence=1.0))
+
+ with self.assertRaisesRegex(ValueError, 'unknown destination'):
+ route_message(result, min_confidence=0.5)
+
+ def test_missing_confidence_goes_to_review(self):
+ result = result_for(ChoiceAnswer(choice='sales'))
+
+ row = route_message(result, min_confidence=0.5)
+
+ self.assertEqual('review', row['destination'])
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git
a/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt
b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt
new file mode 100644
index 00000000000..6b99d3ce92c
--- /dev/null
+++
b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt
@@ -0,0 +1,19 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# Optional client for apache_beam.ml.inference.typesafe_inference.
+-r ../../../ml/inference/typesafe_tests_requirements.txt
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py
b/sdks/python/apache_beam/examples/snippets/snippets_test.py
index 5a47d928901..8898cc8555f 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_test.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_test.py
@@ -50,6 +50,7 @@ from apache_beam.metrics.metric import MetricsFilter
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
+from apache_beam.options.value_provider import RuntimeValueProvider
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_stream import TestStream
from apache_beam.testing.util import assert_that
@@ -1509,8 +1510,11 @@ class SlowlyChangingSideInputsTest(unittest.TestCase):
class ValueProviderInfoTest(unittest.TestCase):
"""Tests for accessing value provider info after run."""
def test_accessing_valueprovider_info_after_run(self):
- with self.assertLogs(level='INFO') as log_capture:
- snippets.accessing_valueprovider_info_after_run()
+ # Prism's loopback worker runs in this process and does not initialize
+ # RuntimeValueProvider through sdk_worker_main.
+ with mock.patch.object(RuntimeValueProvider, 'runtime_options', {}):
+ with self.assertLogs(level='INFO') as log_capture:
+ snippets.accessing_valueprovider_info_after_run()
expected_log_message = "The string value is"
self.assertTrue(
any(expected_log_message in log for log in log_capture.output),
diff --git a/sdks/python/apache_beam/ml/inference/decision.py
b/sdks/python/apache_beam/ml/inference/decision.py
new file mode 100644
index 00000000000..fdfdc5801dc
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/decision.py
@@ -0,0 +1,246 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Typed questions and a Beam transform for decision models.
+
+A decision model evaluates named questions against each input state. Choice
+selects a label, Boolean estimates the probability of a true statement, and
+Score evaluates an ordered rubric. Adapters translate these types to a local
+model or a service's API.
+
+Use the same questions and transform when changing providers::
+
+ from apache_beam.ml.inference.decision import ChoiceQuestion
+ from apache_beam.ml.inference.decision import EvaluateDecisions
+
+ questions = {
+ 'team': ChoiceQuestion(
+ 'Which team should handle this message?',
+ {'billing': 'Invoices and refunds', 'support': 'Technical problems'})
+ }
+ decisions = events | EvaluateDecisions(model, questions)
+
+Each output is a :class:`DecisionResult` containing the original state,
+normalized answers, provider metadata, and request latency. Applications apply
+thresholds and route these results with ordinary Beam transforms.
+"""
+
+import abc
+import time
+from collections.abc import Mapping
+from collections.abc import Sequence
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+from typing import Union
+
+import apache_beam as beam
+from apache_beam.metrics import Metrics
+
+__all__ = [
+ 'ChoiceQuestion',
+ 'BooleanQuestion',
+ 'ScoreQuestion',
+ 'Question',
+ 'ChoiceAnswer',
+ 'BooleanAnswer',
+ 'ScoreAnswer',
+ 'Answer',
+ 'DecisionResponse',
+ 'DecisionResult',
+ 'DecisionModel',
+ 'EvaluateDecisions',
+]
+
+
+@dataclass(frozen=True)
+class ChoiceQuestion:
+ """Select one label from ``criteria``, a mapping of labels to
descriptions."""
+ instructions: str
+ criteria: Mapping[str, str]
+
+
+@dataclass(frozen=True)
+class BooleanQuestion:
+ """Estimate the probability that the statement in ``instructions`` is
true."""
+ instructions: str
+
+
+@dataclass(frozen=True)
+class ScoreQuestion:
+ """Score an ordered rubric whose ``criteria`` levels are indexed from zero.
+
+ Providers may interpolate between levels, returning a fractional score.
+ """
+ instructions: str
+ criteria: Sequence[str]
+
+
+Question = Union[ChoiceQuestion, BooleanQuestion, ScoreQuestion]
+
+
+@dataclass(frozen=True)
+class ChoiceAnswer:
+ """A selected label, with optional label probabilities and confidence.
+
+ Probabilities and confidence use [0, 1]. Confidence is a provider's estimate
+ of answer certainty; ``None`` means the provider did not supply it.
+ """
+ choice: str
+ probabilities: Mapping[str, float] = field(default_factory=dict)
+ confidence: float | None = None
+
+
+@dataclass(frozen=True)
+class BooleanAnswer:
+ """Probability in [0, 1] that the statement is true.
+
+ Applications choose their own threshold for acting on this probability.
+ """
+ probability: float
+
+
+@dataclass(frozen=True)
+class ScoreAnswer:
+ """Rubric score, optional probabilities by level index, and confidence.
+
+ The score uses the question's zero-based scale. Probabilities and confidence
+ use [0, 1]; ``None`` confidence means the provider did not supply it.
+ """
+ score: float
+ probabilities: Mapping[int, float] = field(default_factory=dict)
+ confidence: float | None = None
+
+
+Answer = Union[ChoiceAnswer, BooleanAnswer, ScoreAnswer]
+
+
+@dataclass(frozen=True)
+class DecisionResponse:
+ """One answer per question name, plus the model and provider identifiers."""
+ answers: Mapping[str, Answer]
+ model: str
+ provider: str
+ request_id: str | None = None
+
+
+@dataclass(frozen=True)
+class DecisionResult:
+ """An input state, its decision response, and evaluation time in
milliseconds."""
+ state: Any
+ response: DecisionResponse
+ latency_ms: float
+
+
+class DecisionModel(abc.ABC):
+ """Adapter contract for hosted and local decision models.
+
+ Implement :meth:`evaluate` to return the shared answer types. Override the
+ context manager methods to open and close clients or load model weights.
+ Constructors hold serializable configuration; Beam enters the context on
+ the worker and reuses it across elements and bundles on that DoFn instance.
+
+ An adapter owns provider timeouts, retries, and conversion to the shared
+ types. Input states must be serializable by Beam and accepted by the adapter.
+ """
+ def __enter__(self) -> 'DecisionModel':
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback) -> None:
+ pass
+
+ @abc.abstractmethod
+ def evaluate(
+ self, state: Any, questions: Mapping[str, Question]) -> DecisionResponse:
+ """Evaluate each named question against the same state.
+
+ Return exactly one answer of the corresponding type for each question.
+ Choice labels must belong to the question's criteria; score levels use
+ their zero-based indices. Raise an exception if evaluation fails.
+ """
+ raise NotImplementedError
+
+
+class EvaluateDecisions(beam.PTransform):
+ """Evaluate typed questions for every input element with a reusable model.
+
+ Args:
+ model: A serializable :class:`DecisionModel` adapter. Its context is
entered
+ once per DoFn instance and closed during teardown, which is best effort.
+ questions: A mapping of question names to Choice, Boolean, or Score
+ questions. All questions are evaluated against each input state.
+
+ Returns:
+ A PCollection of :class:`DecisionResult` objects, preserving each element's
+ timestamp and window. Calls run synchronously on the processing thread.
+ The ``EvaluateDecisions`` metrics namespace records ``requests``,
+ ``failures``, and the ``request_latency_ms`` distribution for successful
+ evaluations. Exceptions propagate to the runner.
+ """
+ def __init__(self, model: DecisionModel, questions: Mapping[str, Question]):
+ super().__init__()
+ self._model = model
+ self._questions = dict(questions)
+ self._exception_handling = None
+
+ def with_exception_handling(self, **kwargs):
+ """Return successful decisions and failed states as tagged PCollections.
+
+ Uses :meth:`apache_beam.ParDo.with_exception_handling` with the supplied
+ options. The default tags are ``good`` and ``bad``; bad records contain
+ the original state and exception details. Only processing failures are
+ caught. Model setup and teardown failures still propagate to the runner.
+ """
+ self._exception_handling = kwargs
+ return self
+
+ def expand(self, states):
+ evaluate = beam.ParDo(_EvaluateDecisionsDoFn(self._model, self._questions))
+ if self._exception_handling is not None:
+ evaluate = evaluate.with_exception_handling(**self._exception_handling)
+ return states | evaluate
+
+
+class _EvaluateDecisionsDoFn(beam.DoFn):
+ def __init__(self, model, questions):
+ self._model = model
+ self._questions = questions
+ self._active_model = None
+ self._requests = Metrics.counter('EvaluateDecisions', 'requests')
+ self._failures = Metrics.counter('EvaluateDecisions', 'failures')
+ self._latency = Metrics.distribution(
+ 'EvaluateDecisions', 'request_latency_ms')
+
+ def setup(self):
+ self._active_model = self._model.__enter__()
+
+ def process(self, state):
+ self._requests.inc()
+ started = time.perf_counter()
+ try:
+ response = self._active_model.evaluate(state, self._questions)
+ except Exception:
+ self._failures.inc()
+ raise
+ elapsed_ms = (time.perf_counter() - started) * 1000
+ self._latency.update(int(elapsed_ms))
+ yield DecisionResult(state, response, elapsed_ms)
+
+ def teardown(self):
+ if self._active_model is not None:
+ self._model.__exit__(None, None, None)
+ self._active_model = None
diff --git a/sdks/python/apache_beam/ml/inference/decision_test.py
b/sdks/python/apache_beam/ml/inference/decision_test.py
new file mode 100644
index 00000000000..f596acd6e3e
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/decision_test.py
@@ -0,0 +1,254 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+from unittest import mock
+
+import apache_beam as beam
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.decision import _EvaluateDecisionsDoFn
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.test_stream import TestStream
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+from apache_beam.transforms.window import FixedWindows
+
+_QUESTIONS = {
+ 'team': ChoiceQuestion(
+ 'Which team should handle this message?', {
+ 'billing': 'Invoices and refunds',
+ 'support': 'Technical problems',
+ }),
+ 'eligible': BooleanQuestion('Is this message eligible for automation?'),
+ 'severity': ScoreQuestion(
+ 'How severe is this message?', ('low', 'medium', 'high')),
+}
+
+
+def _response(model, provider, state, eligible, severity):
+ return DecisionResponse(
+ answers={
+ 'team': ChoiceAnswer(
+ 'support' if eligible > .5 else 'billing',
+ probabilities={
+ 'support': eligible,
+ 'billing': 1 - eligible,
+ }),
+ 'eligible': BooleanAnswer(eligible),
+ 'severity': ScoreAnswer(float(severity)),
+ },
+ model=model,
+ provider=provider,
+ request_id=state['id'])
+
+
+class _KeywordModel(DecisionModel):
+ def evaluate(self, state, questions):
+ if set(questions) != set(_QUESTIONS):
+ raise AssertionError('the model received a different question set')
+ eligible = .9 if 'bug' in state['text'] else .2
+ return _response(
+ 'keyword-v1', 'local-keyword', state, eligible, state['priority'])
+
+
+class _PriorityModel(DecisionModel):
+ def evaluate(self, state, questions):
+ if set(questions) != set(_QUESTIONS):
+ raise AssertionError('the model received a different question set')
+ eligible = .8 if state['priority'] >= 2 else .1
+ return _response(
+ 'priority-v1', 'local-priority', state, eligible, state['priority'])
+
+
+class _RejectingModel(_KeywordModel):
+ def evaluate(self, state, questions):
+ if state['id'] == 'bad':
+ raise ValueError('Invalid answer probability')
+ return super().evaluate(state, questions)
+
+
+class _SummarizeDecision(beam.DoFn):
+ def process(
+ self,
+ result,
+ timestamp=beam.DoFn.TimestampParam,
+ window=beam.DoFn.WindowParam):
+ answers = result.response.answers
+ yield (
+ result.state,
+ result.response,
+ answers['team'].choice,
+ answers['eligible'].probability >= .5,
+ answers['severity'].score,
+ timestamp.micros, (window.start.micros, window.end.micros))
+
+
+class _LifecycleModel(DecisionModel):
+ def __init__(self):
+ self.enter_count = 0
+ self.exit_calls = []
+ self.active = mock.Mock()
+
+ def __enter__(self):
+ self.enter_count += 1
+ return self.active
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ self.exit_calls.append((exc_type, exc_value, traceback))
+
+ def evaluate(self, state, questions):
+ raise AssertionError('the context returned model must evaluate')
+
+
+def _summarize_failure(
+ failure, timestamp=beam.DoFn.TimestampParam, window=beam.DoFn.WindowParam):
+ return (
+ failure[0],
+ failure[1][1],
+ timestamp.micros, (window.start.micros, window.end.micros))
+
+
+class DecisionTest(unittest.TestCase):
+ def test_dead_letter_output_preserves_failed_state_and_window(self):
+ good = {'id': 'good', 'text': 'bug in invoice', 'priority': 2}
+ bad = {'id': 'bad', 'text': 'refund request', 'priority': 0}
+ with TestPipeline() as pipeline:
+ states = (
+ pipeline
+ | beam.Create([
+ beam.window.TimestampedValue(good, 1),
+ beam.window.TimestampedValue(bad, 11),
+ ])
+ | beam.WindowInto(FixedWindows(10)))
+ decisions, failures = states | EvaluateDecisions(
+ _RejectingModel(), _QUESTIONS).with_exception_handling(
+ exc_class=ValueError)
+ assert_that(
+ decisions | beam.Map(lambda result: result.state),
+ equal_to([good]),
+ label='successful decisions')
+ assert_that(
+ failures | beam.Map(_summarize_failure),
+ equal_to([(
+ bad,
+ "ValueError('Invalid answer probability')",
+ 11_000_000, (10_000_000, 20_000_000))]),
+ label='failed decisions')
+
+ def
test_models_share_questions_downstream_and_preserve_window_metadata(self):
+ first = {'id': 'first', 'text': 'bug in invoice', 'priority': 2}
+ second = {'id': 'second', 'text': 'refund request', 'priority': 0}
+ stream = (
+ TestStream().add_elements([
+ beam.window.TimestampedValue(first, 1),
+ beam.window.TimestampedValue(second, 11),
+ ]).advance_watermark_to_infinity())
+
+ expected = {
+ 'keyword': [
+ (
+ first,
+ _response('keyword-v1', 'local-keyword', first, .9, 2),
+ 'support',
+ True,
+ 2.0,
+ 1_000_000, (0, 10_000_000)),
+ (
+ second,
+ _response('keyword-v1', 'local-keyword', second, .2, 0),
+ 'billing',
+ False,
+ 0.0,
+ 11_000_000, (10_000_000, 20_000_000)),
+ ],
+ 'priority': [
+ (
+ first,
+ _response('priority-v1', 'local-priority', first, .8, 2),
+ 'support',
+ True,
+ 2.0,
+ 1_000_000, (0, 10_000_000)),
+ (
+ second,
+ _response('priority-v1', 'local-priority', second, .1, 0),
+ 'billing',
+ False,
+ 0.0,
+ 11_000_000, (10_000_000, 20_000_000)),
+ ],
+ }
+
+ with TestPipeline() as pipeline:
+ states = pipeline | stream | beam.WindowInto(FixedWindows(10))
+ for label, model in (
+ ('keyword', _KeywordModel()),
+ ('priority', _PriorityModel()),
+ ):
+ decisions = states | f'{label} decisions' >> EvaluateDecisions(
+ model, _QUESTIONS)
+ summarized = decisions | f'{label} downstream' >> beam.ParDo(
+ _SummarizeDecision())
+ assert_that(
+ summarized, equal_to(expected[label]), label=f'check {label}')
+
+ def test_context_model_is_reused_across_bundles_and_closed(self):
+ model = _LifecycleModel()
+ model.active.evaluate.side_effect = [
+ _response('test', 'mock', {'id': 'one'}, .9, 1),
+ _response('test', 'mock', {'id': 'two'}, .1, 0),
+ ]
+ dofn = _EvaluateDecisionsDoFn(model, _QUESTIONS)
+
+ dofn.setup()
+ dofn.start_bundle()
+ first = list(dofn.process({'id': 'one'}))
+ dofn.finish_bundle()
+ dofn.start_bundle()
+ second = list(dofn.process({'id': 'two'}))
+ dofn.finish_bundle()
+ dofn.teardown()
+
+ self.assertEqual(1, model.enter_count)
+ self.assertEqual(2, model.active.evaluate.call_count)
+ self.assertEqual(1, len(model.exit_calls))
+ self.assertEqual((None, None, None), model.exit_calls[0])
+ self.assertEqual('one', first[0].state['id'])
+ self.assertEqual('two', second[0].state['id'])
+
+ def test_model_exception_propagates_from_process(self):
+ model = _LifecycleModel()
+ model.active.evaluate.side_effect = RuntimeError('provider failed')
+ dofn = _EvaluateDecisionsDoFn(model, _QUESTIONS)
+
+ dofn.setup()
+ with self.assertRaisesRegex(RuntimeError, 'provider failed'):
+ list(dofn.process({'id': 'failed'}))
+ dofn.teardown()
+ self.assertEqual(1, len(model.exit_calls))
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference.py
b/sdks/python/apache_beam/ml/inference/typesafe_inference.py
new file mode 100644
index 00000000000..130a0cafb5e
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_inference.py
@@ -0,0 +1,164 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Optional TypeSafe adapter for Beam decision models."""
+
+from collections.abc import Mapping
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import DecisionModel
+from apache_beam.ml.inference.decision import DecisionResponse
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+
+__all__ = ['JevDecisionModel']
+
+
+def _number(name, field, value, maximum=1):
+ if (isinstance(value, bool) or not isinstance(value, (int, float)) or
+ not 0 <= value <= maximum):
+ raise ValueError(
+ f'Invalid {field} for question {name!r}: '
+ f'expected a finite number in [0, {maximum}]')
+ return value
+
+
+def _probabilities(name, probabilities, keys, key_type):
+ if not isinstance(probabilities, Mapping):
+ raise ValueError(
+ f'Invalid probabilities for question {name!r}: '
+ 'expected a mapping')
+ for key, probability in probabilities.items():
+ if (not isinstance(key, key_type) or isinstance(key, bool) or
+ key not in keys):
+ raise ValueError(
+ f'Invalid probability key for question {name!r}: {key!r}')
+ _number(name, 'probability', probability)
+ return dict(probabilities)
+
+
+def _confidence(name, confidence):
+ return None if confidence is None else _number(name, 'confidence',
confidence)
+
+
+def _answer_type(name, answer, expected_type):
+ if not isinstance(answer, expected_type):
+ raise ValueError(
+ f'Invalid answer type for question {name!r}: '
+ f'expected {expected_type.__name__}')
+
+
+class JevDecisionModel(DecisionModel):
+ """Evaluate Beam decision questions with Jev through the TypeSafe SDK.
+
+ Install the optional SDK from ``typesafe_tests_requirements.txt`` and set
+ ``TYPESAFE_API_KEY`` on the workers. The client is created when Beam
+ enters the adapter context; its HTTP pool is reused until teardown.
+ Choice and Score map directly to the SDK primitives. Boolean maps to Noul,
+ whose ``noul`` field becomes :attr:`BooleanAnswer.probability`.
+
+ The SDK retries 429 and server errors with exponential backoff, jitter, and
+ ``Retry-After`` support. Calls and retries block the processing thread; they
+ do not coordinate rate limits across workers. Invalid answer keys, types,
+ labels, numeric ranges, or probability keys raise ``ValueError``.
+
+ Args:
+ model: TypeSafe model identifier, defaulting to ``jev-latest``.
+ timeout: Timeout in seconds for HTTP operations. The SDK handles retries.
+ """
+ def __init__(self, model='jev-latest', timeout=10.0):
+ self.model = model
+ self.timeout = timeout
+ self.client = None
+
+ def __enter__(self):
+ try:
+ from typesafe_sdk import TypeSafeClient
+ except ImportError as error:
+ raise ImportError(
+ 'Install typesafe-sdk to use JevDecisionModel') from error
+ self.client = TypeSafeClient(timeout=self.timeout)
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ self.client.close()
+ self.client = None
+
+ def evaluate(self, state, questions):
+ if self.client is None:
+ raise RuntimeError('Enter JevDecisionModel before evaluating questions')
+ import typesafe_sdk
+
+ sdk_questions = {}
+ for name, question in questions.items():
+ if isinstance(question, ChoiceQuestion):
+ sdk_questions[name] = typesafe_sdk.Choice(
+ instructions=question.instructions, criteria=question.criteria)
+ elif isinstance(question, BooleanQuestion):
+ sdk_questions[name] = typesafe_sdk.Noul(
+ instructions=question.instructions)
+ elif isinstance(question, ScoreQuestion):
+ sdk_questions[name] = typesafe_sdk.Score(
+ instructions=question.instructions, criteria=question.criteria)
+ else:
+ raise TypeError('Unsupported decision question: %r' % type(question))
+
+ try:
+ response = self.client.system_one(
+ state=state, questions=sdk_questions, model=self.model)
+ except typesafe_sdk.TypeSafeAPIResponseValidationError as error:
+ raise ValueError('Invalid Jev response') from error
+ if (not isinstance(response.answers, Mapping) or
+ response.answers.keys() != questions.keys()):
+ raise ValueError(
+ 'Jev answer question names must match the requested '
+ 'question names')
+ answers = {}
+ for name, question in questions.items():
+ answer = response.answers[name]
+ if isinstance(question, ChoiceQuestion):
+ _answer_type(name, answer, typesafe_sdk.ChoiceAnswer)
+ if (not isinstance(answer.choice, str) or
+ answer.choice not in question.criteria):
+ raise ValueError(
+ f'Invalid choice for question {name!r}: '
+ f'{answer.choice!r}')
+ answers[name] = ChoiceAnswer(
+ choice=answer.choice,
+ probabilities=_probabilities(
+ name, answer.probabilities, question.criteria, str),
+ confidence=_confidence(name, answer.confidence))
+ elif isinstance(question, BooleanQuestion):
+ _answer_type(name, answer, typesafe_sdk.NoulAnswer)
+ answers[name] = BooleanAnswer(
+ probability=_number(name, 'probability', answer.noul))
+ else:
+ _answer_type(name, answer, typesafe_sdk.ScoreAnswer)
+ answers[name] = ScoreAnswer(
+ score=_number(
+ name, 'score', answer.score, len(question.criteria) - 1),
+ probabilities=_probabilities(
+ name, answer.probabilities, range(len(question.criteria)),
int),
+ confidence=_confidence(name, answer.confidence))
+ return DecisionResponse(
+ answers=answers,
+ model=response.model,
+ provider='typesafe',
+ request_id=response.request_id)
diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py
b/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py
new file mode 100644
index 00000000000..acd2234f2dd
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py
@@ -0,0 +1,93 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Live Jev integration test, enabled with TYPESAFE_API_KEY."""
+
+import os
+import unittest
+
+import pytest
+
+try:
+ import typesafe_sdk
+except ImportError:
+ typesafe_sdk = None
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import EvaluateDecisions
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.test_stream import TestStream
+from apache_beam.testing.util import assert_that
+
+
[email protected]_postcommit
[email protected](typesafe_sdk is None, 'typesafe-sdk is not installed')
[email protected](
+ os.environ.get('TYPESAFE_API_KEY'),
+ 'TYPESAFE_API_KEY environment variable is not set')
+class JevInferenceIT(unittest.TestCase):
+ def test_typed_decisions_in_streaming_pipeline(self):
+ state = {
+ 'message': 'My invoice was charged twice. Please refund the duplicate.'
+ }
+ questions = {
+ 'team': ChoiceQuestion(
+ 'Which team should handle this message?',
+ {
+ 'billing': 'Invoices and refunds',
+ 'technical': 'Bugs and outages',
+ }),
+ 'refund': BooleanQuestion('Does this message ask for a refund?'),
+ 'urgency': ScoreQuestion(
+ 'How urgent is this message?', ('Routine', 'Urgent', 'Emergency')),
+ }
+
+ def check_results(results):
+ assert len(results) == 1, results
+ result = results[0]
+ assert result.state == state
+ assert result.response.provider == 'typesafe'
+ assert result.response.model
+ assert set(result.response.answers) == set(questions)
+ choice = result.response.answers['team']
+ boolean = result.response.answers['refund']
+ score = result.response.answers['urgency']
+ assert isinstance(choice, ChoiceAnswer)
+ assert choice.choice in questions['team'].criteria
+ assert isinstance(boolean, BooleanAnswer)
+ assert 0 <= boolean.probability <= 1
+ assert isinstance(score, ScoreAnswer)
+ assert 0 <= score.score <= 2
+ assert set(score.probabilities) <= {0, 1, 2}
+
+ stream = TestStream().add_elements([state]).advance_watermark_to_infinity()
+ with TestPipeline() as pipeline:
+ results = (
+ pipeline
+ | stream
+ | EvaluateDecisions(JevDecisionModel(), questions))
+ assert_that(results, check_results)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py
b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py
new file mode 100644
index 00000000000..5776fc9133e
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py
@@ -0,0 +1,256 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for TypeSafe's translation to the shared decision types."""
+
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+from parameterized import parameterized
+
+try:
+ import typesafe_sdk
+except ImportError:
+ raise unittest.SkipTest('typesafe-sdk is not installed')
+
+from apache_beam.ml.inference.decision import BooleanAnswer
+from apache_beam.ml.inference.decision import BooleanQuestion
+from apache_beam.ml.inference.decision import ChoiceAnswer
+from apache_beam.ml.inference.decision import ChoiceQuestion
+from apache_beam.ml.inference.decision import ScoreAnswer
+from apache_beam.ml.inference.decision import ScoreQuestion
+from apache_beam.ml.inference.typesafe_inference import JevDecisionModel
+
+
+def _choice_answer(**fields):
+ return typesafe_sdk.ChoiceAnswer.model_construct(
+ **{
+ 'choice': 'billing',
+ 'probabilities': {
+ 'billing': 1.0
+ },
+ 'confidence': 0.9,
+ **fields,
+ })
+
+
+def _score_answer(**fields):
+ return typesafe_sdk.ScoreAnswer.model_construct(
+ **{
+ 'score': 0.7,
+ 'probabilities': {
+ 0: 0.3, 1: 0.7
+ },
+ 'confidence': 0.6,
+ **fields,
+ })
+
+
+class JevDecisionModelTest(unittest.TestCase):
+ def test_normalizes_all_primitives_and_reuses_client(self):
+ client = mock.Mock()
+ client.system_one.return_value = SimpleNamespace(
+ answers={
+ 'team': _choice_answer(),
+ 'review': typesafe_sdk.NoulAnswer.model_construct(noul=0.8),
+ 'urgency': _score_answer(),
+ },
+ model='jev-test',
+ request_id='request-1')
+ questions = {
+ 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}),
+ 'review': BooleanQuestion('Does this need review?'),
+ 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']),
+ }
+ state = {'message': 'Review my invoice'}
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient',
+ return_value=client) as factory:
+ adapter = JevDecisionModel(model='jev-test', timeout=2.0)
+ factory.assert_not_called()
+ with adapter:
+ response = adapter.evaluate(state, questions)
+ adapter.evaluate(state, questions)
+ factory.assert_called_once_with(timeout=2.0)
+ client.close.assert_called_once_with()
+ self.assertEqual(2, client.system_one.call_count)
+ request = client.system_one.call_args.kwargs
+ self.assertEqual(state, request['state'])
+ self.assertEqual('jev-test', request['model'])
+ sent = request['questions']
+ self.assertIsInstance(sent['team'], typesafe_sdk.Choice)
+ self.assertEqual(questions['team'].criteria, sent['team'].criteria)
+ self.assertIsInstance(sent['review'], typesafe_sdk.Noul)
+ self.assertEqual(
+ questions['review'].instructions, sent['review'].instructions)
+ self.assertIsInstance(sent['urgency'], typesafe_sdk.Score)
+ self.assertEqual(questions['urgency'].criteria, sent['urgency'].criteria)
+ self.assertEqual(
+ ChoiceAnswer('billing', {'billing': 1.0}, 0.9),
+ response.answers['team'])
+ self.assertEqual(BooleanAnswer(0.8), response.answers['review'])
+ self.assertEqual(
+ ScoreAnswer(0.7, {
+ 0: 0.3, 1: 0.7
+ }, 0.6), response.answers['urgency'])
+ self.assertEqual('jev-test', response.model)
+ self.assertEqual('typesafe', response.provider)
+ self.assertEqual('request-1', response.request_id)
+
+ @parameterized.expand([
+ ('choice_label', 'team', _choice_answer(choice='sales')),
+ ('choice_type', 'team',
typesafe_sdk.NoulAnswer.model_construct(noul=.5)),
+ (
+ 'choice_probability_key',
+ 'team',
+ _choice_answer(probabilities={'sales': 1.0})),
+ ('choice_probabilities_type', 'team', _choice_answer(probabilities=[])),
+ (
+ 'choice_probability_range',
+ 'team',
+ _choice_answer(probabilities={'billing': 1.2})),
+ (
+ 'choice_probability_nan',
+ 'team',
+ _choice_answer(probabilities={'billing': float('nan')})),
+ ('choice_confidence', 'team', _choice_answer(confidence=1.3)),
+ ('choice_confidence_type', 'team', _choice_answer(confidence=True)),
+ (
+ 'boolean_range',
+ 'review',
+ typesafe_sdk.NoulAnswer.model_construct(noul=1.4)),
+ (
+ 'boolean_nan',
+ 'review',
+ typesafe_sdk.NoulAnswer.model_construct(noul=float('nan'))),
+ (
+ 'boolean_type',
+ 'review',
+ typesafe_sdk.NoulAnswer.model_construct(noul='0.5')),
+ ('boolean_answer_type', 'review', _choice_answer()),
+ ('score_range', 'urgency', _score_answer(score=3.5)),
+ ('score_negative', 'urgency', _score_answer(score=-.1)),
+ ('score_infinite', 'urgency', _score_answer(score=float('inf'))),
+ ('score_boolean', 'urgency', _score_answer(score=True)),
+ ('score_answer_type', 'urgency', _choice_answer()),
+ (
+ 'score_probability_key',
+ 'urgency',
+ _score_answer(probabilities={999: 1.0})),
+ (
+ 'score_probability_key_type',
+ 'urgency',
+ _score_answer(probabilities={'1': 1.0})),
+ (
+ 'score_probability_fractional_key',
+ 'urgency',
+ _score_answer(probabilities={.5: 1.0})),
+ (
+ 'score_probability_negative',
+ 'urgency',
+ _score_answer(probabilities={0: -.1})),
+ ('score_confidence', 'urgency', _score_answer(confidence=-.1)),
+ ])
+ def test_rejects_malformed_answer(self, unused_name, name, answer):
+ questions = {
+ 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}),
+ 'review': BooleanQuestion('Does this need review?'),
+ 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']),
+ }
+ client = mock.Mock()
+ client.system_one.return_value = SimpleNamespace(
+ answers={name: answer}, model='jev-test', request_id='request-1')
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient',
return_value=client):
+ with JevDecisionModel() as adapter:
+ with self.assertRaisesRegex(ValueError, name):
+ adapter.evaluate('Hello', {name: questions[name]})
+
+ @parameterized.expand([
+ ('missing', {}),
+ (
+ 'extra',
+ {
+ 'review': typesafe_sdk.NoulAnswer.model_construct(noul=.5),
+ 'unexpected': typesafe_sdk.NoulAnswer.model_construct(noul=.5),
+ }),
+ ('wrong_container', []),
+ ])
+ def test_rejects_answer_keys(self, unused_name, answers):
+ client = mock.Mock()
+ client.system_one.return_value = SimpleNamespace(
+ answers=answers, model='jev-test', request_id='request-1')
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient',
return_value=client):
+ with JevDecisionModel() as adapter:
+ with self.assertRaisesRegex(ValueError, 'question names'):
+ adapter.evaluate(
+ 'Hello', {'review': BooleanQuestion('Needs review?')})
+
+ def test_accepts_boundaries_and_optional_probabilities(self):
+ questions = {
+ 'team': ChoiceQuestion('Pick a team', {'billing': 'Invoices'}),
+ 'review': BooleanQuestion('Does this need review?'),
+ 'urgency': ScoreQuestion('How urgent is it?', ['routine', 'urgent']),
+ }
+ client = mock.Mock()
+ for probability in (0, 1):
+ with self.subTest(probability=probability):
+ client.system_one.return_value = SimpleNamespace(
+ answers={
+ 'team': _choice_answer(probabilities={}, confidence=None),
+ 'review': typesafe_sdk.NoulAnswer.model_construct(
+ noul=probability),
+ 'urgency': _score_answer(
+ score=probability, probabilities={}, confidence=None),
+ },
+ model='jev-test',
+ request_id='request-1')
+ with mock.patch.object(typesafe_sdk,
+ 'TypeSafeClient',
+ return_value=client):
+ with JevDecisionModel() as adapter:
+ response = adapter.evaluate('Hello', questions)
+ self.assertEqual(probability, response.answers['review'].probability)
+ self.assertEqual(probability, response.answers['urgency'].score)
+ self.assertIsNone(response.answers['team'].confidence)
+
+ def test_provider_error_propagates_and_client_closes(self):
+ client = mock.Mock()
+ client.system_one.side_effect = TimeoutError('request timed out')
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient',
return_value=client):
+ with self.assertRaisesRegex(TimeoutError, 'request timed out'):
+ with JevDecisionModel() as adapter:
+ adapter.evaluate(
+ 'Hello', {'review': BooleanQuestion('Needs review?')})
+ client.close.assert_called_once_with()
+
+ def test_sdk_response_validation_error_is_a_value_error(self):
+ error = typesafe_sdk.TypeSafeAPIResponseValidationError(
+ status=200, body={}, headers={}, field_path='answers.review.noul')
+ client = mock.Mock()
+ client.system_one.side_effect = error
+ with mock.patch.object(typesafe_sdk, 'TypeSafeClient',
return_value=client):
+ with JevDecisionModel() as adapter:
+ with self.assertRaisesRegex(ValueError,
+ 'Invalid Jev response') as raised:
+ adapter.evaluate(
+ 'Hello', {'review': BooleanQuestion('Needs review?')})
+ self.assertIs(error, raised.exception.__cause__)
+ client.close.assert_called_once_with()
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git
a/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt
b/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt
new file mode 100644
index 00000000000..982081593f5
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt
@@ -0,0 +1,18 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+typesafe-sdk>=0.7.1,<0.8
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 08c5964898d..a4888617a52 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -613,6 +613,7 @@ if __name__ == '__main__':
] + ml_base + qdrant_dependency,
'p312_ml_test': [
'datatable',
+ 'typesafe-sdk>=0.7.1,<0.8',
] + ml_base + qdrant_dependency,
# maintainer: milvus tests only run with this extension. Make sure it
# is covered by docker-in-docker test when changing py version
