From 54f6ca63d01d0d69120e6c53bd6754053208a013 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Tue, 22 Sep 2026 02:05:58 -0400 Subject: [PATCH 1/7] [Python] Add decision model routing examples --- .../inference/decision_models/README.md | 221 ++++++++++++++++++ .../inference/decision_models/__init__.py | 18 ++ .../inference/decision_models/fraud_review.py | 111 +++++++++ .../decision_models/message_router.py | 219 +++++++++++++++++ .../decision_models/message_router.svg | 99 ++++++++ .../decision_models/message_router_test.py | 71 ++++++ .../inference/decision_models/model.py | 213 +++++++++++++++++ .../decision_models/requirements.txt | 21 ++ 8 files changed, 973 insertions(+) create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/README.md create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/__init__.py create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/message_router.py create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/message_router.svg create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/model.py create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/requirements.txt 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 000000000000..4b73cb9a5021 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md @@ -0,0 +1,221 @@ + + +# Beam decision model examples + +This small proof of concept connects Beam's Python `RequestResponseIO` with +structured decision models. The main example routes a short stream of support +messages with a Choice question. A second example uses a Noul question to flag +messages that deserve fraud review. + +The examples keep the Beam graph independent of a model vendor. A decision +model accepts JSON-like `state` and a set of typed questions, then returns +typed answers plus request metadata. `JevDecisionModel` is the TypeSafe adapter +used by the live path. `LocalDecisionModel` is a deterministic adapter for +running the graph without a network call. Future Laya or Kev integrations can +implement the same adapter contract; this PoC makes no claims about APIs for +those systems. + +## Architecture + +```mermaid +flowchart LR + source["TestStream
or Pub/Sub"] --> decode["Decode JSON event"] + decode --> rr["RequestResponseIO
DecisionModel Caller"] + rr --> choice["Choice answer
destination + probabilities + confidence"] + choice --> gate{"confidence >= threshold?"} + gate -->|yes| route["billing / technical / sales"] + gate -->|no| review["review"] + route --> sink["local JSONL files
or dynamic BigQuery tables"] + review --> sink +``` + +The rendered graph below was generated from the same pipeline with +`--graph-only --bq-dataset=PROJECT:DATASET`. Graph-only mode constructs and +renders the BigQuery sink but does not contact BigQuery or write rows. + +![Rendered Beam message router graph](message_router.svg) + +`RequestResponseIO` isolates the external call behind Beam's request/response +boundary. The pipeline remains responsible for policy: it validates the +returned label, applies the confidence threshold, and chooses the sink. This +is useful for a decision model because the selected label answers *what* the +model thinks, while confidence answers whether the application should act. +Choice and Score also expose the full probability distribution. Noul exposes +the probability of a yes/true answer and does not have a separate confidence +field. + +## Install + +From a Beam checkout, in an activated Python virtual environment, install the +Python SDK and the optional Jev client: + +```sh +python3 -m pip install -e sdks/python +python3 -m pip install -r sdks/python/apache_beam/examples/inference/decision_models/requirements.txt +``` + +The `typesafe-sdk` requirement is only needed for `--model=jev`. The local +adapter does not make an API request and is useful for graph and sink checks. +Install Beam's GCP extra before using `--bq-dataset`, including graph-only mode: + +```sh +python3 -m pip install -e 'sdks/python[gcp]' +``` + +## Main demo: route messages with Choice + +The default source is a three-element `TestStream`, so the demo finishes +quickly. The local adapter produces the same typed response shape without +credentials: + +```sh +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/pipeline.svg" +find "$output_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \; +``` + +The graph is exported by Beam's Python `RenderRunner` while the pipeline is +constructed. Use a `.svg` or `.png` extension when Graphviz's `dot` executable +is installed. A `.dot` output is useful when only the raw graph is available. +This is a graph of the DirectRunner-compatible pipeline construction, not a +runtime profile. + +To render the graph without making Jev or BigQuery calls, use +`--graph-only`. The sink option is still required because it is part of the +constructed graph: + +```sh +python3 -m apache_beam.examples.inference.decision_models.message_router \ + --model=jev --graph-only \ + --bq-dataset=PROJECT:DATASET \ + --graph=/tmp/beam-decision-model.dot +``` + +The `.dot` export works without Graphviz. SVG and PNG exports require `dot` on +`PATH`. + +For the Jev path, put the key in the environment and run the same short +stream: + +```sh +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/pipeline.svg" +``` + +One three-message Beam DirectRunner smoke run against Jev measured the +following caller wall time with the local JSONL sink: + +| message | routed destination | caller latency | +| --- | --- | ---: | +| billing message | `billing` | 406.26 ms | +| technical message | `technical` | 249.70 ms | +| sales message | `sales` | 226.34 ms | + +These are observations from one local environment with `n=3` Choice calls and +one local file sink, not a latency guarantee or a benchmark. They include the +provider call inside the Beam caller and exclude BigQuery write latency. + +The key is read by the SDK through `TYPESAFE_API_KEY`; it is never a pipeline +element or a command-line argument. Do not put it in source, shell history, +logs, or a checked-in options file. + +To use a real streaming source, provide a Pub/Sub subscription containing JSON +objects with `event_id` and `message` string fields: + +```sh +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`, the destination is selected in code and +rows are appended to `messages_billing`, `messages_technical`, +`messages_sales`, or `messages_review` in that dataset. The process needs +Application Default Credentials with Pub/Sub subscriber access and permission +to create/write the BigQuery tables. BigQuery has not been live-verified by +this PoC. For a local run, omit `--bq-dataset` and use `--output-dir` instead. +The sink uses streaming inserts and Beam's normal at-least-once delivery +behavior. A production design needs an idempotency key and an explicit +retry/dead-letter policy. + +## Smaller demo: fraud review with Noul + +Noul asks one yes/no question and returns the probability that the answer is +yes. The example uses that value as a deliberately visible policy threshold: + +```sh +python3 -m apache_beam.examples.inference.decision_models.fraud_review \ + --model=local --primitive=noul + +export TYPESAFE_API_KEY='replace-with-your-key' +python3 -m apache_beam.examples.inference.decision_models.fraud_review \ + --model=jev --primitive=noul +``` + +Rows include `fraud_review_probability` and `needs_review`. A probability near +0.5 is uncertainty, not confidence: Noul has no separate confidence value. +For comparison, run `--primitive=score` or `--primitive=both` to see a rubric +score, its probability distribution, and its confidence. Combining several +atomic questions in one request is the intended way to add policy dimensions +without asking one broad, ambiguous question. + +A three-element Beam DirectRunner Jev run with both questions produced Noul +probabilities `0.46` for the billing-address message, `0.32` for the +duplicate-charge message, and `0.97` for the request to bypass verification. +With the `0.7` threshold, only the last row was flagged. The Score answers +were `0.80`, `0.88`, and `3.00`, respectively. These are sample values. + +## Short benchmark + +Each demo records request elapsed time in `latency_ms`. Run only the small +`TestStream` to measure a handful of calls and inspect the generated JSONL: + +```sh +benchmark_dir="$(mktemp -d /tmp/beam-decision-model-benchmark.XXXXXX)" +time python3 -m apache_beam.examples.inference.decision_models.message_router \ + --model=jev \ + --output-dir="$benchmark_dir" +rg -I -o '"latency_ms": [0-9.]*' "$benchmark_dir" \ + | sort -t: -k2n +``` + +This is an illustrative smoke benchmark, not a performance claim. It mixes +API/network latency, SDK work, and one local runner invocation, and three +messages are too few for a stable percentile. For a meaningful comparison, +repeat with a fixed corpus, record the model and request IDs, and report the +sample size and percentile method. Do not commit credentials or benchmark +output. + +## Scope and next steps + +This PoC demonstrates the adapter boundary and one Beam integration point. It +does not provide multi-message batching, model warmup, circuit breaking, +schema evolution, or production deduplication. The Beam repeater is disabled; +the Jev SDK handles provider retries. A follow-up can add quota controls, +route metrics, and a runner-specific integration test before a production sink. 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 000000000000..fc00cf8225fd --- /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 000000000000..cfc82c021591 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py @@ -0,0 +1,111 @@ +# +# 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 time + +import apache_beam as beam +from apache_beam.io.requestresponse import Caller +from apache_beam.io.requestresponse import RequestResponseIO +from apache_beam.options.pipeline_options import PipelineOptions + +from apache_beam.examples.inference.decision_models.model import JevDecisionModel +from apache_beam.examples.inference.decision_models.model import LocalDecisionModel +from apache_beam.examples.inference.decision_models.model import NoulQuestion +from apache_beam.examples.inference.decision_models.model import ScoreQuestion + +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 = NoulQuestion( + 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')) + + +class ReviewFraudCues(Caller): + def __init__(self, model, primitive): + self.model = model + self.primitive = primitive + + def __enter__(self): + self.model.__enter__() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + return self.model.__exit__(exc_type, exc_val, exc_tb) + + def __call__(self, message): + questions = {} + if self.primitive in ('noul', 'both'): + questions['fraud_cue'] = FRAUD_NOUL + if self.primitive in ('score', 'both'): + questions['risk_level'] = FRAUD_SCORE + started = time.perf_counter() + response = self.model.evaluate( + state={'message': message}, questions=questions) + row = { + 'message': message, + 'model': response.model, + 'provider': response.provider, + 'latency_ms': round((time.perf_counter() - started) * 1000, 2), + } + if 'fraud_cue' in response.answers: + probability = response.answers['fraud_cue'].noul + 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 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') + 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() + with beam.Pipeline(options=PipelineOptions(pipeline_args)) as pipeline: + ( + pipeline + | 'Sample payment messages' >> beam.Create(SAMPLE_MESSAGES) + | 'Evaluate fraud cues' >> RequestResponseIO( + ReviewFraudCues(model, args.primitive), timeout=15, repeater=None) + | '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/message_router.py b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py new file mode 100644 index 000000000000..cf341cf87254 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py @@ -0,0 +1,219 @@ +# +# 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 time + +import apache_beam as beam +from apache_beam.io import fileio +from apache_beam.io.requestresponse import Caller +from apache_beam.io.requestresponse import RequestResponseIO +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 + +from apache_beam.examples.inference.decision_models.model import ChoiceQuestion +from apache_beam.examples.inference.decision_models.model import JevDecisionModel +from apache_beam.examples.inference.decision_models.model import LocalDecisionModel + +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') + + +class ClassifyMessage(Caller): + """Use any DecisionModel as a RequestResponseIO caller.""" + def __init__(self, model, min_confidence): + self.model = model + self.min_confidence = min_confidence + + def __enter__(self): + self.model.__enter__() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + return self.model.__exit__(exc_type, exc_val, exc_tb) + + def __call__(self, event): + started = time.perf_counter() + response = self.model.evaluate( + state={'message': event['message']}, + questions={'destination': QUESTION}) + latency_ms = (time.perf_counter() - started) * 1000 + 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 >= self.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(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 + | 'Classify with decision model' >> RequestResponseIO( + ClassifyMessage(model, args.min_confidence), + timeout=15, + repeater=None) + | '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=^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 000000000000..27bc6575092f --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg @@ -0,0 +1,99 @@ + + + + + + + +G + + +cluster_ref_AppliedPTransform_Classify-with-decision-model_3 + + +Classify with decision model + + + + +cluster_ref_AppliedPTransform_Classify-with-decision-model-_Call_4 + + +_Call + + + + + +ref_AppliedPTransform_Classify-with-decision-model-_Call-ParDo-_CallDoFn-_5 + + +ParDo(_CallDoFn) + + + + + +ref_AppliedPTransform_Window-routed-messages_6 + + +Window routed messages + + + + + +ref_AppliedPTransform_Classify-with-decision-model-_Call-ParDo-_CallDoFn-_5->ref_AppliedPTransform_Window-routed-messages_6 + + + + + +ref_AppliedPTransform_Sample-stream_2 + + +Sample stream + + + + + +ref_AppliedPTransform_Sample-stream_2->ref_AppliedPTransform_Classify-with-decision-model-_Call-ParDo-_CallDoFn-_5 + + + + + +ref_AppliedPTransform_Write-routed-BigQuery-tables_7 + + +Write routed BigQuery tables + + + + + +ref_AppliedPTransform_Window-routed-messages_6->ref_AppliedPTransform_Write-routed-BigQuery-tables_7 + + + + + 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 000000000000..9f9b6311af49 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py @@ -0,0 +1,71 @@ +# +# 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 ClassifyMessage +from apache_beam.examples.inference.decision_models.model import ChoiceAnswer +from apache_beam.examples.inference.decision_models.model import DecisionResponse + + +class FixedModel: + def __init__(self, answer): + self.answer = answer + + def evaluate(self, state, questions): + return DecisionResponse( + answers={'destination': self.answer}, model='fixed', provider='test') + + +class ClassifyMessageTest(unittest.TestCase): + def test_uncertain_choice_goes_to_review(self): + caller = ClassifyMessage( + FixedModel( + ChoiceAnswer( + choice='billing', + probabilities={ + 'billing': 0.6, 'sales': 0.4 + }, + confidence=0.4)), + min_confidence=0.65) + + row = caller({'event_id': '1', 'message': 'I have an account question'}) + + self.assertEqual('billing', row['decision']) + self.assertEqual('review', row['destination']) + + def test_unknown_choice_cannot_select_a_table(self): + caller = ClassifyMessage( + FixedModel(ChoiceAnswer(choice='other', confidence=1.0)), + min_confidence=0.5) + + with self.assertRaisesRegex(ValueError, 'unknown destination'): + caller({'event_id': '1', 'message': 'Hello'}) + + def test_missing_confidence_goes_to_review(self): + caller = ClassifyMessage( + FixedModel(ChoiceAnswer(choice='sales')), min_confidence=0.5) + + row = caller({'event_id': '1', 'message': 'What does it cost?'}) + + self.assertEqual('review', row['destination']) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/examples/inference/decision_models/model.py b/sdks/python/apache_beam/examples/inference/decision_models/model.py new file mode 100644 index 000000000000..5808942a3259 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/model.py @@ -0,0 +1,213 @@ +# +# 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 decision questions and a small Jev adapter for Beam examples.""" + +from dataclasses import dataclass +from dataclasses import field +from typing import Any +from typing import Mapping +from typing import Protocol +from typing import Sequence +from typing import Union + + +@dataclass(frozen=True) +class ChoiceQuestion: + instructions: str + criteria: Mapping[str, str] + + +@dataclass(frozen=True) +class NoulQuestion: + instructions: str + + +@dataclass(frozen=True) +class ScoreQuestion: + instructions: str + criteria: Sequence[str] + + +Question = Union[ChoiceQuestion, NoulQuestion, ScoreQuestion] + + +@dataclass(frozen=True) +class ChoiceAnswer: + choice: str + probabilities: Mapping[str, float] = field(default_factory=dict) + confidence: float | None = None + + +@dataclass(frozen=True) +class NoulAnswer: + noul: float + + +@dataclass(frozen=True) +class ScoreAnswer: + score: float + probabilities: Mapping[str, float] = field(default_factory=dict) + confidence: float = 0.0 + + +Answer = Union[ChoiceAnswer, NoulAnswer, ScoreAnswer] + + +@dataclass(frozen=True) +class DecisionResponse: + answers: Mapping[str, Answer] + model: str + provider: str + request_id: str | None = None + + +class DecisionModel(Protocol): + """A provider owns evaluation; Beam owns routing and persistence.""" + def __enter__(self) -> 'DecisionModel': + ... + + def __exit__(self, exc_type, exc_value, traceback) -> None: + ... + + def evaluate( + self, state: Any, questions: Mapping[str, Question]) -> DecisionResponse: + ... + + +class JevDecisionModel: + """TypeSafe SDK adapter. Each Beam caller creates one client per worker.""" + def __init__(self, model='jev-latest'): + self.model = model + 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() + 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') + from typesafe_sdk import Choice + from typesafe_sdk import Noul + from typesafe_sdk import Score + + sdk_questions = {} + for name, question in questions.items(): + if isinstance(question, ChoiceQuestion): + sdk_questions[name] = Choice( + instructions=question.instructions, criteria=question.criteria) + elif isinstance(question, NoulQuestion): + sdk_questions[name] = Noul(instructions=question.instructions) + elif isinstance(question, ScoreQuestion): + sdk_questions[name] = Score( + instructions=question.instructions, criteria=question.criteria) + else: + raise TypeError('Unsupported decision question: %r' % type(question)) + + response = self.client.system_one( + state=state, questions=sdk_questions, model=self.model) + answers = {} + for name, question in questions.items(): + answer = response.answers[name] + if isinstance(question, ChoiceQuestion): + answers[name] = ChoiceAnswer( + choice=answer.choice, + probabilities=dict(answer.probabilities), + confidence=answer.confidence) + elif isinstance(question, NoulQuestion): + answers[name] = NoulAnswer(noul=answer.noul) + else: + answers[name] = ScoreAnswer( + score=answer.score, + probabilities={ + str(level): probability + for level, probability in answer.probabilities.items() + }, + confidence=answer.confidence) + return DecisionResponse( + answers=answers, + model=response.model, + provider='typesafe', + request_id=response.request_id) + + +_ROUTE_TERMS = { + 'billing': ('invoice', 'charge', 'refund', 'payment', 'subscription'), + 'technical': ('api', '503', 'error', 'integration', 'outage'), + 'sales': ('quote', 'seat', 'pricing', 'price', 'plan'), +} + + +class LocalDecisionModel: + """Transparent rules for running these examples without network access.""" + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return None + + 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, NoulQuestion): + answers[name] = NoulAnswer(noul=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={ + str(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/requirements.txt b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt new file mode 100644 index 000000000000..d7280f85b93e --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt @@ -0,0 +1,21 @@ +# +# 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: required only for --model=jev. The local adapter has no extra +# dependency. Keep this compatible range aligned with the Jev SDK API used by +# model.py. +typesafe-sdk>=0.7.1,<0.8 From 19420631647647a65974ba40ec16ea5e6b5e337f Mon Sep 17 00:00:00 2001 From: bvolpato Date: Tue, 22 Sep 2026 08:56:39 -0400 Subject: [PATCH 2/7] [Python] Reuse decision model clients in a lightweight DoFn --- .../inference/decision_models/README.md | 216 ++++++++---------- .../inference/decision_models/fraud_review.py | 73 +++--- .../decision_models/message_router.py | 69 ++---- .../decision_models/message_router.svg | 97 ++++---- .../decision_models/message_router_test.py | 58 ++--- .../inference/decision_models/model.py | 7 +- .../inference/decision_models/transforms.py | 66 ++++++ .../decision_models/transforms_test.py | 62 +++++ 8 files changed, 363 insertions(+), 285 deletions(-) create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/transforms.py create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py diff --git a/sdks/python/apache_beam/examples/inference/decision_models/README.md b/sdks/python/apache_beam/examples/inference/decision_models/README.md index 4b73cb9a5021..ebc83f2c9c7b 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/README.md +++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md @@ -19,61 +19,59 @@ # Beam decision model examples -This small proof of concept connects Beam's Python `RequestResponseIO` with -structured decision models. The main example routes a short stream of support -messages with a Choice question. A second example uses a Noul question to flag -messages that deserve fraud review. - -The examples keep the Beam graph independent of a model vendor. A decision -model accepts JSON-like `state` and a set of typed questions, then returns -typed answers plus request metadata. `JevDecisionModel` is the TypeSafe adapter -used by the live path. `LocalDecisionModel` is a deterministic adapter for -running the graph without a network call. Future Laya or Kev integrations can -implement the same adapter contract; this PoC makes no claims about APIs for -those systems. +These examples connect Beam pipelines to typed decision models. The main +pipeline uses Choice to route support messages. The fraud pipeline uses Noul +and optionally Score. `LocalDecisionModel` runs the graph without a network +call, while `JevDecisionModel` adapts the TypeSafe SDK. ## Architecture ```mermaid flowchart LR source["TestStream
or Pub/Sub"] --> decode["Decode JSON event"] - decode --> rr["RequestResponseIO
DecisionModel Caller"] - rr --> choice["Choice answer
destination + probabilities + confidence"] - choice --> gate{"confidence >= threshold?"} - gate -->|yes| route["billing / technical / sales"] - gate -->|no| review["review"] - route --> sink["local JSONL files
or dynamic BigQuery tables"] + decode --> evaluate["EvaluateDecisions
(beam.DoFn)"] + evaluate --> result["DecisionResult"] + result --> route["route_message
validate Choice + confidence"] + route --> sink["JSONL files
or dynamic BigQuery tables"] + route --> review["review"] review --> sink ``` -The rendered graph below was generated from the same pipeline with -`--graph-only --bq-dataset=PROJECT:DATASET`. Graph-only mode constructs and -renders the BigQuery sink but does not contact BigQuery or write rows. +`EvaluateDecisions` accepts a model and typed questions. `setup` enters the +model context and `teardown` closes it. Each DoFn instance owns one context. +For Jev, that context creates one TypeSafe SDK client with a persistent HTTP +connection pool used across elements and bundles. `process` calls `evaluate` +synchronously and yields `DecisionResult(state, response, latency_ms)`, +preserving Beam backpressure without a per-element executor. -![Rendered Beam message router graph](message_router.svg) +`JevDecisionModel` uses a 10 second HTTP operation timeout and the SDK's retry +behavior. The SDK keepalive default is 5 seconds, and connection reuse follows +the SDK pool lifecycle. + +`route_message` checks that the Choice answer is `billing`, `technical`, or +`sales`. A missing or low confidence answer goes to `review`. The selected +destination controls the JSONL file name or BigQuery table name. -`RequestResponseIO` isolates the external call behind Beam's request/response -boundary. The pipeline remains responsible for policy: it validates the -returned label, applies the confidence threshold, and chooses the sink. This -is useful for a decision model because the selected label answers *what* the -model thinks, while confidence answers whether the application should act. -Choice and Score also expose the full probability distribution. Noul exposes -the probability of a yes/true answer and does not have a separate confidence -field. +![Rendered Beam message router graph](message_router.svg) ## Install -From a Beam checkout, in an activated Python virtual environment, install the -Python SDK and the optional Jev client: +From a Beam checkout, activate a Python virtual environment and install the +core SDK: ```sh python3 -m pip install -e sdks/python -python3 -m pip install -r sdks/python/apache_beam/examples/inference/decision_models/requirements.txt ``` -The `typesafe-sdk` requirement is only needed for `--model=jev`. The local -adapter does not make an API request and is useful for graph and sink checks. -Install Beam's GCP extra before using `--bq-dataset`, including graph-only mode: +Install the optional Jev client for `--model=jev`: + +```sh +python3 -m pip install -r \ + sdks/python/apache_beam/examples/inference/decision_models/requirements.txt +``` + +Install Beam's GCP extra for Pub/Sub or BigQuery, including graph-only renders +that contain a BigQuery sink: ```sh python3 -m pip install -e 'sdks/python[gcp]' @@ -81,41 +79,36 @@ python3 -m pip install -e 'sdks/python[gcp]' ## Main demo: route messages with Choice -The default source is a three-element `TestStream`, so the demo finishes -quickly. The local adapter produces the same typed response shape without -credentials: +The default source is a three element `TestStream`. The local adapter writes +one JSONL file per destination: ```sh 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/pipeline.svg" + --graph="$output_dir/message_router.svg" find "$output_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \; ``` -The graph is exported by Beam's Python `RenderRunner` while the pipeline is -constructed. Use a `.svg` or `.png` extension when Graphviz's `dot` executable -is installed. A `.dot` output is useful when only the raw graph is available. -This is a graph of the DirectRunner-compatible pipeline construction, not a -runtime profile. +The checked-in SVG uses Beam's `RenderRunner`. Regenerate it with the +graph-only command below and a `.svg` path when Graphviz's `dot` is on `PATH`. -To render the graph without making Jev or BigQuery calls, use -`--graph-only`. The sink option is still required because it is part of the -constructed graph: +Use `--graph-only` to render before the pipeline starts. The sink remains part +of the graph, so the command includes a BigQuery dataset: ```sh python3 -m apache_beam.examples.inference.decision_models.message_router \ - --model=jev --graph-only \ + --model=jev \ + --graph-only \ --bq-dataset=PROJECT:DATASET \ - --graph=/tmp/beam-decision-model.dot + --graph=/tmp/beam-decision-model.svg ``` The `.dot` export works without Graphviz. SVG and PNG exports require `dot` on `PATH`. -For the Jev path, put the key in the environment and run the same short -stream: +Run the Jev path against the same short stream with `TYPESAFE_API_KEY` set: ```sh export TYPESAFE_API_KEY='replace-with-your-key' @@ -124,28 +117,11 @@ 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/pipeline.svg" + --graph="$jev_output_dir/message_router.svg" ``` -One three-message Beam DirectRunner smoke run against Jev measured the -following caller wall time with the local JSONL sink: - -| message | routed destination | caller latency | -| --- | --- | ---: | -| billing message | `billing` | 406.26 ms | -| technical message | `technical` | 249.70 ms | -| sales message | `sales` | 226.34 ms | - -These are observations from one local environment with `n=3` Choice calls and -one local file sink, not a latency guarantee or a benchmark. They include the -provider call inside the Beam caller and exclude BigQuery write latency. - -The key is read by the SDK through `TYPESAFE_API_KEY`; it is never a pipeline -element or a command-line argument. Do not put it in source, shell history, -logs, or a checked-in options file. - -To use a real streaming source, provide a Pub/Sub subscription containing JSON -objects with `event_id` and `message` string fields: +For a streaming source, provide a Pub/Sub subscription containing JSON objects +with `event_id` and `message` string fields: ```sh python3 -m apache_beam.examples.inference.decision_models.message_router \ @@ -154,68 +130,70 @@ python3 -m apache_beam.examples.inference.decision_models.message_router \ --bq-dataset=PROJECT:DATASET ``` -With `--bq-dataset=PROJECT:DATASET`, the destination is selected in code and -rows are appended to `messages_billing`, `messages_technical`, -`messages_sales`, or `messages_review` in that dataset. The process needs -Application Default Credentials with Pub/Sub subscriber access and permission -to create/write the BigQuery tables. BigQuery has not been live-verified by -this PoC. For a local run, omit `--bq-dataset` and use `--output-dir` instead. -The sink uses streaming inserts and Beam's normal at-least-once delivery -behavior. A production design needs an idempotency key and an explicit -retry/dead-letter policy. +With `--bq-dataset=PROJECT:DATASET`, rows go to `messages_billing`, +`messages_technical`, `messages_sales`, or `messages_review`. Application +Default Credentials need Pub/Sub subscriber access and permission to create and +write BigQuery tables. `WriteToBigQuery` uses streaming inserts and appends +rows. For deployment, add an idempotency key and explicit retry and dead-letter +policy. + +Use `--output-dir` instead of `--bq-dataset` for local JSONL files. A normal +run accepts exactly one sink option. ## Smaller demo: fraud review with Noul -Noul asks one yes/no question and returns the probability that the answer is -yes. The example uses that value as a deliberately visible policy threshold: +The fraud example prints one JSON row per sample message. Noul supplies the +probability of a yes answer, and values at or above `0.7` are flagged: ```sh python3 -m apache_beam.examples.inference.decision_models.fraud_review \ --model=local --primitive=noul +``` -export TYPESAFE_API_KEY='replace-with-your-key' +Add `--primitive=score` for the rubric score, or `--primitive=both` for both +questions. Run the Jev version with `--model=jev` after setting the key: + +```sh python3 -m apache_beam.examples.inference.decision_models.fraud_review \ - --model=jev --primitive=noul + --model=jev --primitive=both ``` -Rows include `fraud_review_probability` and `needs_review`. A probability near -0.5 is uncertainty, not confidence: Noul has no separate confidence value. -For comparison, run `--primitive=score` or `--primitive=both` to see a rubric -score, its probability distribution, and its confidence. Combining several -atomic questions in one request is the intended way to add policy dimensions -without asking one broad, ambiguous question. +A final three element DirectRunner run returned these values. Beam may print +the rows in a different order: -A three-element Beam DirectRunner Jev run with both questions produced Noul -probabilities `0.46` for the billing-address message, `0.32` for the -duplicate-charge message, and `0.97` for the request to bypass verification. -With the `0.7` threshold, only the last row was flagged. The Score answers -were `0.80`, `0.88`, and `3.00`, respectively. These are sample values. +| message | Noul probability | Score | +| --- | ---: | ---: | +| billing address | 0.45 | 0.86 | +| duplicate charge | 0.31 | 0.87 | +| bypass verification | 0.97 | 3.00 | -## Short benchmark +Only the bypass verification message crossed the `0.7` review threshold. -Each demo records request elapsed time in `latency_ms`. Run only the small -`TestStream` to measure a handful of calls and inspect the generated JSONL: +## Measured Jev performance -```sh -benchmark_dir="$(mktemp -d /tmp/beam-decision-model-benchmark.XXXXXX)" -time python3 -m apache_beam.examples.inference.decision_models.message_router \ - --model=jev \ - --output-dir="$benchmark_dir" -rg -I -o '"latency_ms": [0-9.]*' "$benchmark_dir" \ - | sort -t: -k2n -``` +These measurements were captured on September 22, 2026 with Jev 1.13.0, +TypeSafe SDK 0.7.1, and DirectRunner. The final trace used one SDK client, one +TCP connection, and one TLS handshake for the three events: + +| event | request latency | +| --- | ---: | +| billing | 439.72 ms | +| technical | 169.52 ms | +| sales | 292.01 ms | + +Total pipeline time was `1052.15 ms`, including pipeline construction. + +A separate transport experiment compared fresh clients with one pooled client: -This is an illustrative smoke benchmark, not a performance claim. It mixes -API/network latency, SDK work, and one local runner invocation, and three -messages are too few for a stable percentile. For a meaningful comparison, -repeat with a fixed corpus, record the model and request IDs, and report the -sample size and percentile method. Do not commit credentials or benchmark -output. +| 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 | -## Scope and next steps +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. -This PoC demonstrates the adapter boundary and one Beam integration point. It -does not provide multi-message batching, model warmup, circuit breaking, -schema evolution, or production deduplication. The Beam repeater is disabled; -the Jev SDK handles provider retries. A follow-up can add quota controls, -route metrics, and a runner-specific integration test before a production sink. +Each output row includes `latency_ms`. Provider adapters implement the +`DecisionModel` protocol with `__enter__`, `__exit__`, and `evaluate`. Laya and +Kev adapters can implement the same protocol. Beam owns the question flow, +policy gate, and destination sinks. 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 index cfc82c021591..a5d459b0bceb 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py +++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py @@ -20,17 +20,15 @@ import argparse import json import os -import time import apache_beam as beam -from apache_beam.io.requestresponse import Caller -from apache_beam.io.requestresponse import RequestResponseIO from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.examples.inference.decision_models.model import JevDecisionModel from apache_beam.examples.inference.decision_models.model import LocalDecisionModel from apache_beam.examples.inference.decision_models.model import NoulQuestion from apache_beam.examples.inference.decision_models.model import ScoreQuestion +from apache_beam.examples.inference.decision_models.transforms import EvaluateDecisions SAMPLE_MESSAGES = [ 'Please update my billing address before the next renewal.', @@ -49,43 +47,24 @@ 'An explicit attempt to bypass verification')) -class ReviewFraudCues(Caller): - def __init__(self, model, primitive): - self.model = model - self.primitive = primitive - - def __enter__(self): - self.model.__enter__() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - return self.model.__exit__(exc_type, exc_val, exc_tb) - - def __call__(self, message): - questions = {} - if self.primitive in ('noul', 'both'): - questions['fraud_cue'] = FRAUD_NOUL - if self.primitive in ('score', 'both'): - questions['risk_level'] = FRAUD_SCORE - started = time.perf_counter() - response = self.model.evaluate( - state={'message': message}, questions=questions) - row = { - 'message': message, - 'model': response.model, - 'provider': response.provider, - 'latency_ms': round((time.perf_counter() - started) * 1000, 2), - } - if 'fraud_cue' in response.answers: - probability = response.answers['fraud_cue'].noul - 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 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'].noul + 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 run(argv=None): @@ -97,12 +76,20 @@ def run(argv=None): 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: ( pipeline - | 'Sample payment messages' >> beam.Create(SAMPLE_MESSAGES) - | 'Evaluate fraud cues' >> RequestResponseIO( - ReviewFraudCues(model, args.primitive), timeout=15, repeater=None) + | 'Sample payment messages' >> beam.Create([{ + 'message': message + } for message in SAMPLE_MESSAGES]) + | 'Evaluate fraud cues' >> beam.ParDo( + EvaluateDecisions(model, questions)) + | 'Apply review threshold' >> beam.Map(review_fraud_cues) | 'Show review decisions' >> beam.Map(lambda row: print(json.dumps(row, sort_keys=True)))) 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 index cf341cf87254..a5bc2c4fb6a9 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/message_router.py +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py @@ -20,12 +20,9 @@ import argparse import json import os -import time import apache_beam as beam from apache_beam.io import fileio -from apache_beam.io.requestresponse import Caller -from apache_beam.io.requestresponse import RequestResponseIO from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.runners import render @@ -36,6 +33,7 @@ from apache_beam.examples.inference.decision_models.model import ChoiceQuestion from apache_beam.examples.inference.decision_models.model import JevDecisionModel from apache_beam.examples.inference.decision_models.model import LocalDecisionModel +from apache_beam.examples.inference.decision_models.transforms import EvaluateDecisions DESTINATIONS = { 'billing': 'Invoices, charges, payments, and refunds', @@ -70,43 +68,27 @@ 'provider:STRING,request_id:STRING,latency_ms:FLOAT') -class ClassifyMessage(Caller): - """Use any DecisionModel as a RequestResponseIO caller.""" - def __init__(self, model, min_confidence): - self.model = model - self.min_confidence = min_confidence - - def __enter__(self): - self.model.__enter__() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - return self.model.__exit__(exc_type, exc_val, exc_tb) - - def __call__(self, event): - started = time.perf_counter() - response = self.model.evaluate( - state={'message': event['message']}, - questions={'destination': QUESTION}) - latency_ms = (time.perf_counter() - started) * 1000 - 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 >= self.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(latency_ms, 2), - } +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): @@ -150,10 +132,9 @@ def build_pipeline(pipeline, args): model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel() rows = ( events - | 'Classify with decision model' >> RequestResponseIO( - ClassifyMessage(model, args.min_confidence), - timeout=15, - repeater=None) + | 'Ask Choice' >> beam.ParDo( + EvaluateDecisions(model, {'destination': QUESTION})) + | 'Select destination' >> beam.Map(route_message, args.min_confidence) | 'Window routed messages' >> beam.WindowInto(FixedWindows(60))) if args.bq_dataset: 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 index 27bc6575092f..34462df68677 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg @@ -20,78 +20,77 @@ - - + + G - - -cluster_ref_AppliedPTransform_Classify-with-decision-model_3 - - -Classify with decision model + + + +ref_AppliedPTransform_Sample-stream_2 + + +Sample stream - -cluster_ref_AppliedPTransform_Classify-with-decision-model-_Call_4 - - -_Call + + +ref_AppliedPTransform_Ask-Choice_3 + + +Ask Choice - - -ref_AppliedPTransform_Classify-with-decision-model-_Call-ParDo-_CallDoFn-_5 - - -ParDo(_CallDoFn) - - + + +ref_AppliedPTransform_Sample-stream_2->ref_AppliedPTransform_Ask-Choice_3 + + - + -ref_AppliedPTransform_Window-routed-messages_6 - - -Window routed messages +ref_AppliedPTransform_Select-destination_4 + + +Select destination - + -ref_AppliedPTransform_Classify-with-decision-model-_Call-ParDo-_CallDoFn-_5->ref_AppliedPTransform_Window-routed-messages_6 - - +ref_AppliedPTransform_Ask-Choice_3->ref_AppliedPTransform_Select-destination_4 + + - - -ref_AppliedPTransform_Sample-stream_2 - - -Sample stream + + +ref_AppliedPTransform_Window-routed-messages_5 + + +Window routed messages - - -ref_AppliedPTransform_Sample-stream_2->ref_AppliedPTransform_Classify-with-decision-model-_Call-ParDo-_CallDoFn-_5 - - + + +ref_AppliedPTransform_Select-destination_4->ref_AppliedPTransform_Window-routed-messages_5 + + - - -ref_AppliedPTransform_Write-routed-BigQuery-tables_7 - + + +ref_AppliedPTransform_Write-routed-BigQuery-tables_6 + Write routed BigQuery tables - - -ref_AppliedPTransform_Window-routed-messages_6->ref_AppliedPTransform_Write-routed-BigQuery-tables_7 + + +ref_AppliedPTransform_Window-routed-messages_5->ref_AppliedPTransform_Write-routed-BigQuery-tables_6 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 index 9f9b6311af49..a9d3ce2fb63d 100644 --- 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 @@ -19,50 +19,54 @@ import unittest -from apache_beam.examples.inference.decision_models.message_router import ClassifyMessage +from apache_beam.examples.inference.decision_models.message_router import route_message from apache_beam.examples.inference.decision_models.model import ChoiceAnswer from apache_beam.examples.inference.decision_models.model import DecisionResponse +from apache_beam.examples.inference.decision_models.transforms import DecisionResult -class FixedModel: - def __init__(self, answer): - self.answer = answer +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) - def evaluate(self, state, questions): - return DecisionResponse( - answers={'destination': self.answer}, model='fixed', provider='test') - -class ClassifyMessageTest(unittest.TestCase): +class RouteMessageTest(unittest.TestCase): def test_uncertain_choice_goes_to_review(self): - caller = ClassifyMessage( - FixedModel( - ChoiceAnswer( - choice='billing', - probabilities={ - 'billing': 0.6, 'sales': 0.4 - }, - confidence=0.4)), - min_confidence=0.65) - - row = caller({'event_id': '1', 'message': 'I have an account question'}) + 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): - caller = ClassifyMessage( - FixedModel(ChoiceAnswer(choice='other', confidence=1.0)), - min_confidence=0.5) + result = result_for(ChoiceAnswer(choice='other', confidence=1.0)) with self.assertRaisesRegex(ValueError, 'unknown destination'): - caller({'event_id': '1', 'message': 'Hello'}) + route_message(result, min_confidence=0.5) def test_missing_confidence_goes_to_review(self): - caller = ClassifyMessage( - FixedModel(ChoiceAnswer(choice='sales')), min_confidence=0.5) + result = result_for(ChoiceAnswer(choice='sales')) - row = caller({'event_id': '1', 'message': 'What does it cost?'}) + row = route_message(result, min_confidence=0.5) self.assertEqual('review', row['destination']) diff --git a/sdks/python/apache_beam/examples/inference/decision_models/model.py b/sdks/python/apache_beam/examples/inference/decision_models/model.py index 5808942a3259..6612e48494ee 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/model.py +++ b/sdks/python/apache_beam/examples/inference/decision_models/model.py @@ -90,9 +90,10 @@ def evaluate( class JevDecisionModel: - """TypeSafe SDK adapter. Each Beam caller creates one client per worker.""" - def __init__(self, model='jev-latest'): + """TypeSafe SDK adapter with a persistent HTTP connection pool.""" + def __init__(self, model='jev-latest', timeout=10.0): self.model = model + self.timeout = timeout self.client = None def __enter__(self): @@ -101,7 +102,7 @@ def __enter__(self): except ImportError as error: raise ImportError( 'Install typesafe-sdk to use JevDecisionModel') from error - self.client = TypeSafeClient() + self.client = TypeSafeClient(timeout=self.timeout) return self def __exit__(self, exc_type, exc_value, traceback): diff --git a/sdks/python/apache_beam/examples/inference/decision_models/transforms.py b/sdks/python/apache_beam/examples/inference/decision_models/transforms.py new file mode 100644 index 000000000000..28805052bdb8 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/transforms.py @@ -0,0 +1,66 @@ +# +# 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. +# + +"""Evaluate typed questions with a model connection owned by each DoFn instance.""" + +import time +from dataclasses import dataclass +from typing import Any + +import apache_beam as beam +from apache_beam.metrics import Metrics + +from apache_beam.examples.inference.decision_models.model import DecisionResponse + + +@dataclass(frozen=True) +class DecisionResult: + state: Any + response: DecisionResponse + latency_ms: float + + +class EvaluateDecisions(beam.DoFn): + """Reuse a model client across elements and bundles on this DoFn instance. + + The adapter controls network timeouts and retries. Keeping the call on the + processing thread avoids a per-element executor and preserves backpressure. + """ + def __init__(self, model, questions): + self.model = model + self.questions = questions + self.requests = Metrics.counter(self.__class__, 'requests') + self.failures = Metrics.counter(self.__class__, 'failures') + self.latency = Metrics.distribution(self.__class__, 'request_latency_ms') + + def setup(self): + self.model.__enter__() + + def process(self, state): + self.requests.inc() + started = time.perf_counter() + try: + response = self.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): + self.model.__exit__(None, None, None) diff --git a/sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py b/sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py new file mode 100644 index 000000000000..747721080b64 --- /dev/null +++ b/sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py @@ -0,0 +1,62 @@ +# +# 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 decision model lifecycle and evaluation failures.""" + +import unittest +from unittest.mock import Mock + +from apache_beam.examples.inference.decision_models.model import LocalDecisionModel +from apache_beam.examples.inference.decision_models.model import NoulQuestion +from apache_beam.examples.inference.decision_models.transforms import EvaluateDecisions + + +class EvaluateDecisionsTest(unittest.TestCase): + def test_reuses_model_across_bundles_and_preserves_state(self): + model = Mock(wraps=LocalDecisionModel()) + model.__enter__ = Mock() + model.__exit__ = Mock() + question = NoulQuestion('Does the message ask to bypass verification?') + transform = EvaluateDecisions(model, {'fraud': question}) + transform.setup() + try: + for message, expected in [('Pay my invoice', 0.0), + ('Please bypass verification', 1.0)]: + state = {'message': message, 'event_id': message} + transform.start_bundle() + result, = transform.process(state) + transform.finish_bundle() + self.assertEqual(state, result.state) + self.assertEqual(expected, result.response.answers['fraud'].noul) + self.assertGreaterEqual(result.latency_ms, 0) + model.__enter__.assert_called_once_with() + self.assertEqual(2, model.evaluate.call_count) + finally: + transform.teardown() + model.__exit__.assert_called_once_with(None, None, None) + + def test_evaluation_failure_propagates(self): + model = Mock() + model.evaluate.side_effect = TimeoutError('decision timed out') + transform = EvaluateDecisions(model, {}) + + with self.assertRaisesRegex(TimeoutError, 'decision timed out'): + list(transform.process({'message': 'Hello'})) + + +if __name__ == '__main__': + unittest.main() From d04ce54454cb5d47000751ced5c5d7029e0c7fe2 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Tue, 22 Sep 2026 13:16:55 -0400 Subject: [PATCH 3/7] [Python] Add reusable decision model primitives and adapters --- .../apache_beam/examples/inference/README.md | 8 +- .../inference/decision_models/README.md | 181 ++++++++------ .../inference/decision_models/fraud_review.py | 18 +- .../inference/decision_models/local_model.py | 78 ++++++ .../decision_models/message_router.py | 13 +- .../decision_models/message_router.svg | 30 +-- .../decision_models/message_router_test.py | 6 +- .../inference/decision_models/model.py | 214 ---------------- .../decision_models/requirements.txt | 4 +- .../inference/decision_models/transforms.py | 66 ----- .../decision_models/transforms_test.py | 62 ----- .../apache_beam/ml/inference/decision.py | 232 ++++++++++++++++++ .../apache_beam/ml/inference/decision_test.py | 213 ++++++++++++++++ .../ml/inference/typesafe_inference.py | 107 ++++++++ .../ml/inference/typesafe_inference_test.py | 106 ++++++++ 15 files changed, 883 insertions(+), 455 deletions(-) create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/local_model.py delete mode 100644 sdks/python/apache_beam/examples/inference/decision_models/model.py delete mode 100644 sdks/python/apache_beam/examples/inference/decision_models/transforms.py delete mode 100644 sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py create mode 100644 sdks/python/apache_beam/ml/inference/decision.py create mode 100644 sdks/python/apache_beam/ml/inference/decision_test.py create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_inference.py create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_inference_test.py diff --git a/sdks/python/apache_beam/examples/inference/README.md b/sdks/python/apache_beam/examples/inference/README.md index c43ad4583479..f14b531c21a2 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. diff --git a/sdks/python/apache_beam/examples/inference/decision_models/README.md b/sdks/python/apache_beam/examples/inference/decision_models/README.md index ebc83f2c9c7b..3750b491fa78 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/README.md +++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md @@ -17,47 +17,101 @@ under the License. --> -# Beam decision model examples +# Beam decision models -These examples connect Beam pipelines to typed decision models. The main -pipeline uses Choice to route support messages. The fraud pipeline uses Noul -and optionally Score. `LocalDecisionModel` runs the graph without a network -call, while `JevDecisionModel` adapts the TypeSafe SDK. +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. ## Architecture ```mermaid flowchart LR - source["TestStream
or Pub/Sub"] --> decode["Decode JSON event"] - decode --> evaluate["EvaluateDecisions
(beam.DoFn)"] - evaluate --> result["DecisionResult"] - result --> route["route_message
validate Choice + confidence"] - route --> sink["JSONL files
or dynamic BigQuery tables"] - route --> review["review"] - review --> sink + input["PCollection[state]"] --> evaluate["EvaluateDecisions(model, questions)
PTransform"] + model["DecisionModel adapter"] --> evaluate + evaluate --> result["PCollection[DecisionResult]"] + result --> policy["route_message
policy and sink selection"] + policy --> sink["JSONL files
or BigQuery tables"] ``` -`EvaluateDecisions` accepts a model and typed questions. `setup` enters the -model context and `teardown` closes it. Each DoFn instance owns one context. -For Jev, that context creates one TypeSafe SDK client with a persistent HTTP -connection pool used across elements and bundles. `process` calls `evaluate` -synchronously and yields `DecisionResult(state, response, latency_ms)`, -preserving Beam backpressure without a per-element executor. +![Rendered Beam message router graph](message_router.svg) -`JevDecisionModel` uses a 10 second HTTP operation timeout and the SDK's retry -behavior. The SDK keepalive default is 5 seconds, and connection reuse follows -the SDK pool lifecycle. +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. -`route_message` checks that the Choice answer is `billing`, `technical`, or -`sales`. A missing or low confidence answer goes to `review`. The selected -destination controls the JSONL file name or BigQuery table name. +## Integration -![Rendered Beam message router graph](message_router.svg) +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 existing example requirements and set +`TYPESAFE_API_KEY` when running it. ## Install -From a Beam checkout, activate a Python virtual environment and install the -core SDK: +From a Beam checkout, activate a virtual environment and install the core SDK: ```sh python3 -m pip install -e sdks/python @@ -70,17 +124,16 @@ python3 -m pip install -r \ sdks/python/apache_beam/examples/inference/decision_models/requirements.txt ``` -Install Beam's GCP extra for Pub/Sub or BigQuery, including graph-only renders -that contain a BigQuery sink: +Install Beam's GCP extra for Pub/Sub or BigQuery sinks: ```sh python3 -m pip install -e 'sdks/python[gcp]' ``` -## Main demo: route messages with Choice +## Run examples -The default source is a three element `TestStream`. The local adapter writes -one JSONL file per destination: +The main demo uses a three-element `TestStream` and writes one JSONL file per +destination: ```sh output_dir="$(mktemp -d /tmp/beam-decision-model-output.XXXXXX)" @@ -91,11 +144,7 @@ python3 -m apache_beam.examples.inference.decision_models.message_router \ find "$output_dir" -name '*.jsonl' -print -exec sed -n '1,3p' {} \; ``` -The checked-in SVG uses Beam's `RenderRunner`. Regenerate it with the -graph-only command below and a `.svg` path when Graphviz's `dot` is on `PATH`. - -Use `--graph-only` to render before the pipeline starts. The sink remains part -of the graph, so the command includes a BigQuery dataset: +Render the graph without running the job. The sink stays in the graph: ```sh python3 -m apache_beam.examples.inference.decision_models.message_router \ @@ -105,10 +154,9 @@ python3 -m apache_beam.examples.inference.decision_models.message_router \ --graph=/tmp/beam-decision-model.svg ``` -The `.dot` export works without Graphviz. SVG and PNG exports require `dot` on -`PATH`. +The `.dot` export works without Graphviz; SVG and PNG require `dot` on `PATH`. -Run the Jev path against the same short stream with `TYPESAFE_API_KEY` set: +Run the Jev path against the same stream with `TYPESAFE_API_KEY` set: ```sh export TYPESAFE_API_KEY='replace-with-your-key' @@ -121,7 +169,7 @@ python3 -m apache_beam.examples.inference.decision_models.message_router \ ``` For a streaming source, provide a Pub/Sub subscription containing JSON objects -with `event_id` and `message` string fields: +with `event_id` and `message` fields: ```sh python3 -m apache_beam.examples.inference.decision_models.message_router \ @@ -131,59 +179,49 @@ python3 -m apache_beam.examples.inference.decision_models.message_router \ ``` With `--bq-dataset=PROJECT:DATASET`, rows go to `messages_billing`, -`messages_technical`, `messages_sales`, or `messages_review`. Application -Default Credentials need Pub/Sub subscriber access and permission to create and -write BigQuery tables. `WriteToBigQuery` uses streaming inserts and appends -rows. For deployment, add an idempotency key and explicit retry and dead-letter -policy. - -Use `--output-dir` instead of `--bq-dataset` for local JSONL files. A normal -run accepts exactly one sink option. - -## Smaller demo: fraud review with Noul +`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 -probability of a yes answer, and values at or above `0.7` are flagged: +boolean probability; values at or above `0.7` are flagged: ```sh python3 -m apache_beam.examples.inference.decision_models.fraud_review \ --model=local --primitive=noul -``` - -Add `--primitive=score` for the rubric score, or `--primitive=both` for both -questions. Run the Jev version with `--model=jev` after setting the key: - -```sh +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 ``` -A final three element DirectRunner run returned these values. Beam may print +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.45 | 0.86 | -| duplicate charge | 0.31 | 0.87 | +| 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. -## Measured Jev performance +## 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 one SDK client, one -TCP connection, and one TLS handshake for the three events: +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 | 439.72 ms | -| technical | 169.52 ms | -| sales | 292.01 ms | +| billing | 408.27 ms | +| technical | 217.50 ms | +| sales | 190.43 ms | + +Total pipeline time was `986.68 ms`, including pipeline construction. -Total pipeline time was `1052.15 ms`, including pipeline construction. +### Standalone transport experiment, September 22, 2026 -A separate transport experiment compared fresh clients with one pooled client: +An earlier standalone experiment compared fresh clients with one pooled client: | client setup | observations | | --- | --- | @@ -192,8 +230,3 @@ A separate transport experiment compared fresh clients with one pooled client: 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. - -Each output row includes `latency_ms`. Provider adapters implement the -`DecisionModel` protocol with `__enter__`, `__exit__`, and `evaluate`. Laya and -Kev adapters can implement the same protocol. Beam owns the question flow, -policy gate, and destination sinks. 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 index a5d459b0bceb..3368cb641a7b 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py +++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py @@ -22,21 +22,20 @@ 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 -from apache_beam.examples.inference.decision_models.model import JevDecisionModel -from apache_beam.examples.inference.decision_models.model import LocalDecisionModel -from apache_beam.examples.inference.decision_models.model import NoulQuestion -from apache_beam.examples.inference.decision_models.model import ScoreQuestion -from apache_beam.examples.inference.decision_models.transforms import EvaluateDecisions - 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 = NoulQuestion( +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?', @@ -56,7 +55,7 @@ def review_fraud_cues(result): 'latency_ms': round(result.latency_ms, 2), } if 'fraud_cue' in response.answers: - probability = response.answers['fraud_cue'].noul + probability = response.answers['fraud_cue'].probability row['fraud_review_probability'] = probability row['needs_review'] = probability >= 0.7 if 'risk_level' in response.answers: @@ -87,8 +86,7 @@ def run(argv=None): | 'Sample payment messages' >> beam.Create([{ 'message': message } for message in SAMPLE_MESSAGES]) - | 'Evaluate fraud cues' >> beam.ParDo( - EvaluateDecisions(model, questions)) + | 'Evaluate fraud cues' >> EvaluateDecisions(model, questions) | 'Apply review threshold' >> beam.Map(review_fraud_cues) | 'Show review decisions' >> beam.Map(lambda row: print(json.dumps(row, sort_keys=True)))) 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 000000000000..55c30216f2c2 --- /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 index a5bc2c4fb6a9..8413d1e0b0ef 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/message_router.py +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.py @@ -22,7 +22,11 @@ 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 @@ -30,11 +34,6 @@ from apache_beam.transforms.window import FixedWindows from apache_beam.transforms.window import TimestampedValue -from apache_beam.examples.inference.decision_models.model import ChoiceQuestion -from apache_beam.examples.inference.decision_models.model import JevDecisionModel -from apache_beam.examples.inference.decision_models.model import LocalDecisionModel -from apache_beam.examples.inference.decision_models.transforms import EvaluateDecisions - DESTINATIONS = { 'billing': 'Invoices, charges, payments, and refunds', 'technical': 'Bugs, outages, and integration problems', @@ -132,8 +131,7 @@ def build_pipeline(pipeline, args): model = JevDecisionModel() if args.model == 'jev' else LocalDecisionModel() rows = ( events - | 'Ask Choice' >> beam.ParDo( - EvaluateDecisions(model, {'destination': QUESTION})) + | 'Ask Choice' >> EvaluateDecisions(model, {'destination': QUESTION}) | 'Select destination' >> beam.Map(route_message, args.min_confidence) | 'Window routed messages' >> beam.WindowInto(FixedWindows(60))) @@ -186,6 +184,7 @@ def run(argv=None): 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) 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 index 34462df68677..e37964497cf8 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg +++ b/sdks/python/apache_beam/examples/inference/decision_models/message_router.svg @@ -49,48 +49,48 @@ - + -ref_AppliedPTransform_Select-destination_4 - +ref_AppliedPTransform_Select-destination_5 + Select destination - + -ref_AppliedPTransform_Ask-Choice_3->ref_AppliedPTransform_Select-destination_4 +ref_AppliedPTransform_Ask-Choice_3->ref_AppliedPTransform_Select-destination_5 - + -ref_AppliedPTransform_Window-routed-messages_5 - +ref_AppliedPTransform_Window-routed-messages_6 + Window routed messages - + -ref_AppliedPTransform_Select-destination_4->ref_AppliedPTransform_Window-routed-messages_5 +ref_AppliedPTransform_Select-destination_5->ref_AppliedPTransform_Window-routed-messages_6 - + -ref_AppliedPTransform_Write-routed-BigQuery-tables_6 - +ref_AppliedPTransform_Write-routed-BigQuery-tables_7 + Write routed BigQuery tables - + -ref_AppliedPTransform_Window-routed-messages_5->ref_AppliedPTransform_Write-routed-BigQuery-tables_6 +ref_AppliedPTransform_Window-routed-messages_6->ref_AppliedPTransform_Write-routed-BigQuery-tables_7 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 index a9d3ce2fb63d..c5040599b081 100644 --- 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 @@ -20,9 +20,9 @@ import unittest from apache_beam.examples.inference.decision_models.message_router import route_message -from apache_beam.examples.inference.decision_models.model import ChoiceAnswer -from apache_beam.examples.inference.decision_models.model import DecisionResponse -from apache_beam.examples.inference.decision_models.transforms import DecisionResult +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): diff --git a/sdks/python/apache_beam/examples/inference/decision_models/model.py b/sdks/python/apache_beam/examples/inference/decision_models/model.py deleted file mode 100644 index 6612e48494ee..000000000000 --- a/sdks/python/apache_beam/examples/inference/decision_models/model.py +++ /dev/null @@ -1,214 +0,0 @@ -# -# 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 decision questions and a small Jev adapter for Beam examples.""" - -from dataclasses import dataclass -from dataclasses import field -from typing import Any -from typing import Mapping -from typing import Protocol -from typing import Sequence -from typing import Union - - -@dataclass(frozen=True) -class ChoiceQuestion: - instructions: str - criteria: Mapping[str, str] - - -@dataclass(frozen=True) -class NoulQuestion: - instructions: str - - -@dataclass(frozen=True) -class ScoreQuestion: - instructions: str - criteria: Sequence[str] - - -Question = Union[ChoiceQuestion, NoulQuestion, ScoreQuestion] - - -@dataclass(frozen=True) -class ChoiceAnswer: - choice: str - probabilities: Mapping[str, float] = field(default_factory=dict) - confidence: float | None = None - - -@dataclass(frozen=True) -class NoulAnswer: - noul: float - - -@dataclass(frozen=True) -class ScoreAnswer: - score: float - probabilities: Mapping[str, float] = field(default_factory=dict) - confidence: float = 0.0 - - -Answer = Union[ChoiceAnswer, NoulAnswer, ScoreAnswer] - - -@dataclass(frozen=True) -class DecisionResponse: - answers: Mapping[str, Answer] - model: str - provider: str - request_id: str | None = None - - -class DecisionModel(Protocol): - """A provider owns evaluation; Beam owns routing and persistence.""" - def __enter__(self) -> 'DecisionModel': - ... - - def __exit__(self, exc_type, exc_value, traceback) -> None: - ... - - def evaluate( - self, state: Any, questions: Mapping[str, Question]) -> DecisionResponse: - ... - - -class JevDecisionModel: - """TypeSafe SDK adapter with a persistent HTTP connection pool.""" - 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') - from typesafe_sdk import Choice - from typesafe_sdk import Noul - from typesafe_sdk import Score - - sdk_questions = {} - for name, question in questions.items(): - if isinstance(question, ChoiceQuestion): - sdk_questions[name] = Choice( - instructions=question.instructions, criteria=question.criteria) - elif isinstance(question, NoulQuestion): - sdk_questions[name] = Noul(instructions=question.instructions) - elif isinstance(question, ScoreQuestion): - sdk_questions[name] = Score( - instructions=question.instructions, criteria=question.criteria) - else: - raise TypeError('Unsupported decision question: %r' % type(question)) - - response = self.client.system_one( - state=state, questions=sdk_questions, model=self.model) - answers = {} - for name, question in questions.items(): - answer = response.answers[name] - if isinstance(question, ChoiceQuestion): - answers[name] = ChoiceAnswer( - choice=answer.choice, - probabilities=dict(answer.probabilities), - confidence=answer.confidence) - elif isinstance(question, NoulQuestion): - answers[name] = NoulAnswer(noul=answer.noul) - else: - answers[name] = ScoreAnswer( - score=answer.score, - probabilities={ - str(level): probability - for level, probability in answer.probabilities.items() - }, - confidence=answer.confidence) - return DecisionResponse( - answers=answers, - model=response.model, - provider='typesafe', - request_id=response.request_id) - - -_ROUTE_TERMS = { - 'billing': ('invoice', 'charge', 'refund', 'payment', 'subscription'), - 'technical': ('api', '503', 'error', 'integration', 'outage'), - 'sales': ('quote', 'seat', 'pricing', 'price', 'plan'), -} - - -class LocalDecisionModel: - """Transparent rules for running these examples without network access.""" - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - return None - - 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, NoulQuestion): - answers[name] = NoulAnswer(noul=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={ - str(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/requirements.txt b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt index d7280f85b93e..91f59eb9f7f2 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt +++ b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt @@ -15,7 +15,5 @@ # limitations under the License. # -# Optional: required only for --model=jev. The local adapter has no extra -# dependency. Keep this compatible range aligned with the Jev SDK API used by -# model.py. +# Optional client for apache_beam.ml.inference.typesafe_inference. typesafe-sdk>=0.7.1,<0.8 diff --git a/sdks/python/apache_beam/examples/inference/decision_models/transforms.py b/sdks/python/apache_beam/examples/inference/decision_models/transforms.py deleted file mode 100644 index 28805052bdb8..000000000000 --- a/sdks/python/apache_beam/examples/inference/decision_models/transforms.py +++ /dev/null @@ -1,66 +0,0 @@ -# -# 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. -# - -"""Evaluate typed questions with a model connection owned by each DoFn instance.""" - -import time -from dataclasses import dataclass -from typing import Any - -import apache_beam as beam -from apache_beam.metrics import Metrics - -from apache_beam.examples.inference.decision_models.model import DecisionResponse - - -@dataclass(frozen=True) -class DecisionResult: - state: Any - response: DecisionResponse - latency_ms: float - - -class EvaluateDecisions(beam.DoFn): - """Reuse a model client across elements and bundles on this DoFn instance. - - The adapter controls network timeouts and retries. Keeping the call on the - processing thread avoids a per-element executor and preserves backpressure. - """ - def __init__(self, model, questions): - self.model = model - self.questions = questions - self.requests = Metrics.counter(self.__class__, 'requests') - self.failures = Metrics.counter(self.__class__, 'failures') - self.latency = Metrics.distribution(self.__class__, 'request_latency_ms') - - def setup(self): - self.model.__enter__() - - def process(self, state): - self.requests.inc() - started = time.perf_counter() - try: - response = self.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): - self.model.__exit__(None, None, None) diff --git a/sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py b/sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py deleted file mode 100644 index 747721080b64..000000000000 --- a/sdks/python/apache_beam/examples/inference/decision_models/transforms_test.py +++ /dev/null @@ -1,62 +0,0 @@ -# -# 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 decision model lifecycle and evaluation failures.""" - -import unittest -from unittest.mock import Mock - -from apache_beam.examples.inference.decision_models.model import LocalDecisionModel -from apache_beam.examples.inference.decision_models.model import NoulQuestion -from apache_beam.examples.inference.decision_models.transforms import EvaluateDecisions - - -class EvaluateDecisionsTest(unittest.TestCase): - def test_reuses_model_across_bundles_and_preserves_state(self): - model = Mock(wraps=LocalDecisionModel()) - model.__enter__ = Mock() - model.__exit__ = Mock() - question = NoulQuestion('Does the message ask to bypass verification?') - transform = EvaluateDecisions(model, {'fraud': question}) - transform.setup() - try: - for message, expected in [('Pay my invoice', 0.0), - ('Please bypass verification', 1.0)]: - state = {'message': message, 'event_id': message} - transform.start_bundle() - result, = transform.process(state) - transform.finish_bundle() - self.assertEqual(state, result.state) - self.assertEqual(expected, result.response.answers['fraud'].noul) - self.assertGreaterEqual(result.latency_ms, 0) - model.__enter__.assert_called_once_with() - self.assertEqual(2, model.evaluate.call_count) - finally: - transform.teardown() - model.__exit__.assert_called_once_with(None, None, None) - - def test_evaluation_failure_propagates(self): - model = Mock() - model.evaluate.side_effect = TimeoutError('decision timed out') - transform = EvaluateDecisions(model, {}) - - with self.assertRaisesRegex(TimeoutError, 'decision timed out'): - list(transform.process({'message': 'Hello'})) - - -if __name__ == '__main__': - unittest.main() 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 000000000000..dcf476b42edb --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/decision.py @@ -0,0 +1,232 @@ +# +# 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) + + def expand(self, states): + return states | beam.ParDo( + _EvaluateDecisionsDoFn(self._model, self._questions)) + + +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 000000000000..ad4d604743ec --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/decision_test.py @@ -0,0 +1,213 @@ +# +# 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 _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') + + +class DecisionTest(unittest.TestCase): + 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 000000000000..6188f9713f1c --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/typesafe_inference.py @@ -0,0 +1,107 @@ +# +# 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 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'] + + +class JevDecisionModel(DecisionModel): + """Evaluate Beam decision questions with Jev through the TypeSafe SDK. + + Install ``typesafe-sdk>=0.7.1,<0.8`` on the workers and set + ``TYPESAFE_API_KEY`` in their environment. 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`. + + 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') + from typesafe_sdk import Choice + from typesafe_sdk import Noul + from typesafe_sdk import Score + + sdk_questions = {} + for name, question in questions.items(): + if isinstance(question, ChoiceQuestion): + sdk_questions[name] = Choice( + instructions=question.instructions, criteria=question.criteria) + elif isinstance(question, BooleanQuestion): + sdk_questions[name] = Noul(instructions=question.instructions) + elif isinstance(question, ScoreQuestion): + sdk_questions[name] = Score( + instructions=question.instructions, criteria=question.criteria) + else: + raise TypeError('Unsupported decision question: %r' % type(question)) + + response = self.client.system_one( + state=state, questions=sdk_questions, model=self.model) + answers = {} + for name, question in questions.items(): + answer = response.answers[name] + if isinstance(question, ChoiceQuestion): + answers[name] = ChoiceAnswer( + choice=answer.choice, + probabilities=dict(answer.probabilities), + confidence=answer.confidence) + elif isinstance(question, BooleanQuestion): + answers[name] = BooleanAnswer(probability=answer.noul) + else: + answers[name] = ScoreAnswer( + score=answer.score, + probabilities={ + int(level): probability + for level, probability in answer.probabilities.items() + }, + confidence=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_test.py b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py new file mode 100644 index 000000000000..2711029e4b17 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py @@ -0,0 +1,106 @@ +# +# 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 + +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 + + +class JevDecisionModelTest(unittest.TestCase): + def test_normalizes_all_primitives_and_reuses_client(self): + client = mock.Mock() + client.system_one.return_value = SimpleNamespace( + answers={ + 'team': SimpleNamespace( + choice='billing', + probabilities={'billing': 1.0}, + confidence=0.9), + 'review': SimpleNamespace(noul=0.8), + 'urgency': SimpleNamespace( + score=0.7, probabilities={ + '0': 0.3, '1': 0.7 + }, confidence=0.6), + }, + 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) + + 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() + + +if __name__ == '__main__': + unittest.main() From 81a3c48770b828c516ff5ee3824629c7f414aae4 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Tue, 22 Sep 2026 19:15:39 -0400 Subject: [PATCH 4/7] [Python] Gate live Jev tests and isolate runtime value options --- .../inference/decision_models/README.md | 30 ++++++ .../examples/snippets/snippets_test.py | 8 +- .../inference/typesafe_inference_it_test.py | 93 +++++++++++++++++++ 3 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py diff --git a/sdks/python/apache_beam/examples/inference/decision_models/README.md b/sdks/python/apache_beam/examples/inference/decision_models/README.md index 3750b491fa78..9975a0d2735b 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/README.md +++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md @@ -205,6 +205,36 @@ the rows in a different order: Only the bypass verification message crossed the `0.7` review threshold. +## Tests + +Install Beam's test dependencies in the same virtual environment: + +```sh +python3 -m pip install -e 'sdks/python[test]' +``` + +The unit tests cover the shared API, routing policy, and mocked Jev responses: + +```sh +python3 -m pytest -q \ + sdks/python/apache_beam/ml/inference/decision_test.py \ + sdks/python/apache_beam/ml/inference/typesafe_inference_test.py \ + sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py +``` + +With the optional Jev SDK installed and `TYPESAFE_API_KEY` set, run the live +integration test: + +```sh +python3 -m pytest -q -rs -m it_postcommit \ + sdks/python/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, diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py b/sdks/python/apache_beam/examples/snippets/snippets_test.py index 5a47d9289012..8898cc8555f9 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.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 @@ def check_enriched_elements(actual): 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/typesafe_inference_it_test.py b/sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py new file mode 100644 index 000000000000..acd2234f2ddf --- /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 + + +@pytest.mark.it_postcommit +@unittest.skipIf(typesafe_sdk is None, 'typesafe-sdk is not installed') +@unittest.skipUnless( + 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() From c8b3a97d54c0c80b99995c12f38e77edf5c2cb5d Mon Sep 17 00:00:00 2001 From: bvolpato Date: Sat, 10 Oct 2026 18:21:42 -0400 Subject: [PATCH 5/7] [Python] Validate decision responses and add dead-letter handling --- .../inference/decision_models/README.md | 62 +++++-- .../inference/decision_models/fraud_review.py | 44 ++++- .../decision_models/fraud_review_test.py | 86 +++++++++ .../decision_models/requirements.txt | 2 +- .../apache_beam/ml/inference/decision.py | 18 +- .../apache_beam/ml/inference/decision_test.py | 41 +++++ .../ml/inference/typesafe_inference.py | 95 ++++++++-- .../ml/inference/typesafe_inference_test.py | 168 +++++++++++++++++- .../inference/typesafe_tests_requirements.txt | 18 ++ 9 files changed, 486 insertions(+), 48 deletions(-) create mode 100644 sdks/python/apache_beam/examples/inference/decision_models/fraud_review_test.py create mode 100644 sdks/python/apache_beam/ml/inference/typesafe_tests_requirements.txt diff --git a/sdks/python/apache_beam/examples/inference/decision_models/README.md b/sdks/python/apache_beam/examples/inference/decision_models/README.md index 9975a0d2735b..637b8159c555 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/README.md +++ b/sdks/python/apache_beam/examples/inference/decision_models/README.md @@ -52,6 +52,12 @@ during DoFn setup and closes it during teardown for each DoFn instance. 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 @@ -106,7 +112,7 @@ with beam.Pipeline() as pipeline: ``` The downstream policy stays unchanged when swapping adapters. The Jev client -is optional; install the existing example requirements and set +is optional; install the shared SDK requirements listed below and set `TYPESAFE_API_KEY` when running it. ## Install @@ -114,28 +120,34 @@ is optional; install the existing example requirements and set From a Beam checkout, activate a virtual environment and install the core SDK: ```sh -python3 -m pip install -e sdks/python +cd sdks/python +python3 -m pip install -e . ``` -Install the optional Jev client for `--model=jev`: +Install the optional Jev SDK used by the adapter and examples: ```sh +cd sdks/python python3 -m pip install -r \ - sdks/python/apache_beam/examples/inference/decision_models/requirements.txt + apache_beam/ml/inference/typesafe_tests_requirements.txt ``` Install Beam's GCP extra for Pub/Sub or BigQuery sinks: ```sh -python3 -m pip install -e 'sdks/python[gcp]' +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 \ @@ -147,6 +159,7 @@ 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 \ @@ -159,6 +172,7 @@ 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 \ @@ -172,6 +186,7 @@ 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 \ @@ -186,6 +201,7 @@ 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 \ @@ -194,6 +210,28 @@ 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: @@ -210,24 +248,28 @@ Only the bypass verification message crossed the `0.7` review threshold. Install Beam's test dependencies in the same virtual environment: ```sh -python3 -m pip install -e 'sdks/python[test]' +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 \ - sdks/python/apache_beam/ml/inference/decision_test.py \ - sdks/python/apache_beam/ml/inference/typesafe_inference_test.py \ - sdks/python/apache_beam/examples/inference/decision_models/message_router_test.py + 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 \ - sdks/python/apache_beam/ml/inference/typesafe_inference_it_test.py + apache_beam/ml/inference/typesafe_inference_it_test.py ``` It sends one streaming event through `EvaluateDecisions`, asking Choice, 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 index 3368cb641a7b..5afcb6f34f2d 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py +++ b/sdks/python/apache_beam/examples/inference/decision_models/fraud_review.py @@ -22,7 +22,8 @@ import os import apache_beam as beam -from apache_beam.examples.inference.decision_models.local_model import LocalDecisionModel +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 @@ -32,7 +33,8 @@ 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.', + 'Transfer the refund to this new account immediately and ' + 'skip verification.', ] FRAUD_NOUL = BooleanQuestion( @@ -66,11 +68,23 @@ def review_fraud_cues(result): 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') @@ -81,12 +95,28 @@ def run(argv=None): 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 ( - pipeline - | 'Sample payment messages' >> beam.Create([{ - 'message': message - } for message in SAMPLE_MESSAGES]) - | 'Evaluate fraud cues' >> EvaluateDecisions(model, questions) + decisions | 'Apply review threshold' >> beam.Map(review_fraud_cues) | 'Show review decisions' >> beam.Map(lambda row: print(json.dumps(row, sort_keys=True)))) 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 000000000000..ec6f8160ea6e --- /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/requirements.txt b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt index 91f59eb9f7f2..6b99d3ce92cd 100644 --- a/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt +++ b/sdks/python/apache_beam/examples/inference/decision_models/requirements.txt @@ -16,4 +16,4 @@ # # Optional client for apache_beam.ml.inference.typesafe_inference. -typesafe-sdk>=0.7.1,<0.8 +-r ../../../ml/inference/typesafe_tests_requirements.txt diff --git a/sdks/python/apache_beam/ml/inference/decision.py b/sdks/python/apache_beam/ml/inference/decision.py index dcf476b42edb..fdfdc5801dc4 100644 --- a/sdks/python/apache_beam/ml/inference/decision.py +++ b/sdks/python/apache_beam/ml/inference/decision.py @@ -195,10 +195,24 @@ 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): - return states | beam.ParDo( - _EvaluateDecisionsDoFn(self._model, self._questions)) + 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): diff --git a/sdks/python/apache_beam/ml/inference/decision_test.py b/sdks/python/apache_beam/ml/inference/decision_test.py index ad4d604743ec..f596acd6e3ec 100644 --- a/sdks/python/apache_beam/ml/inference/decision_test.py +++ b/sdks/python/apache_beam/ml/inference/decision_test.py @@ -82,6 +82,13 @@ def evaluate(self, state, questions): '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, @@ -115,7 +122,41 @@ 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} diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference.py b/sdks/python/apache_beam/ml/inference/typesafe_inference.py index 6188f9713f1c..130a0cafb5e9 100644 --- a/sdks/python/apache_beam/ml/inference/typesafe_inference.py +++ b/sdks/python/apache_beam/ml/inference/typesafe_inference.py @@ -17,6 +17,8 @@ """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 @@ -29,15 +31,54 @@ __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 ``typesafe-sdk>=0.7.1,<0.8`` on the workers and set - ``TYPESAFE_API_KEY`` in their environment. The client is created when Beam + 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. @@ -63,43 +104,59 @@ def __exit__(self, exc_type, exc_value, traceback): def evaluate(self, state, questions): if self.client is None: raise RuntimeError('Enter JevDecisionModel before evaluating questions') - from typesafe_sdk import Choice - from typesafe_sdk import Noul - from typesafe_sdk import Score + import typesafe_sdk sdk_questions = {} for name, question in questions.items(): if isinstance(question, ChoiceQuestion): - sdk_questions[name] = Choice( + sdk_questions[name] = typesafe_sdk.Choice( instructions=question.instructions, criteria=question.criteria) elif isinstance(question, BooleanQuestion): - sdk_questions[name] = Noul(instructions=question.instructions) + sdk_questions[name] = typesafe_sdk.Noul( + instructions=question.instructions) elif isinstance(question, ScoreQuestion): - sdk_questions[name] = Score( + sdk_questions[name] = typesafe_sdk.Score( instructions=question.instructions, criteria=question.criteria) else: raise TypeError('Unsupported decision question: %r' % type(question)) - response = self.client.system_one( - state=state, questions=sdk_questions, model=self.model) + 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=dict(answer.probabilities), - confidence=answer.confidence) + probabilities=_probabilities( + name, answer.probabilities, question.criteria, str), + confidence=_confidence(name, answer.confidence)) elif isinstance(question, BooleanQuestion): - answers[name] = BooleanAnswer(probability=answer.noul) + _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=answer.score, - probabilities={ - int(level): probability - for level, probability in answer.probabilities.items() - }, - confidence=answer.confidence) + 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, diff --git a/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py index 2711029e4b17..5776fc9133e8 100644 --- a/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/typesafe_inference_test.py @@ -21,6 +21,8 @@ from types import SimpleNamespace from unittest import mock +from parameterized import parameterized + try: import typesafe_sdk except ImportError: @@ -35,20 +37,38 @@ 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': SimpleNamespace( - choice='billing', - probabilities={'billing': 1.0}, - confidence=0.9), - 'review': SimpleNamespace(noul=0.8), - 'urgency': SimpleNamespace( - score=0.7, probabilities={ - '0': 0.3, '1': 0.7 - }, confidence=0.6), + 'team': _choice_answer(), + 'review': typesafe_sdk.NoulAnswer.model_construct(noul=0.8), + 'urgency': _score_answer(), }, model='jev-test', request_id='request-1') @@ -91,6 +111,122 @@ def test_normalizes_all_primitives_and_reuses_client(self): 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') @@ -101,6 +237,20 @@ def test_provider_error_propagates_and_client_closes(self): '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 000000000000..982081593f5e --- /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 From 7ba1e16464e0da87faa2df5a5cc6236ba1d03223 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Tue, 13 Oct 2026 01:05:26 -0400 Subject: [PATCH 6/7] [Python] Run decision adapter tests in ML CI --- CHANGES.md | 1 + sdks/python/tox.ini | 1 + 2 files changed, 2 insertions(+) diff --git a/CHANGES.md b/CHANGES.md index 87b8ad672ebb..a12edb609c99 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -69,6 +69,7 @@ ## New Features / Improvements * (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)). ## Breaking Changes diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 8f59766612cd..994b3362d4b3 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -121,6 +121,7 @@ commands = # Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests. deps = accelerate>=1.6.0 + -r{toxinidir}/apache_beam/ml/inference/typesafe_tests_requirements.txt setenv = extras = test,gcp,dataframe,p312_ml_test commands = From 5655d4a65e683699dc3f4f58703209fb69bf1ad6 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Thu, 15 Oct 2026 01:02:04 -0400 Subject: [PATCH 7/7] [Python] Move decision SDK tests into the ML extra --- sdks/python/setup.py | 1 + sdks/python/tox.ini | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 24e9440361ca..d18375d6129d 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -613,6 +613,7 @@ def get_portability_package_data(): ] + 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 diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 66f8bddf734b..c88e5e2fea6a 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -121,7 +121,6 @@ commands = # Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests. deps = accelerate>=1.6.0 - -r{toxinidir}/apache_beam/ml/inference/typesafe_tests_requirements.txt setenv = extras = test,gcp,dataframe,p312_ml_test commands =