Package org.apache.beam.sdk.ml.inference.openai

Class OpenAIModelHandler

java.lang.Object

org.apache.beam.sdk.ml.inference.openai.OpenAIModelHandler

All Implemented Interfaces:

BaseModelHandler<OpenAIModelParameters,OpenAIModelInput,OpenAIModelResponse>

public class OpenAIModelHandler
extends Object
implements BaseModelHandler<OpenAIModelParameters,OpenAIModelInput,OpenAIModelResponse>

Model handler for OpenAI API inference requests.

This handler manages communication with OpenAI's API, including client initialization, request formatting, and response parsing. It uses OpenAI's structured output feature to ensure reliable input-output pairing.

Usage

 OpenAIModelParameters params = OpenAIModelParameters.builder()
     .apiKey("sk-...")
     .modelName("gpt-4")
     .instructionPrompt("Classify the following text into one of the categories: {CATEGORIES}")
     .build();

 PCollection<OpenAIModelInput> inputs = ...;
 PCollection<Iterable<PredictionResult<OpenAIModelInput, OpenAIModelResponse>>> results =
     inputs.apply(
         RemoteInference.<OpenAIModelInput, OpenAIModelResponse>invoke()
             .handler(OpenAIModelHandler.class)
             .withParameters(params)
     );
 

Nested Class Summary

Nested Classes

Modifier and Type	Class	Description
static class	OpenAIModelHandler.Response	Schema class for structured output response.
static class	OpenAIModelHandler.StructuredInputOutput	Schema class for structured output containing multiple responses.

Constructor Summary

Constructors

Constructor	Description
OpenAIModelHandler()

Method Summary

Modifier and Type	Method	Description
void	createClient(OpenAIModelParameters parameters)	Initializes the OpenAI client with the provided parameters.
Iterable<PredictionResult<OpenAIModelInput,OpenAIModelResponse>>	request(List<OpenAIModelInput> input)	Performs inference on a batch of inputs using the OpenAI Client.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Details

OpenAIModelHandler

public OpenAIModelHandler()

Method Details

createClient

public void createClient(OpenAIModelParameters parameters)

Initializes the OpenAI client with the provided parameters.

This method is called once during setup. It creates an authenticated OpenAI client using the API key from the parameters.

Specified by:

createClient in interface BaseModelHandler<OpenAIModelParameters,OpenAIModelInput,OpenAIModelResponse>

Parameters:

parameters - the configuration parameters including API key and model name

request

public Iterable<PredictionResult<OpenAIModelInput,OpenAIModelResponse>> request(List<OpenAIModelInput> input)

Performs inference on a batch of inputs using the OpenAI Client.

This method serializes the input batch to JSON string, sends it to OpenAI with structured output requirements, and parses the response into PredictionResult objects that pair each input with its corresponding output.

Specified by:

request in interface BaseModelHandler<OpenAIModelParameters,OpenAIModelInput,OpenAIModelResponse>

Parameters:

input - the list of inputs to process

Returns:

an iterable of model results and input pairs