diff --git a/sdks/python/apache_beam/ml/inference/openai_inference.py b/sdks/python/apache_beam/ml/inference/openai_inference.py new file mode 100644 index 000000000000..a32b57f56ea7 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/openai_inference.py @@ -0,0 +1,392 @@ +# +# 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. +# + +"""A ModelHandler for OpenAI models using the OpenAI Python SDK. + +This module provides an integration between Apache Beam's RunInference +transform and OpenAI's API, enabling batch inference and embeddings in +Beam pipelines. + +Example usage:: + + import apache_beam as beam + from apache_beam.ml.inference.base import RunInference + from apache_beam.ml.inference.openai_inference import ( + OpenAIModelHandler, + chat_completion_from_string, + ) + + # Basic text generation with chat completions + model_handler = OpenAIModelHandler( + model_name='gpt-4o-mini', + api_key='your-api-key', + request_fn=chat_completion_from_string, + ) + + # With system prompt and structured output passed via inference_args + inference_args = { + 'system': 'You are a helpful assistant that responds concisely.', + 'response_format': { + 'type': 'json_schema', + 'json_schema': { + 'name': 'answer_response', + 'schema': { + 'type': 'object', + 'properties': { + 'answer': {'type': 'string'}, + 'confidence': {'type': 'number'}, + }, + 'required': ['answer', 'confidence'], + 'additionalProperties': False, + }, + 'strict': True, + }, + }, + } + + with beam.Pipeline() as p: + results = ( + p + | beam.Create(['What is Apache Beam?', 'Explain MapReduce.']) + | RunInference(model_handler, inference_args=inference_args) + ) +""" + +import logging +from collections.abc import Callable +from collections.abc import Iterable +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from typing import Any +from typing import Optional +from typing import TypeVar + +from openai import APIConnectionError +from openai import APIStatusError +from openai import OpenAI + +from apache_beam.ml.inference import utils +from apache_beam.ml.inference.base import PredictionResult +from apache_beam.ml.inference.base import RemoteModelHandler + +__all__ = [ + 'OpenAIModelHandler', + 'chat_completion_from_string', + 'chat_completion_from_conversation', + 'embedding_from_string', +] + +LOGGER = logging.getLogger("OpenAIModelHandler") + +# Upper bound on the number of in-flight HTTP requests issued concurrently for +# a single batch by request functions that call per-item endpoints (e.g. chat +# completions). Bounds thread usage when adaptive batching produces large +# batches; use max_batch_size to further limit per-batch concurrency. +_MAX_CONCURRENT_REQUESTS_PER_BATCH = 32 + +_InputT = TypeVar('_InputT') +_OutputT = TypeVar('_OutputT') + + +def _retry_on_appropriate_error(exception: Exception) -> bool: + """Retry filter that returns True for retriable OpenAI API errors. + + Retries on HTTP 429 (rate limiting), HTTP 5xx (server errors), and + connection / timeout errors. + + Args: + exception: the exception encountered during the request/response loop. + + Returns: + True if the exception is retriable (429, 5xx, or connection error), + False otherwise. + """ + if isinstance(exception, APIConnectionError): + return True + if isinstance(exception, APIStatusError): + return exception.status_code == 429 or exception.status_code >= 500 + return False + + +def _fan_out(fn: Callable[[_InputT], _OutputT], + items: Sequence[_InputT]) -> list[_OutputT]: + """Applies fn to each item concurrently and returns results in input order. + + The synchronous OpenAI client is safe to share across threads (its + underlying connection pool is guarded by thread locks), so per-item requests + are dispatched on a bounded thread pool while reusing the client's pooled + connections. If any request raises, pending requests are cancelled and the + first exception (in input order) is propagated so Beam's retry logic can + handle the batch. + """ + if len(items) <= 1: + return [fn(item) for item in items] + executor = ThreadPoolExecutor( + max_workers=min(len(items), _MAX_CONCURRENT_REQUESTS_PER_BATCH), + thread_name_prefix='OpenAIModelHandler') + try: + return list(executor.map(fn, items)) + finally: + executor.shutdown(wait=True, cancel_futures=True) + + +def chat_completion_from_string( + model_name: str, + batch: Sequence[str], + client: OpenAI, + inference_args: dict[str, Any]) -> list[Any]: + """Request function that sends string prompts to OpenAI's Chat Completions API. + + The Chat Completions endpoint accepts a single conversation per HTTP + request, so each string in the batch is sent as a separate user message + request, dispatched concurrently. If a 'system' parameter is provided in + inference_args, a system message is prepended. The results are returned as a + list of ChatCompletion response objects matching the batch order. + + Args: + model_name: the OpenAI model to use (e.g. 'gpt-4o', 'gpt-4o-mini'). + batch: the string prompts to send to OpenAI. + client: the OpenAI client instance. + inference_args: additional arguments passed to the chat.completions.create + call (e.g. 'temperature', 'max_tokens', 'response_format', 'system'). + """ + inf_args = dict(inference_args) + system = inf_args.pop('system', None) + + def _request(prompt: str) -> Any: + messages: list[dict[str, Any]] = [] + if system is not None: + messages.append({"role": "system", "content": system}) + messages.append({"role": "user", "content": prompt}) + return client.chat.completions.create( + model=model_name, messages=messages, **inf_args) + + return _fan_out(_request, batch) + + +def chat_completion_from_conversation( + model_name: str, + batch: Sequence[list[dict[str, Any]]], + client: OpenAI, + inference_args: dict[str, Any]) -> list[Any]: + """Request function that sends multi-turn conversations to OpenAI. + + Each element in the batch is a list of message dicts (e.g. with 'role' and + 'content' keys), representing a multi-turn conversation. Conversations are + sent as separate requests, dispatched concurrently. If a 'system' parameter + is provided in inference_args, a system message is prepended. + + Args: + model_name: the OpenAI model to use. + batch: a sequence of conversations (each a list of message dicts). + client: the OpenAI client instance. + inference_args: additional arguments passed to the chat.completions.create + call. + """ + inf_args = dict(inference_args) + system = inf_args.pop('system', None) + + def _request(conversation: list[dict[str, Any]]) -> Any: + messages: list[dict[str, Any]] = [] + if system is not None: + messages.append({"role": "system", "content": system}) + messages.extend(conversation) + return client.chat.completions.create( + model=model_name, messages=messages, **inf_args) + + return _fan_out(_request, batch) + + +def embedding_from_string( + model_name: str, + batch: Sequence[str], + client: OpenAI, + inference_args: dict[str, Any]) -> list[Any]: + """Request function that sends string inputs to OpenAI's Embeddings API. + + The Embeddings endpoint natively accepts a list of inputs, so the entire + batch is sent in a single request to ``client.embeddings.create``. The + returned embeddings are sorted by their index to guarantee ordering matches + the batch. + + Args: + model_name: the OpenAI embedding model to use (e.g. 'text-embedding-3-small'). + batch: the string inputs to embed. + client: the OpenAI client instance. + inference_args: additional arguments passed to the embeddings.create call + (e.g. 'dimensions', 'encoding_format'). + + Returns: + A list of Embedding objects matching the batch order. + """ + inf_args = dict(inference_args) + response = client.embeddings.create( + model=model_name, input=list(batch), **inf_args) + return sorted(response.data, key=lambda x: x.index) + + +class OpenAIModelHandler(RemoteModelHandler[Any, PredictionResult, OpenAI]): + def __init__( + self, + model_name: str, + request_fn: Callable[[str, Sequence[Any], OpenAI, dict[str, Any]], Any], + api_key: Optional[str] = None, + *, + organization: Optional[str] = None, + project: Optional[str] = None, + base_url: Optional[str] = None, + client_args: Optional[dict[str, Any]] = None, + min_batch_size: Optional[int] = None, + max_batch_size: Optional[int] = None, + max_batch_duration_secs: Optional[int] = None, + max_batch_weight: Optional[int] = None, + element_size_fn: Optional[Callable[[Any], int]] = None, + batch_length_fn: Optional[Callable[[Any], int]] = None, + batch_bucket_boundaries: Optional[list[int]] = None, + **kwargs): + """Implementation of the ModelHandler interface for OpenAI models. + + **NOTE:** This API and its implementation are under development and + do not provide backward compatibility guarantees. + + This handler connects to the OpenAI API using the OpenAI Python SDK + to run inference using models such as GPT-4o, GPT-4o-mini, or embedding + models. It supports chat completions from string prompts or multi-turn + conversations (with each batch's per-item requests dispatched + concurrently), natively batched embeddings, system prompts, structured + outputs, and custom OpenAI-compatible endpoints via `base_url`. + + A single thread-safe ``OpenAI`` client is created per worker process and + shared across bundle-processing threads, so pooled HTTP connections are + reused across batches. + + Args: + model_name: the OpenAI model to send requests to (e.g. + 'gpt-4o', 'gpt-4o-mini', 'text-embedding-3-small'). + request_fn: the function to use to send requests. Should take the + model name, batch, client, and inference_args and return the + responses from OpenAI. Built-in options are chat_completion_from_string, + chat_completion_from_conversation, and embedding_from_string. + api_key: the OpenAI API key. If not provided, the client will + look for the OPENAI_API_KEY environment variable. + organization: optional OpenAI organization ID. If not provided, + the client will look for the OPENAI_ORG_ID environment variable. + project: optional OpenAI project ID. If not provided, the client + will look for the OPENAI_PROJECT_ID environment variable. + base_url: optional base URL for the API requests. This can be used + to target OpenAI-compatible servers (such as vLLM, Ollama, or + Azure OpenAI endpoints). + client_args: optional dictionary of additional keyword arguments + passed when instantiating the OpenAI client (e.g. timeout, + default_headers). + min_batch_size: optional. the minimum batch size to use when + batching inputs. + max_batch_size: optional. the maximum batch size to use when + batching inputs. + max_batch_duration_secs: optional. the maximum amount of time to + buffer a batch before emitting; used in streaming contexts. + max_batch_weight: optional. the maximum total weight of a batch. + element_size_fn: optional. a function that returns the size + (weight) of an element. + batch_length_fn: optional. a callable that returns the length of an + element for length-aware batching. + batch_bucket_boundaries: optional. a sorted list of positive boundary + values for length-aware batching buckets. + **kwargs: optional arguments passed to RemoteModelHandler (e.g. + num_retries, throttle_delay_secs, rate_limiter). + """ + self._batching_kwargs = {} + self._env_vars = kwargs.get('env_vars', {}) + if min_batch_size is not None: + self._batching_kwargs["min_batch_size"] = min_batch_size + if max_batch_size is not None: + self._batching_kwargs["max_batch_size"] = max_batch_size + if max_batch_duration_secs is not None: + self._batching_kwargs["max_batch_duration_secs"] = max_batch_duration_secs + if max_batch_weight is not None: + self._batching_kwargs["max_batch_weight"] = max_batch_weight + if element_size_fn is not None: + self._batching_kwargs['element_size_fn'] = element_size_fn + if batch_length_fn is not None: + self._batching_kwargs['length_fn'] = batch_length_fn + if batch_bucket_boundaries is not None: + self._batching_kwargs['bucket_boundaries'] = batch_bucket_boundaries + + self.model_name = model_name + self.request_fn = request_fn + self.api_key = api_key + self.organization = organization + self.project = project + self.base_url = base_url + self.client_args = client_args + + retry_filter = kwargs.pop('retry_filter', _retry_on_appropriate_error) + + super().__init__( + namespace='OpenAIModelHandler', retry_filter=retry_filter, **kwargs) + + def batch_elements_kwargs(self): + return self._batching_kwargs + + def create_client(self) -> OpenAI: + """Creates the OpenAI client used to send requests. + + Sets ``max_retries=0`` by default so that Apache Beam's + ``RemoteModelHandler`` retry and client-side throttling mechanics are the + sole retry layer. If api_key, organization, project, base_url, or + client_args were provided at construction time, they are passed to the + OpenAI constructor. Otherwise, the client falls back to standard + environment variables (OPENAI_API_KEY, OPENAI_ORG_ID, etc.). + """ + params: dict[str, Any] = {'max_retries': 0} + if self.api_key is not None: + params['api_key'] = self.api_key + if self.organization is not None: + params['organization'] = self.organization + if self.project is not None: + params['project'] = self.project + if self.base_url is not None: + params['base_url'] = self.base_url + if self.client_args is not None: + params.update(self.client_args) + return OpenAI(**params) + + def request( + self, + batch: Sequence[Any], + model: OpenAI, + inference_args: Optional[dict[str, Any]] = None + ) -> Iterable[PredictionResult]: + """Sends a prediction request to the OpenAI API. + + Args: + batch: a sequence of inputs to be passed to the request function. + model: an OpenAI client instance. + inference_args: additional arguments to send as part of the + prediction request (e.g. temperature, max_tokens, system, + response_format). + + Returns: + An iterable of PredictionResults. + """ + if inference_args is None: + inference_args = {} + else: + inference_args = dict(inference_args) + + responses = self.request_fn(self.model_name, batch, model, inference_args) + return utils._convert_to_result(batch, responses, self.model_name) diff --git a/sdks/python/apache_beam/ml/inference/openai_inference_it_test.py b/sdks/python/apache_beam/ml/inference/openai_inference_it_test.py new file mode 100644 index 000000000000..88b50436ff3d --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/openai_inference_it_test.py @@ -0,0 +1,211 @@ +# +# 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. +# + +"""End-to-End test for OpenAI Remote Inference""" + +import logging +import os +import unittest + +import pytest + +try: + from openai import OpenAI # pylint: disable=unused-import + + from apache_beam.ml.inference.openai_inference import OpenAIModelHandler + from apache_beam.ml.inference.openai_inference import chat_completion_from_conversation + from apache_beam.ml.inference.openai_inference import chat_completion_from_string + from apache_beam.ml.inference.openai_inference import embedding_from_string +except ImportError: + raise unittest.SkipTest("OpenAI dependencies are not installed") + +import apache_beam as beam +from apache_beam.ml.inference.base import RunInference +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import is_not_empty + +_OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY', None) +_TEST_CHAT_MODEL = 'gpt-4o-mini' +_TEST_EMBEDDING_MODEL = 'text-embedding-3-small' + + +def _extract_chat_text(prediction_result): + return prediction_result.inference.choices[0].message.content + + +def _extract_embedding(prediction_result): + return prediction_result.inference.embedding + + +@pytest.mark.openai_postcommit +class OpenAIInferenceIT(unittest.TestCase): + @unittest.skipIf( + _OPENAI_API_KEY is None, 'OPENAI_API_KEY environment variable is not set') + def test_openai_chat_text_generation(self): + handler = OpenAIModelHandler( + model_name=_TEST_CHAT_MODEL, + request_fn=chat_completion_from_string, + api_key=_OPENAI_API_KEY, + max_batch_size=2, + ) + + prompts = [ + 'What is Apache Beam in one sentence?', + 'Name three distributed computing frameworks.', + ] + + with TestPipeline() as p: + results = ( + p + | beam.Create(prompts) + | RunInference(handler) + | beam.Map(_extract_chat_text)) + assert_that(results, is_not_empty()) + + @unittest.skipIf( + _OPENAI_API_KEY is None, 'OPENAI_API_KEY environment variable is not set') + def test_openai_conversation(self): + handler = OpenAIModelHandler( + model_name=_TEST_CHAT_MODEL, + request_fn=chat_completion_from_conversation, + api_key=_OPENAI_API_KEY, + max_batch_size=1, + ) + + conversations = [ + [ + { + "role": "user", "content": "What is 2 + 2?" + }, + { + "role": "assistant", "content": "4" + }, + { + "role": "user", "content": "Add 3 to that." + }, + ], + ] + + with TestPipeline() as p: + results = ( + p + | beam.Create(conversations) + | RunInference(handler) + | beam.Map(_extract_chat_text)) + assert_that(results, is_not_empty()) + + @unittest.skipIf( + _OPENAI_API_KEY is None, 'OPENAI_API_KEY environment variable is not set') + def test_openai_with_system_prompt(self): + handler = OpenAIModelHandler( + model_name=_TEST_CHAT_MODEL, + request_fn=chat_completion_from_string, + api_key=_OPENAI_API_KEY, + max_batch_size=1, + ) + + prompts = ['What is your name?'] + inference_args = { + 'system': 'You are a pirate. Respond only in pirate speak.', + } + + with TestPipeline() as p: + results = ( + p + | beam.Create(prompts) + | RunInference(handler, inference_args=inference_args) + | beam.Map(_extract_chat_text)) + assert_that(results, is_not_empty()) + + @unittest.skipIf( + _OPENAI_API_KEY is None, 'OPENAI_API_KEY environment variable is not set') + def test_openai_system_prompt_with_structured_output(self): + handler = OpenAIModelHandler( + model_name=_TEST_CHAT_MODEL, + request_fn=chat_completion_from_string, + api_key=_OPENAI_API_KEY, + max_batch_size=1, + ) + + prompts = ['Count from 1 to 5.'] + inference_args = { + 'system': 'You respond only with JSON conforming to the schema.', + 'response_format': { + 'type': 'json_schema', + 'json_schema': { + 'name': 'fizz_buzz_response', + 'schema': { + 'type': 'object', + 'properties': { + 'items': { + 'type': 'array', + 'items': { + 'type': 'object', + 'properties': { + 'value': { + 'type': 'string' + }, + }, + 'required': ['value'], + 'additionalProperties': False, + }, + }, + }, + 'required': ['items'], + 'additionalProperties': False, + }, + 'strict': True, + }, + }, + } + + with TestPipeline() as p: + results = ( + p + | beam.Create(prompts) + | RunInference(handler, inference_args=inference_args) + | beam.Map(_extract_chat_text)) + assert_that(results, is_not_empty()) + + @unittest.skipIf( + _OPENAI_API_KEY is None, 'OPENAI_API_KEY environment variable is not set') + def test_openai_embeddings(self): + handler = OpenAIModelHandler( + model_name=_TEST_EMBEDDING_MODEL, + request_fn=embedding_from_string, + api_key=_OPENAI_API_KEY, + max_batch_size=2, + ) + + prompts = [ + 'What is Apache Beam?', + 'Distributed stream and batch processing.', + ] + + with TestPipeline() as p: + results = ( + p + | beam.Create(prompts) + | RunInference(handler) + | beam.Map(_extract_embedding)) + assert_that(results, is_not_empty()) + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/openai_inference_test.py b/sdks/python/apache_beam/ml/inference/openai_inference_test.py new file mode 100644 index 000000000000..23a13980ea8b --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/openai_inference_test.py @@ -0,0 +1,624 @@ +# +# 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. +# +# pytype: skip-file + +import threading +import time +import unittest +from dataclasses import dataclass +from unittest import mock + +import httpx + +try: + from openai import APIConnectionError + from openai import APIStatusError + from openai import APITimeoutError + + from apache_beam.ml.inference.openai_inference import OpenAIModelHandler + from apache_beam.ml.inference.openai_inference import _retry_on_appropriate_error + from apache_beam.ml.inference.openai_inference import chat_completion_from_conversation + from apache_beam.ml.inference.openai_inference import chat_completion_from_string + from apache_beam.ml.inference.openai_inference import embedding_from_string +except ImportError: + raise unittest.SkipTest('OpenAI dependencies are not installed') + +import apache_beam as beam +from apache_beam.ml.inference.base import PredictionResult +from apache_beam.ml.inference.base import RunInference +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to + +_TEST_MODEL = 'gpt-4o-mini' +_TEST_EMBEDDING_MODEL = 'text-embedding-3-small' + + +@dataclass +class FakeChatMessage: + content: str + role: str = 'assistant' + + +@dataclass +class FakeChoice: + message: FakeChatMessage + finish_reason: str = 'stop' + index: int = 0 + + +@dataclass +class FakeChatCompletion: + """Picklable stand-in for ChatCompletion.""" + choices: list + model: str = _TEST_MODEL + id: str = 'fake-id' + + +@dataclass +class FakeEmbedding: + """Picklable stand-in for Embedding.""" + embedding: list + index: int + object: str = 'embedding' + + +@dataclass +class FakeEmbeddingResponse: + """Picklable stand-in for CreateEmbeddingResponse.""" + data: list + model: str = _TEST_EMBEDDING_MODEL + + +def _make_fake_chat_response(text): + return FakeChatCompletion( + choices=[FakeChoice(message=FakeChatMessage(content=text))]) + + +def _make_fake_embedding_response(batch, embedding_dim=4): + data = [ + FakeEmbedding(embedding=[0.1 * (i + 1)] * embedding_dim, index=i) + for i in range(len(batch)) + ] + return FakeEmbeddingResponse(data=data) + + +def _make_mock_http_response(status_code: int) -> httpx.Response: + req = httpx.Request('POST', 'https://api.openai.com/v1/chat/completions') + return httpx.Response(status_code=status_code, request=req) + + +def _make_mock_client() -> mock.MagicMock: + client = mock.MagicMock() + client.chat.completions.create = mock.MagicMock() + client.embeddings.create = mock.MagicMock() + return client + + +def _echo_last_message(**kwargs): + """Returns a fake chat response derived from the request's last message. + + Chat requests within a batch are issued concurrently, so call order is + nondeterministic; keying the response on the request contents lets tests + check that results are returned in batch order. + """ + return _make_fake_chat_response( + f"answer for: {kwargs['messages'][-1]['content']}") + + +class RetryOnErrorTest(unittest.TestCase): + def test_retry_on_rate_limit(self): + e = APIStatusError( + message="Rate limited", + response=_make_mock_http_response(429), + body=None) + self.assertTrue(_retry_on_appropriate_error(e)) + + def test_retry_on_server_error(self): + e = APIStatusError( + message="Internal server error", + response=_make_mock_http_response(500), + body=None) + self.assertTrue(_retry_on_appropriate_error(e)) + + def test_retry_on_503(self): + e = APIStatusError( + message="Service unavailable", + response=_make_mock_http_response(503), + body=None) + self.assertTrue(_retry_on_appropriate_error(e)) + + def test_retry_on_connection_error(self): + req = httpx.Request('POST', 'https://api.openai.com/v1/chat/completions') + e = APIConnectionError(request=req) + self.assertTrue(_retry_on_appropriate_error(e)) + + def test_retry_on_timeout_error(self): + req = httpx.Request('POST', 'https://api.openai.com/v1/chat/completions') + e = APITimeoutError(request=req) + self.assertTrue(_retry_on_appropriate_error(e)) + + def test_no_retry_on_400(self): + e = APIStatusError( + message="Bad request", + response=_make_mock_http_response(400), + body=None) + self.assertFalse(_retry_on_appropriate_error(e)) + + def test_no_retry_on_401(self): + e = APIStatusError( + message="Unauthorized", + response=_make_mock_http_response(401), + body=None) + self.assertFalse(_retry_on_appropriate_error(e)) + + def test_no_retry_on_403(self): + e = APIStatusError( + message="Forbidden", response=_make_mock_http_response(403), body=None) + self.assertFalse(_retry_on_appropriate_error(e)) + + def test_no_retry_on_non_api_error(self): + self.assertFalse(_retry_on_appropriate_error(ValueError("oops"))) + self.assertFalse(_retry_on_appropriate_error(RuntimeError("fail"))) + + +class ChatCompletionFromStringTest(unittest.TestCase): + def test_sends_each_prompt(self): + client = _make_mock_client() + client.chat.completions.create.side_effect = _echo_last_message + results = chat_completion_from_string( + _TEST_MODEL, ['hello', 'world'], client, {}) + self.assertEqual(len(results), 2) + self.assertEqual(results[0].choices[0].message.content, "answer for: hello") + self.assertEqual(results[1].choices[0].message.content, "answer for: world") + self.assertEqual(client.chat.completions.create.call_count, 2) + + calls = client.chat.completions.create.call_args_list + sent_messages = sorted((c.kwargs['messages'] for c in calls), + key=lambda m: m[0]['content']) + self.assertEqual( + sent_messages, + [[{ + "role": "user", "content": "hello" + }], [{ + "role": "user", "content": "world" + }]]) + for c in calls: + self.assertEqual(c.kwargs['model'], _TEST_MODEL) + + def test_requests_run_concurrently_and_preserve_order(self): + lock = threading.Lock() + state = {'in_flight': 0, 'max_in_flight': 0} + completion_order = [] + prompts = [f'p{i}' for i in range(4)] + + def slow_create(**kwargs): + prompt = kwargs['messages'][-1]['content'] + with lock: + state['in_flight'] += 1 + state['max_in_flight'] = max(state['max_in_flight'], state['in_flight']) + # Earlier prompts sleep longer so they finish last. + time.sleep(0.05 * (len(prompts) - prompts.index(prompt))) + with lock: + state['in_flight'] -= 1 + completion_order.append(prompt) + return _make_fake_chat_response(f"answer for: {prompt}") + + client = _make_mock_client() + client.chat.completions.create.side_effect = slow_create + results = chat_completion_from_string(_TEST_MODEL, prompts, client, {}) + + self.assertGreater(state['max_in_flight'], 1) + self.assertNotEqual(completion_order, prompts) + self.assertEqual([r.choices[0].message.content for r in results], + [f"answer for: {p}" for p in prompts]) + + def test_concurrency_is_capped(self): + lock = threading.Lock() + state = {'in_flight': 0, 'max_in_flight': 0} + + def slow_create(**kwargs): + with lock: + state['in_flight'] += 1 + state['max_in_flight'] = max(state['max_in_flight'], state['in_flight']) + time.sleep(0.02) + with lock: + state['in_flight'] -= 1 + return _make_fake_chat_response("ok") + + client = _make_mock_client() + client.chat.completions.create.side_effect = slow_create + with mock.patch('apache_beam.ml.inference.openai_inference.' + '_MAX_CONCURRENT_REQUESTS_PER_BATCH', + 2): + results = chat_completion_from_string( + _TEST_MODEL, [f'p{i}' for i in range(6)], client, {}) + + self.assertEqual(len(results), 6) + self.assertLessEqual(state['max_in_flight'], 2) + + def test_propagates_request_error(self): + rate_limited = APIStatusError( + message="Rate limited", + response=_make_mock_http_response(429), + body=None) + + def create(**kwargs): + if kwargs['messages'][-1]['content'] == 'bad': + raise rate_limited + return _make_fake_chat_response("ok") + + client = _make_mock_client() + client.chat.completions.create.side_effect = create + with self.assertRaises(APIStatusError) as ctx: + chat_completion_from_string( + _TEST_MODEL, ['good', 'bad', 'good'], client, {}) + self.assertIs(ctx.exception, rate_limited) + self.assertTrue(_retry_on_appropriate_error(ctx.exception)) + + def test_passes_inference_args(self): + client = _make_mock_client() + client.chat.completions.create.return_value = _make_fake_chat_response("ok") + chat_completion_from_string( + _TEST_MODEL, ['test'], client, { + 'max_tokens': 2048, 'temperature': 0.5 + }) + call_args = client.chat.completions.create.call_args + self.assertEqual(call_args.kwargs['max_tokens'], 2048) + self.assertEqual(call_args.kwargs['temperature'], 0.5) + + def test_prepends_system_prompt(self): + client = _make_mock_client() + client.chat.completions.create.return_value = _make_fake_chat_response("ok") + chat_completion_from_string( + _TEST_MODEL, ['test'], client, {'system': 'You are helpful.'}) + call_args = client.chat.completions.create.call_args + self.assertEqual( + call_args.kwargs['messages'], + [ + { + "role": "system", "content": "You are helpful." + }, + { + "role": "user", "content": "test" + }, + ]) + self.assertNotIn('system', call_args.kwargs) + + +class ChatCompletionFromConversationTest(unittest.TestCase): + def test_sends_conversation(self): + client = _make_mock_client() + client.chat.completions.create.return_value = ( + _make_fake_chat_response("Paris!")) + convo = [ + { + "role": "user", "content": "What is the capital of France?" + }, + ] + results = chat_completion_from_conversation( + _TEST_MODEL, [convo], client, {}) + self.assertEqual(len(results), 1) + call_args = client.chat.completions.create.call_args + self.assertEqual(call_args.kwargs['messages'], convo) + + def test_prepends_system_prompt_to_conversation(self): + client = _make_mock_client() + client.chat.completions.create.return_value = _make_fake_chat_response("ok") + convo = [ + { + "role": "user", "content": "Hello" + }, + ] + chat_completion_from_conversation( + _TEST_MODEL, [convo], client, {'system': 'Be concise.'}) + call_args = client.chat.completions.create.call_args + self.assertEqual( + call_args.kwargs['messages'], + [ + { + "role": "system", "content": "Be concise." + }, + { + "role": "user", "content": "Hello" + }, + ]) + self.assertNotIn('system', call_args.kwargs) + + +class EmbeddingFromStringTest(unittest.TestCase): + def test_sends_batch_to_embeddings(self): + client = _make_mock_client() + fake_resp = _make_fake_embedding_response(['hello', 'world']) + client.embeddings.create.return_value = fake_resp + + results = embedding_from_string( + _TEST_EMBEDDING_MODEL, ['hello', 'world'], client, {}) + self.assertEqual(len(results), 2) + client.embeddings.create.assert_called_once_with( + model=_TEST_EMBEDDING_MODEL, input=['hello', 'world']) + + def test_preserves_order_via_index(self): + client = _make_mock_client() + # Out of order data + fake_resp = FakeEmbeddingResponse( + data=[ + FakeEmbedding(embedding=[0.2], index=1), + FakeEmbedding(embedding=[0.1], index=0), + ]) + client.embeddings.create.return_value = fake_resp + + results = embedding_from_string( + _TEST_EMBEDDING_MODEL, ['first', 'second'], client, {}) + self.assertEqual(results[0].index, 0) + self.assertEqual(results[0].embedding, [0.1]) + self.assertEqual(results[1].index, 1) + self.assertEqual(results[1].embedding, [0.2]) + + def test_passes_inference_args(self): + client = _make_mock_client() + client.embeddings.create.return_value = _make_fake_embedding_response(['a']) + embedding_from_string( + _TEST_EMBEDDING_MODEL, ['a'], client, {'dimensions': 256}) + call_args = client.embeddings.create.call_args + self.assertEqual(call_args.kwargs['dimensions'], 256) + + +class OpenAIModelHandlerTest(unittest.TestCase): + @mock.patch('apache_beam.ml.inference.openai_inference.OpenAI') + def test_create_client_with_api_key(self, mock_openai): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='test-key-123') + handler.create_client() + mock_openai.assert_called_once_with(api_key='test-key-123', max_retries=0) + + @mock.patch('apache_beam.ml.inference.openai_inference.OpenAI') + def test_create_client_from_env(self, mock_openai): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, request_fn=chat_completion_from_string) + handler.create_client() + mock_openai.assert_called_once_with(max_retries=0) + + @mock.patch('apache_beam.ml.inference.openai_inference.OpenAI') + def test_create_client_with_all_options(self, mock_openai): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='test-key-123', + organization='test-org', + project='test-proj', + base_url='https://custom.endpoint.com/v1', + client_args={'timeout': 30.0}) + handler.create_client() + mock_openai.assert_called_once_with( + api_key='test-key-123', + organization='test-org', + project='test-proj', + base_url='https://custom.endpoint.com/v1', + max_retries=0, + timeout=30.0) + + def test_request_returns_prediction_results(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='fake') + mock_client = _make_mock_client() + responses = { + 'q1': _make_fake_chat_response("answer 1"), + 'q2': _make_fake_chat_response("answer 2"), + } + mock_client.chat.completions.create.side_effect = ( + lambda **kwargs: responses[kwargs['messages'][-1]['content']]) + + results = list(handler.request(['q1', 'q2'], mock_client, {})) + + self.assertEqual(len(results), 2) + self.assertIsInstance(results[0], PredictionResult) + self.assertEqual(results[0].example, 'q1') + self.assertEqual(results[0].inference, responses['q1']) + self.assertEqual(results[0].model_id, _TEST_MODEL) + self.assertEqual(results[1].example, 'q2') + self.assertEqual(results[1].inference, responses['q2']) + + def test_batch_elements_kwargs(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='fake', + min_batch_size=2, + max_batch_size=10, + max_batch_duration_secs=5, + max_batch_weight=100) + kwargs = handler.batch_elements_kwargs() + self.assertEqual(kwargs['min_batch_size'], 2) + self.assertEqual(kwargs['max_batch_size'], 10) + self.assertEqual(kwargs['max_batch_duration_secs'], 5) + self.assertEqual(kwargs['max_batch_weight'], 100) + + def test_custom_retry_filter(self): + custom_filter = lambda e: False + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + retry_filter=custom_filter) + self.assertEqual(handler.retry_filter, custom_filter) + + +class SystemPromptTest(unittest.TestCase): + def test_system_prompt_passed_via_inference_args(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='fake') + mock_client = _make_mock_client() + mock_client.chat.completions.create.return_value = ( + _make_fake_chat_response("ok")) + + handler.request(['test'], mock_client, {'system': 'Be concise.'}) + + call_args = mock_client.chat.completions.create.call_args + self.assertEqual( + call_args.kwargs['messages'][0], { + "role": "system", "content": "Be concise." + }) + self.assertNotIn('system', call_args.kwargs) + + def test_no_system_prompt_when_omitted(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='fake') + mock_client = _make_mock_client() + mock_client.chat.completions.create.return_value = ( + _make_fake_chat_response("ok")) + + handler.request(['test'], mock_client, {}) + + call_args = mock_client.chat.completions.create.call_args + self.assertEqual( + call_args.kwargs['messages'], [{ + "role": "user", "content": "test" + }]) + + +class ResponseFormatTest(unittest.TestCase): + _SCHEMA = { + 'type': 'json_schema', + 'json_schema': { + 'name': 'answer_schema', + 'schema': { + 'type': 'object', + 'properties': { + 'answer': { + 'type': 'string' + } + }, + 'required': ['answer'], + 'additionalProperties': False, + }, + 'strict': True, + }, + } + + def test_response_format_passed_via_inference_args(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='fake') + mock_client = _make_mock_client() + mock_client.chat.completions.create.return_value = ( + _make_fake_chat_response('{"answer":"ok"}')) + + handler.request(['test'], mock_client, {'response_format': self._SCHEMA}) + + call_args = mock_client.chat.completions.create.call_args + self.assertEqual(call_args.kwargs['response_format'], self._SCHEMA) + + def test_no_response_format_when_omitted(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=chat_completion_from_string, + api_key='fake') + mock_client = _make_mock_client() + mock_client.chat.completions.create.return_value = ( + _make_fake_chat_response("ok")) + + handler.request(['test'], mock_client, {}) + + call_args = mock_client.chat.completions.create.call_args + self.assertNotIn('response_format', call_args.kwargs) + + +def _fake_chat_request_fn(model_name, batch, client, inference_args): + """A picklable request function that returns fake chat responses.""" + return [ + FakeChatCompletion( + choices=[ + FakeChoice(message=FakeChatMessage(content=f'answer for: {p}')) + ]) for p in batch + ] + + +def _fake_embedding_request_fn(model_name, batch, client, inference_args): + """A picklable request function that returns fake embedding responses.""" + return [ + FakeEmbedding(embedding=[0.1, 0.2], index=i) for i, _ in enumerate(batch) + ] + + +class OpenAIRunInferencePipelineTest(unittest.TestCase): + def test_pipeline_chat_e2e(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=_fake_chat_request_fn, + api_key='fake-key', + max_batch_size=5, + ) + + prompts = ['What is Beam?', 'What is MapReduce?'] + + with TestPipeline() as p: + results = ( + p + | beam.Create(prompts) + | RunInference(handler) + | beam.Map(lambda r: r.example)) + assert_that(results, equal_to(prompts)) + + def test_pipeline_with_system_prompt(self): + handler = OpenAIModelHandler( + model_name=_TEST_MODEL, + request_fn=_fake_chat_request_fn, + api_key='fake-key', + max_batch_size=5, + ) + + prompts = ['Tell me about Beam.'] + + with TestPipeline() as p: + results = ( + p + | beam.Create(prompts) + | RunInference( + handler, inference_args={'system': 'You respond in haiku form.'}) + | beam.Map(lambda r: r.example)) + assert_that(results, equal_to(prompts)) + + def test_pipeline_embeddings_e2e(self): + handler = OpenAIModelHandler( + model_name=_TEST_EMBEDDING_MODEL, + request_fn=_fake_embedding_request_fn, + api_key='fake-key', + max_batch_size=5, + ) + + inputs = ['text one', 'text two'] + + with TestPipeline() as p: + results = ( + p + | beam.Create(inputs) + | RunInference(handler) + | beam.Map(lambda r: r.example)) + assert_that(results, equal_to(inputs)) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/ml/inference/openai_tests_requirements.txt b/sdks/python/apache_beam/ml/inference/openai_tests_requirements.txt new file mode 100644 index 000000000000..17bd0f3c82e2 --- /dev/null +++ b/sdks/python/apache_beam/ml/inference/openai_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. +# + +openai>=1.0.0 diff --git a/sdks/python/pytest.ini b/sdks/python/pytest.ini index 082c3ed634b1..584258902750 100644 --- a/sdks/python/pytest.ini +++ b/sdks/python/pytest.ini @@ -72,6 +72,7 @@ markers = uses_feast: tests that uses feast in some way gemini_postcommit: gemini postcommits that need additional deps. anthropic_postcommit: anthropic postcommits that need additional deps. + openai_postcommit: openai postcommits that need additional deps. require_docker_in_docker: tests that require running Docker inside Docker (Docker-in-Docker), which is not supported on Beam’s self-hosted runners. Context: https://github.com/apache/beam/pull/35585 uses_dill: tests that require dill pickle library. diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index 7e6919d18ab8..23b2ea888200 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -597,6 +597,34 @@ task anthropicInferenceTest { } } +// OpenAI RunInference IT tests +task openaiInferenceTest { + dependsOn 'initializeForDataflowJob' + dependsOn ':sdks:python:sdist' + def requirementsFile = "${rootDir}/sdks/python/apache_beam/ml/inference/openai_tests_requirements.txt" + doFirst { + exec { + executable 'sh' + args '-c', ". ${envdir}/bin/activate && pip install -r $requirementsFile" + } + } + doLast { + def testOpts = basicTestOpts + def argMap = [ + "test_opts": testOpts, + "suite": "OpenAITests-df-py${pythonVersionSuffix}", + "collect": "openai_postcommit" , + "runner": "TestDataflowRunner", + "requirements_file": "$requirementsFile" + ] + def cmdArgs = mapToArgString(argMap) + exec { + executable 'sh' + args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $cmdArgs" + } + } +} + task installTFTRequirements { dependsOn 'initializeForDataflowJob' doLast {