Chatgpt Agent

flytekit/__init__.py

@@ -246,6 +246,8 @@
     StructuredDatasetType,
 )
 
+from flytekit.extend.backend.task_executor import SyncAgentBase  # isort:skip. This is for circular import avoidance.
+
 
 def current_context() -> ExecutionParameters:
     """

flytekit/core/external_api_task.py

@@ -0,0 +1,71 @@
+import collections
+import inspect
+from abc import abstractmethod
+from typing import Any, Dict, Optional, TypeVar
+
+from flyteidl.admin.agent_pb2 import CreateTaskResponse
+from typing_extensions import get_type_hints
+
+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 AsyncAgentExecutorMixin
+
+T = TypeVar("T")
+TASK_MODULE = "task_module"
+TASK_NAME = "task_name"
+TASK_CONFIG = "task_config"
+TASK_TYPE = "api_task"
+
+
+class ExternalApiTask(AsyncAgentExecutorMixin, PythonTask):
+    """
+    Base class for all external API tasks. External API tasks are tasks that are designed to run until they receive a
+    response from an external service. When the response is received, the task will complete. External API tasks are
+    designed to be run by the flyte agent.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        config: Optional[T] = None,
+        task_type: str = TASK_TYPE,
+        return_type: Optional[Any] = None,
+        **kwargs,
+    ):
+        type_hints = get_type_hints(self.do, include_extras=True)
+        signature = inspect.signature(self.do)
+        inputs = collections.OrderedDict()
+        outputs = collections.OrderedDict({"o0": return_type}) if return_type else collections.OrderedDict()
+
+        for k, _ in signature.parameters.items():  # type: ignore
+            annotation = type_hints.get(k, None)
+            inputs[k] = annotation
+
+        super().__init__(
+            task_type=task_type,
+            name=name,
+            task_config=config,
+            interface=Interface(inputs=inputs, outputs=outputs),
+            **kwargs,
+        )
+
+        self._task_config = config
+
+    @abstractmethod
+    async def do(self, **kwargs) -> CreateTaskResponse:
+        """
+        Initiate an HTTP request to an external service such as OpenAI or Vertex AI and retrieve the response.
+        """
+        raise NotImplementedError
+
+    def get_custom(self, settings: SerializationSettings = None) -> Dict[str, Any]:
+        cfg = {
+            TASK_MODULE: type(self).__module__,
+            TASK_NAME: type(self).__name__,
+        }
+
+        if self._task_config is not None:
+            cfg[TASK_CONFIG] = self._task_config
+
+        return cfg

flytekit/extend/backend/agent_service.py

