add openai batch API agent

[samhita-alla]

Tracking issue

Fixes flyteorg/flyte#5274

Why are the changes needed?

The Flyte OpenAI batch API agent can be used to create a batch and poll the status of the OpenAI batch in real-time while remaining dormant. This means there's no need to incur costs when there's no work to be done. OpenAI's 50% model discount, as explained here, further enhances cost savings.

The agent seamlessly integrates with end-to-end Flyte workflows, making it especially useful when downstream tasks depend on batch processing jobs.

What changes were proposed in this pull request?

This PR adds an agent to support OpenAI batch inference. It also adds a JSONL type transformer, allowing users to specify the Iterator[JSON] type, which automatically handles the conversion between iterator and JSONL FlyteFile. The JSONL type can be used when sending input to the OpenAI batch API agent.

from typing import Iterator

from flytekit import workflow, Secret
from flytekit.types.file import JSONLFile
from flytekit.types.iterator import JSON
from flytekitplugins.openai import create_batch, BatchResult


def jsons():
    for x in [
        {
            "custom_id": "request-1",
            "method": "POST",
            "url": "/v1/chat/completions",
            "body": {
                "model": "gpt-3.5-turbo",
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "What is 2+2?"},
                ],
            },
        },
        {
            "custom_id": "request-2",
            "method": "POST",
            "url": "/v1/chat/completions",
            "body": {
                "model": "gpt-3.5-turbo",
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Who won the world series in 2020?"},
                ],
            },
        },
    ]:
        yield x


it_batch = create_batch(
    name="gpt-3.5-turbo",
    openai_organization="your-org",
    secret=Secret(group="openai-secret", key="api-key"),
)

file_batch = create_batch(
    name="gpt-3.5-turbo",
    openai_organization="your-org",
    secret=Secret(group="openai-secret", key="api-key"),
    is_json_iterator=False,
)


@workflow
def json_iterator_wf(json_vals: Iterator[JSON] = jsons()) -> BatchResult:
    return it_batch(jsonl_in=json_vals)


@workflow
def jsonl_wf(jsonl_file: JSONLFile = "data.jsonl") -> BatchResult:
    return batch(jsonl_in=jsonl_file)

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

TODO:

After #2297 PR is merged,

• Run test cases
• Run a real-world example on the Flyte cluster with secrets configured

[kumare3]

This is an awesome start. It's worth a blog!!

[codecov]

Codecov Report

Attention: Patch coverage is 88.75000% with 9 lines in your changes are missing coverage. Please review.

Project coverage is 78.02%. Comparing base (70332db) to head (2152efc).
Report is 3 commits behind head on master.

❗ Current head 2152efc differs from pull request most recent head 7191e39

Please upload reports for the commit 7191e39 to get more accurate results.

Files	Patch %	Lines
flytekit/types/iterator/json_iterator.py	93.10%	3 Missing and 1 partial ⚠️
flytekit/interaction/click_types.py	57.14%	2 Missing and 1 partial ⚠️
flytekit/core/type_engine.py	81.81%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2353      +/-   ##
==========================================
+ Coverage   72.10%   78.02%   +5.91%     
==========================================
  Files         181      182       +1     
  Lines       18395    18470      +75     
  Branches     3601     3609       +8     
==========================================
+ Hits        13264    14411    +1147     
+ Misses       4506     3458    -1048     
+ Partials      625      601      -24     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[pingsutw]

LGTM, just one question about JSONIterator