ChatGPT Agent

.github/workflows/pythonbuild.yml

@@ -239,6 +239,7 @@ jobs:
           # onnx-tensorflow needs a version of tensorflow that does not work with protobuf>4.
           # The issue is being tracked on the tensorflow side in https://github.com/tensorflow/tensorflow/issues/53234#issuecomment-1330111693
           # flytekit-onnx-tensorflow
+          - flytekit-openai
           - flytekit-pandera
           - flytekit-papermill
           - flytekit-polars

Dockerfile.agent

@@ -11,6 +11,7 @@ RUN pip install prometheus-client grpcio-health-checking
 RUN pip install --no-cache-dir -U flytekit==$VERSION \
   flytekitplugins-airflow==$VERSION \
   flytekitplugins-bigquery==$VERSION \
+  flytekitplugins-chatgpt==$VERSION \
   flytekitplugins-snowflake==$VERSION \
   && apt-get clean autoclean \
   && apt-get autoremove --yes \

flytekit/extend/backend/base_agent.py

@@ -251,9 +251,12 @@
         return resource.outputs
 
     async def _do(self: T, agent: SyncAgentBase, template: TaskTemplate, inputs: Dict[str, Any] = None) -> Resource:
-        ctx = FlyteContext.current_context()
-        literal_map = TypeEngine.dict_to_literal_map(ctx, inputs or {}, self.get_input_types())
-        return await mirror_async_methods(agent.do, task_template=template, inputs=literal_map)
+        try:
+            ctx = FlyteContext.current_context()
+            literal_map = TypeEngine.dict_to_literal_map(ctx, inputs or {}, self.get_input_types())
+            return await mirror_async_methods(agent.do, task_template=template, inputs=literal_map)
+        except Exception as error_message:
+            raise FlyteUserException(f"Failed to run the task {self.name} with error: {error_message}")
 
 
 class AsyncAgentExecutorMixin:

plugins/flytekit-openai/README.md