@@ -44,7 +44,11 @@
 def agent_exception_handler(func):
     async def wrapper(
         self,
-        request: typing.Union[CreateTaskRequest, GetTaskRequest, DeleteTaskRequest],
+        request: typing.Union[
+            CreateTaskRequest,
+            GetTaskRequest,
+            DeleteTaskRequest,
+        ],
         context: grpc.ServicerContext,
         *args,
         **kwargs,
@@ -92,7 +96,6 @@ async def CreateTask(self, request: CreateTaskRequest, context: grpc.ServicerCon
         tmp = TaskTemplate.from_flyte_idl(request.template)
         inputs = LiteralMap.from_flyte_idl(request.inputs) if request.inputs else None
         agent = AgentRegistry.get_agent(tmp.type)
-
         logger.info(f"{tmp.type} agent start creating the job")
         if agent.asynchronous:
             return await agent.async_create(

flytekit/extend/backend/base_agent.py

@@ -6,7 +6,7 @@
 from abc import ABC
 from collections import OrderedDict
 from functools import partial
-from types import FrameType
+from types import FrameType, coroutine
 
 import grpc
 from flyteidl.admin.agent_pb2 import (
@@ -138,7 +138,8 @@ def convert_to_flyte_state(state: str) -> State:
     Convert the state from the agent to the state in flyte.
     """
     state = state.lower()
-    if state in ["failed", "timedout", "canceled"]:
+    # timedout is the state of Databricks job. https://docs.databricks.com/en/workflows/jobs/jobs-2.0-api.html#runresultstate
+    if state in ["failed", "timeout", "timedout", "canceled"]:
         return RETRYABLE_FAILURE
     elif state in ["done", "succeeded", "success"]:
         return SUCCEEDED
@@ -158,15 +159,27 @@ def get_agent_secret(secret_key: str) -> str:
     return flytekit.current_context().secrets.get(secret_key)
 
 
+def _get_grpc_context() -> grpc.ServicerContext:
+    from unittest.mock import MagicMock
+
+    grpc_ctx = MagicMock(spec=grpc.ServicerContext)
+    return grpc_ctx
+
+
 class AsyncAgentExecutorMixin:
     """
     This mixin class is used to run the agent task locally, and it's only used for local execution.
     Task should inherit from this class if the task can be run in the agent.
+    It can handle asynchronous tasks and synchronous tasks.
+    Asynchronous tasks are for tasks running long, for example running query job.
+    Synchronous tasks are for tasks running quick, for example, you want to execute something really fast, or even retrieving some metadata from a backend service.
     """
 
-    _is_canceled = None
-    _agent = None
-    _entity = None
+    _clean_up_task: coroutine = None
+    _agent: AgentBase = None
+    _entity: PythonTask = None
+    _ctx: FlyteContext = FlyteContext.current_context()
+    _grpc_ctx: grpc.ServicerContext = _get_grpc_context()
 
     def execute(self, **kwargs) -> typing.Any:
         ctx = FlyteContext.current_context()
@@ -180,6 +193,13 @@ def execute(self, **kwargs) -> typing.Any:
         self._agent = AgentRegistry.get_agent(task_template.type)
 
         res = asyncio.run(self._create(task_template, output_prefix, kwargs))
+
+        # If the task is synchronous, the agent will return the output from the resource literals.
+        if res.HasField("resource"):
+            if res.resource.state != SUCCEEDED:
+                raise FlyteUserException(f"Failed to run the task {self._entity.name}")
+            return LiteralMap.from_flyte_idl(res.resource.outputs)
+
         res = asyncio.run(self._get(resource_meta=res.resource_meta))
 
         if res.resource.state != SUCCEEDED:
@@ -198,7 +218,6 @@ async def _create(
         self, task_template: TaskTemplate, output_prefix: str, inputs: typing.Dict[str, typing.Any] = None
     ) -> CreateTaskResponse:
         ctx = FlyteContext.current_context()
-        grpc_ctx = _get_grpc_context()
 
         # Convert python inputs to literals
         literals = inputs or {}
@@ -213,9 +232,9 @@ async def _create(
             task_template = render_task_template(task_template, output_prefix)
 
         if self._agent.asynchronous:
-            res = await self._agent.async_create(grpc_ctx, output_prefix, task_template, inputs)
+            res = await self._agent.async_create(self._grpc_ctx, output_prefix, task_template, literal_map)
         else:
-            res = self._agent.create(grpc_ctx, output_prefix, task_template, inputs)
+            res = self._agent.create(self._grpc_ctx, output_prefix, task_template, literal_map)
 
         signal.signal(signal.SIGINT, partial(self.signal_handler, res.resource_meta))  # type: ignore
         return res
@@ -232,8 +251,8 @@ async def _get(self, resource_meta: bytes) -> GetTaskResponse:
                 time.sleep(1)
                 if self._agent.asynchronous:
                     res = await self._agent.async_get(grpc_ctx, resource_meta)
-                    if self._is_canceled:
-                        await self._is_canceled
+                    if self._clean_up_task:
+                        await self._clean_up_task
                         sys.exit(1)
                 else:
                     res = self._agent.get(grpc_ctx, resource_meta)
@@ -242,12 +261,11 @@ async def _get(self, resource_meta: bytes) -> GetTaskResponse:
         return res
 
     def signal_handler(self, resource_meta: bytes, signum: int, frame: FrameType) -> typing.Any:
-        grpc_ctx = _get_grpc_context()
         if self._agent.asynchronous:
-            if self._is_canceled is None:
-                self._is_canceled = asyncio.create_task(self._agent.async_delete(grpc_ctx, resource_meta))
+            if self._clean_up_task is None:
+                self._clean_up_task = asyncio.create_task(self._agent.async_delete(self._grpc_ctx, resource_meta))
         else:
-            self._agent.delete(grpc_ctx, resource_meta)
+            self._agent.delete(self._grpc_ctx, resource_meta)
             sys.exit(1)
 
 

flytekit/extend/backend/task_executor.py

@@ -0,0 +1,71 @@
+import importlib
+import typing
+from dataclasses import dataclass
+from typing import final
+
+import grpc
+from flyteidl.admin.agent_pb2 import CreateTaskResponse
+
+from flytekit import FlyteContextManager
+from flytekit.core.external_api_task import TASK_CONFIG, TASK_MODULE, TASK_NAME, TASK_TYPE
+from flytekit.core.type_engine import TypeEngine
+from flytekit.extend.backend.base_agent import AgentBase, AgentRegistry
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+T = typing.TypeVar("T")
+
+
+@dataclass
+class IOContext:
+    inputs: LiteralMap
+    output_prefix: str
+
+
+class SyncAgentBase(AgentBase):
+    """
+    SyncAgentBase is an agent responsible for syncrhounous tasks, which are fast and quick.
+
+    This class is meant to be subclassed when implementing plugins that require
+    an external API to perform the task execution. It provides a routing mechanism
+    to direct the task to the appropriate handler based on the task's specifications.
+    """
+
+    def __init__(self):
+        super().__init__(task_type=TASK_TYPE, asynchronous=True)
+
+    @final
+    async def async_create(
+        self,
+        context: grpc.ServicerContext,
+        output_prefix: str,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+    ) -> CreateTaskResponse:
+        return await self.do(context, output_prefix, task_template, inputs)
+
+    async def do(
+        self,
+        context: grpc.ServicerContext,
+        output_prefix: str,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+    ) -> CreateTaskResponse:
+        python_interface_inputs = {
+            name: TypeEngine.guess_python_type(lt.type) for name, lt in task_template.interface.inputs.items()
+        }
+        ctx = FlyteContextManager.current_context()
+
+        native_inputs = {}
+        if inputs:
+            native_inputs = TypeEngine.literal_map_to_kwargs(ctx, inputs, python_interface_inputs)
+
+        meta = task_template.custom
+
+        task_module = importlib.import_module(name=meta[TASK_MODULE])
+        task_def = getattr(task_module, meta[TASK_NAME])
+        config = meta[TASK_CONFIG] if meta.get(TASK_CONFIG) else None
+        return await task_def(TASK_TYPE, config=config).do(**native_inputs)
+
+
+AgentRegistry.register(SyncAgentBase())

flytekit/models/literals.py

@@ -8,11 +8,10 @@
 from flytekit.exceptions import user as _user_exceptions
 from flytekit.models import common as _common
 from flytekit.models.core import types as _core_types
-from flytekit.models.types import Error
+from flytekit.models.types import Error, StructuredDatasetType
 from flytekit.models.types import LiteralType as _LiteralType
 from flytekit.models.types import OutputReference as _OutputReference
 from flytekit.models.types import SchemaType as _SchemaType
-from flytekit.models.types import StructuredDatasetType
 
 
 class RetryStrategy(_common.FlyteIdlEntity):

flytekit/sensor/base_sensor.py

@@ -15,6 +15,7 @@
 SENSOR_MODULE = "sensor_module"
 SENSOR_NAME = "sensor_name"
 SENSOR_CONFIG_PKL = "sensor_config_pkl"
+SENSOR_TYPE = "sensor"
 INPUTS = "inputs"
 
 
@@ -35,7 +36,7 @@ def __init__(
         type_hints = get_type_hints(self.poke, include_extras=True)
         signature = inspect.signature(self.poke)
         inputs = collections.OrderedDict()
-        for k, v in signature.parameters.items():  # type: ignore
+        for k, _ in signature.parameters.items():  # type: ignore
             annotation = type_hints.get(k, None)
             inputs[k] = annotation
 

flytekit/sensor/sensor_engine.py

@@ -19,7 +19,7 @@
 from flytekit.extend.backend.base_agent import AgentBase, AgentRegistry
 from flytekit.models.literals import LiteralMap
 from flytekit.models.task import TaskTemplate
-from flytekit.sensor.base_sensor import INPUTS, SENSOR_CONFIG_PKL, SENSOR_MODULE, SENSOR_NAME
+from flytekit.sensor.base_sensor import INPUTS, SENSOR_CONFIG_PKL, SENSOR_MODULE, SENSOR_NAME, SENSOR_TYPE
 
 T = typing.TypeVar("T")
 
@@ -39,9 +39,11 @@ async def async_create(
             name: TypeEngine.guess_python_type(lt.type) for name, lt in task_template.interface.inputs.items()
         }
         ctx = FlyteContextManager.current_context()
+
         if inputs:
             native_inputs = TypeEngine.literal_map_to_kwargs(ctx, inputs, python_interface_inputs)
             task_template.custom[INPUTS] = native_inputs
+
         return CreateTaskResponse(resource_meta=cloudpickle.dumps(task_template.custom))
 
     async def async_get(self, context: grpc.ServicerContext, resource_meta: bytes) -> GetTaskResponse:
@@ -52,7 +54,7 @@ async def async_get(self, context: grpc.ServicerContext, resource_meta: bytes) -
         sensor_config = jsonpickle.decode(meta[SENSOR_CONFIG_PKL]) if meta.get(SENSOR_CONFIG_PKL) else None
 
         inputs = meta.get(INPUTS, {})
-        cur_state = SUCCEEDED if await sensor_def("sensor", config=sensor_config).poke(**inputs) else RUNNING
+        cur_state = SUCCEEDED if await sensor_def(SENSOR_TYPE, config=sensor_config).poke(**inputs) else RUNNING
         return GetTaskResponse(resource=Resource(state=cur_state, outputs=None))
 
     async def async_delete(self, context: grpc.ServicerContext, resource_meta: bytes) -> DeleteTaskResponse:

plugins/flytekit-aws-athena/flytekitplugins/athena/task.py

@@ -5,6 +5,7 @@
 
 from flytekit.configuration import SerializationSettings
 from flytekit.extend import SQLTask
+from flytekit.extend.backend.base_agent import ASYNC_PLUGIN
 from flytekit.models.presto import PrestoQuery
 from flytekit.types.schema import FlyteSchema
 
@@ -65,6 +66,7 @@ def __init__(
             inputs=inputs,
             outputs=outputs,
             task_type=self._TASK_TYPE,
+            runtime_flavor=ASYNC_PLUGIN,
             **kwargs,
         )
         self._output_schema_type = output_schema_type

plugins/flytekit-aws-batch/flytekitplugins/awsbatch/task.py

@@ -41,7 +41,10 @@ def __init__(self, task_config: AWSBatchConfig, task_function: Callable, **kwarg
         if task_config is None:
             task_config = AWSBatchConfig()
         super(AWSBatchFunctionTask, self).__init__(
-            task_config=task_config, task_type=self._AWS_BATCH_TASK_TYPE, task_function=task_function, **kwargs
+            task_config=task_config,
+            task_type=self._AWS_BATCH_TASK_TYPE,
+            task_function=task_function,
+            **kwargs,
         )
         self._task_config = task_config
 

plugins/flytekit-openai-chatgpt/dev-requirements.in

@@ -0,0 +1 @@
+pytest-asyncio

plugins/flytekit-openai-chatgpt/dev-requirements.txt

@@ -0,0 +1,20 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile dev-requirements.in
+#
+exceptiongroup==1.1.3
+    # via pytest
+iniconfig==2.0.0
+    # via pytest
+packaging==23.2
+    # via pytest
+pluggy==1.3.0
+    # via pytest
+pytest==7.4.2
+    # via pytest-asyncio
+pytest-asyncio==0.21.1
+    # via -r dev-requirements.in
+tomli==2.0.1
+    # via pytest