@@ -0,0 +1,44 @@
+# Flytekit ChatGPT Plugin
+ChatGPT plugin allows you to run ChatGPT tasks in the Flyte workflow without changing any code.
+
+## Example
+```python
+from flytekit import task, workflow
+from flytekitplugins.chatgpt import ChatGPTTask, ChatGPTConfig
+
+chatgpt_small_job = ChatGPTTask(
+    name="chatgpt gpt-3.5-turbo",
+    openai_organization="org-NayNG68kGnVXMJ8Ak4PMgQv7",
+    chatgpt_config={
+            "model": "gpt-3.5-turbo",
+            "temperature": 0.7,
+    },
+)
+
+chatgpt_big_job = ChatGPTTask(
+    name="chatgpt gpt-4",
+    openai_organization="org-NayNG68kGnVXMJ8Ak4PMgQv7",
+    chatgpt_config={
+            "model": "gpt-4",
+            "temperature": 0.7,
+    },
+)
+
+
+@workflow
+def wf(message: str) -> str:
+    message = chatgpt_small_job(message=message)
+    message = chatgpt_big_job(message=message)
+    return message
+
+
+if __name__ == "__main__":
+    print(wf(message="hi"))
+```
+
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-chatgpt
+```

plugins/flytekit-openai/flytekitplugins/chatgpt/__init__.py

@@ -0,0 +1,12 @@
+"""
+.. currentmodule:: flytekitplugins.chatgpt
+This package contains things that are useful when extending Flytekit.
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+   ChatGPTAgent
+   ChatGPTTask
+"""
+
+from .agent import ChatGPTAgent
+from .task import ChatGPTTask

plugins/flytekit-openai/flytekitplugins/chatgpt/agent.py

@@ -0,0 +1,52 @@
+import asyncio
+import logging
+from typing import Optional
+
+from flyteidl.core.execution_pb2 import TaskExecution
+
+from flytekit import FlyteContextManager, lazy_module
+from flytekit.core.type_engine import TypeEngine
+from flytekit.extend.backend.base_agent import AgentRegistry, Resource, SyncAgentBase
+from flytekit.extend.backend.utils import get_agent_secret
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+openai = lazy_module("openai")
+
+TIMEOUT_SECONDS = 10
+OPENAI_API_KEY = "FLYTE_OPENAI_API_KEY"
+
+
+class ChatGPTAgent(SyncAgentBase):
+    name = "ChatGPT Agent"
+
+    def __init__(self):
+        super().__init__(task_type_name="chatgpt")
+
+    async def do(
+        self,
+        task_template: TaskTemplate,
+        inputs: Optional[LiteralMap] = None,
+    ) -> Resource:
+        ctx = FlyteContextManager.current_context()
+        input_python_value = TypeEngine.literal_map_to_kwargs(ctx, inputs, {"message": str})
+        message = input_python_value["message"]
+
+        custom = task_template.custom
+        custom["chatgpt_config"]["messages"] = [{"role": "user", "content": message}]
+        client = openai.AsyncOpenAI(
+            organization=custom["openai_organization"],
+            api_key=get_agent_secret(secret_key=OPENAI_API_KEY),
+        )
+
+        logger = logging.getLogger("httpx")
+        logger.setLevel(logging.WARNING)
+
+        completion = await asyncio.wait_for(client.chat.completions.create(**custom["chatgpt_config"]), TIMEOUT_SECONDS)
+        message = completion.choices[0].message.content
+        outputs = {"o0": message}
+
+        return Resource(phase=TaskExecution.SUCCEEDED, outputs=outputs)
+
+
+AgentRegistry.register(ChatGPTAgent())

plugins/flytekit-openai/flytekitplugins/chatgpt/task.py

@@ -0,0 +1,44 @@
+from typing import Any, Dict
+
+from flytekit.configuration import SerializationSettings
+from flytekit.core.base_task import PythonTask
+from flytekit.core.interface import Interface
+from flytekit.extend.backend.base_agent import SyncAgentExecutorMixin
+
+
+class ChatGPTTask(SyncAgentExecutorMixin, PythonTask):
+    """
+    This is the simplest form of a ChatGPT Task, you can define the model and the input you want.
+    """
+
+    _TASK_TYPE = "chatgpt"
+
+    def __init__(self, name: str, openai_organization: str, chatgpt_config: Dict[str, Any], **kwargs):
+        """
+        Args:
+            name: Name of this task, should be unique in the project
+            openai_organization: OpenAI Organization. String can be found here. https://platform.openai.com/docs/api-reference/organization-optional
+            chatgpt_config: ChatGPT job configuration. Config structure can be found here. https://platform.openai.com/docs/api-reference/completions/create
+        """
+
+        if "model" not in chatgpt_config:
+            raise ValueError("The 'model' configuration variable is required in chatgpt_config")
+
+        task_config = {"openai_organization": openai_organization, "chatgpt_config": chatgpt_config}
+
+        inputs = {"message": str}
+        outputs = {"o0": str}
+
+        super().__init__(
+            task_type=self._TASK_TYPE,
+            name=name,
+            task_config=task_config,
+            interface=Interface(inputs=inputs, outputs=outputs),
+            **kwargs,
+        )
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        return {
+            "openai_organization": self.task_config["openai_organization"],
+            "chatgpt_config": self.task_config["chatgpt_config"],
+        }

plugins/flytekit-openai/setup.py

@@ -0,0 +1,38 @@
+from setuptools import setup
+
+PLUGIN_NAME = "chatgpt"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit>1.10.7", "openai>=1.12.0", "flyteidl>=1.11.0b0"]
+
+__version__ = "0.0.0+develop"
+
+setup(
+    name=microlib_name,
+    version=__version__,
+    author="flyteorg",
+    author_email="[email protected]",
+    description="This package holds the ChatGPT plugins for flytekit",
+    namespace_packages=["flytekitplugins"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    install_requires=plugin_requires,
+    license="apache2",
+    python_requires=">=3.8",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
+)

plugins/flytekit-openai/tests/test_agent.py

@@ -0,0 +1,69 @@
+from datetime import timedelta
+from unittest import mock
+
+import pytest
+from flyteidl.core.execution_pb2 import TaskExecution
+
+from flytekit.extend.backend.base_agent import AgentRegistry
+from flytekit.interfaces.cli_identifiers import Identifier
+from flytekit.models import literals
+from flytekit.models.core.identifier import ResourceType
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import RuntimeMetadata, TaskMetadata, TaskTemplate
+
+
+async def mock_acreate(*args, **kwargs) -> str:
+    mock_response = mock.MagicMock()
+    mock_choice = mock.MagicMock()
+    mock_choice.message.content = "mocked_message"
+    mock_response.choices = [mock_choice]
+    return mock_response
+
+
+@pytest.mark.asyncio
+async def test_chatgpt_agent():
+    agent = AgentRegistry.get_agent("chatgpt")
+    task_id = Identifier(
+        resource_type=ResourceType.TASK, project="project", domain="domain", name="name", version="version"
+    )
+    task_config = {
+        "openai_organization": "test-openai-orgnization-id",
+        "chatgpt_config": {"model": "gpt-3.5-turbo", "temperature": 0.7},
+    }
+    task_metadata = TaskMetadata(
+        True,
+        RuntimeMetadata(RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"),
+        timedelta(days=1),
+        literals.RetryStrategy(3),
+        True,
+        "0.1.1b0",
+        "This is deprecated!",
+        True,
+        "A",
+    )
+    tmp = TaskTemplate(
+        id=task_id,
+        custom=task_config,
+        metadata=task_metadata,
+        interface=None,
+        type="chatgpt",
+    )
+
+    task_inputs = LiteralMap(
+        {
+            "message": literals.Literal(
+                scalar=literals.Scalar(primitive=literals.Primitive(string_value="Test ChatGPT Plugin"))
+            ),
+        },
+    )
+    message = "mocked_message"
+    mocked_token = "mocked_openai_api_key"
+    mocked_context = mock.patch("flytekit.current_context", autospec=True).start()
+    mocked_context.return_value.secrets.get.return_value = mocked_token
+
+    with mock.patch("openai.resources.chat.completions.AsyncCompletions.create", new=mock_acreate):
+        # Directly await the coroutine without using asyncio.run
+        response = await agent.do(tmp, task_inputs)
+
+    assert response.phase == TaskExecution.SUCCEEDED
+    assert response.outputs == {"o0": message}

plugins/flytekit-openai/tests/test_chatgpt.py

@@ -0,0 +1,42 @@
+from collections import OrderedDict
+
+from flytekitplugins.chatgpt import ChatGPTTask
+
+from flytekit.configuration import Image, ImageConfig, SerializationSettings
+from flytekit.extend import get_serializable
+from flytekit.models.types import SimpleType
+
+
+def test_chatgpt_task():
+    chatgpt_task = ChatGPTTask(
+        name="chatgpt",
+        openai_organization="TEST ORGANIZATION ID",
+        chatgpt_config={
+            "model": "gpt-3.5-turbo",
+            "temperature": 0.7,
+        },
+    )
+
+    assert len(chatgpt_task.interface.inputs) == 1
+    assert len(chatgpt_task.interface.outputs) == 1
+
+    default_img = Image(name="default", fqn="test", tag="tag")
+    serialization_settings = SerializationSettings(
+        project="proj",
+        domain="dom",
+        version="123",
+        image_config=ImageConfig(default_image=default_img, images=[default_img]),
+        env={},
+    )
+
+    chatgpt_task_spec = get_serializable(OrderedDict(), serialization_settings, chatgpt_task)
+    custom = chatgpt_task_spec.template.custom
+    assert custom["openai_organization"] == "TEST ORGANIZATION ID"
+    assert custom["chatgpt_config"]["model"] == "gpt-3.5-turbo"
+    assert custom["chatgpt_config"]["temperature"] == 0.7
+
+    assert len(chatgpt_task_spec.template.interface.inputs) == 1
+    assert len(chatgpt_task_spec.template.interface.outputs) == 1
+
+    assert chatgpt_task_spec.template.interface.inputs["message"].type.simple == SimpleType.STRING
+    assert chatgpt_task_spec.template.interface.outputs["o0"].type.simple == SimpleType.STRING