Chatgpt Agent

flytekit/__init__.py

@@ -245,6 +245,8 @@
     StructuredDatasetType,
 )
 
+from flytekit.extend.backend.task_executor import TaskExecutor  # isort:skip. This is for circular import avoidance.
+
 __version__ = "0.0.0+develop"
 
 

flytekit/clis/sdk_in_container/serve.py

@@ -1,7 +1,10 @@
 from concurrent import futures
 
 import click
-from flyteidl.service.agent_pb2_grpc import add_AsyncAgentServiceServicer_to_server
+from flyteidl.service.agent_pb2_grpc import (
+    add_AsyncAgentServiceServicer_to_server,
+    add_SyncAgentServiceServicer_to_server,
+)
 from grpc import aio
 
 _serve_help = """Start a grpc server for the agent service."""
@@ -42,7 +45,7 @@ def serve(_: click.Context, port, worker, timeout):
 
 async def _start_grpc_server(port: int, worker: int, timeout: int):
     click.secho("Starting up the server to expose the prometheus metrics...", fg="blue")
-    from flytekit.extend.backend.agent_service import AsyncAgentService
+    from flytekit.extend.backend.agent_service import AsyncAgentService, SyncAgentService
 
     try:
         from prometheus_client import start_http_server
@@ -54,6 +57,7 @@ async def _start_grpc_server(port: int, worker: int, timeout: int):
     server = aio.server(futures.ThreadPoolExecutor(max_workers=worker))
 
     add_AsyncAgentServiceServicer_to_server(AsyncAgentService(), server)
+    add_SyncAgentServiceServicer_to_server(SyncAgentService(), server)
 
     server.add_insecure_port(f"[::]:{port}")
     await server.start()

flytekit/core/base_sql_task.py

@@ -27,6 +27,7 @@ def __init__(
         inputs: Optional[Dict[str, Tuple[Type, Any]]] = None,
         metadata: Optional[TaskMetadata] = None,
         outputs: Optional[Dict[str, Type]] = None,
+        is_sync_plugin: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -39,6 +40,7 @@ def __init__(
             interface=Interface(inputs=inputs or {}, outputs=outputs or {}),
             metadata=metadata,
             task_config=task_config,
+            is_sync_plugin=is_sync_plugin,
             **kwargs,
         )
         self._query_template = re.sub(r"\s+", " ", query_template.replace("\n", " ").replace("\t", " ")).strip()

flytekit/core/base_task.py

@@ -26,6 +26,7 @@
 from typing import Any, Coroutine, Dict, Generic, List, Optional, OrderedDict, Tuple, Type, TypeVar, Union, cast
 
 from flyteidl.core import tasks_pb2
+from flyteidl.core.tasks_pb2 import PluginMetadata
 
 from flytekit.configuration import SerializationSettings
 from flytekit.core.context_manager import (
@@ -106,6 +107,7 @@ class TaskMetadata(object):
     retries: int = 0
     timeout: Optional[Union[datetime.timedelta, int]] = None
     pod_template_name: Optional[str] = None
+    plugin_metadata: Optional[PluginMetadata] = None
 
     def __post_init__(self):
         if self.timeout:
@@ -131,7 +133,10 @@ def to_taskmetadata_model(self) -> _task_model.TaskMetadata:
         return _task_model.TaskMetadata(
             discoverable=self.cache,
             runtime=_task_model.RuntimeMetadata(
-                _task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK, __version__, "python"
+                _task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK,
+                __version__,
+                "python",
+                self.plugin_metadata,
             ),
             timeout=self.timeout,
             retries=self.retry_strategy,
@@ -171,12 +176,14 @@ def __init__(
         task_type_version=0,
         security_ctx: Optional[SecurityContext] = None,
         docs: Optional[Documentation] = None,
+        is_sync_plugin: Optional[bool] = None,
         **kwargs,
     ):
         self._task_type = task_type
         self._name = name
         self._interface = interface
-        self._metadata = metadata if metadata else TaskMetadata()
+        plugin_metadata = PluginMetadata(is_sync_plugin=is_sync_plugin) if is_sync_plugin is not None else None
+        self._metadata = metadata if metadata else TaskMetadata(plugin_metadata=plugin_metadata)
         self._task_type_version = task_type_version
         self._security_ctx = security_ctx
         self._docs = docs
@@ -420,6 +427,7 @@ def __init__(
         environment: Optional[Dict[str, str]] = None,
         disable_deck: Optional[bool] = None,
         enable_deck: Optional[bool] = None,
+        is_sync_plugin: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -435,11 +443,13 @@ def __init__(
                 execution of the task. Supplied as a dictionary of key/value pairs
             disable_deck (bool): (deprecated) If true, this task will not output deck html file
             enable_deck (bool): If true, this task will output deck html file
+            is_sync_plugin (bool): If true, plugin task will execute synchronously.
         """
         super().__init__(
             task_type=task_type,
             name=name,
             interface=transform_interface_to_typed_interface(interface),
+            is_sync_plugin=is_sync_plugin,
             **kwargs,
         )
         self._python_interface = interface if interface else Interface()

flytekit/core/external_api_task.py

@@ -0,0 +1,73 @@
+import collections
+import inspect
+from abc import abstractmethod
+from typing import Any, Dict, Optional, TypeVar
+
+import jsonpickle
+from flyteidl.admin.agent_pb2 import DoTaskResponse
+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 SyncAgentExecutorMixin
+
+T = TypeVar("T")
+TASK_MODULE = "task_module"
+TASK_NAME = "task_name"
+TASK_CONFIG_PKL = "task_config_pkl"
+TASK_TYPE = "api_task"
+
+
+class ExternalApiTask(SyncAgentExecutorMixin, 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),
+            is_sync_plugin=True,
+            **kwargs,
+        )
+
+        self._config = config
+
+    @abstractmethod
+    async def do(self, **kwargs) -> DoTaskResponse:
+        """
+        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._config is not None:
+            cfg[TASK_CONFIG_PKL] = jsonpickle.encode(self._config)
+
+        return cfg

flytekit/core/python_auto_container.py

@@ -48,6 +48,7 @@ def __init__(
         pod_template: Optional[PodTemplate] = None,
         pod_template_name: Optional[str] = None,
         accelerator: Optional[BaseAccelerator] = None,
+        is_sync_plugin: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -92,6 +93,7 @@ def __init__(
             name=name,
             task_config=task_config,
             security_ctx=sec_ctx,
+            is_sync_plugin=is_sync_plugin,
             **kwargs,
         )
         self._container_image = container_image

flytekit/core/python_function_task.py

@@ -102,6 +102,7 @@ def __init__(
         ignore_input_vars: Optional[List[str]] = None,
         execution_mode: ExecutionBehavior = ExecutionBehavior.DEFAULT,
         task_resolver: Optional[TaskResolverMixin] = None,
+        is_sync_plugin: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -124,6 +125,7 @@ def __init__(
             interface=mutated_interface,
             task_config=task_config,
             task_resolver=task_resolver,
+            is_sync_plugin=is_sync_plugin,
             **kwargs,
         )
 

flytekit/extend/backend/agent_service.py

@@ -7,10 +7,12 @@
     CreateTaskResponse,
     DeleteTaskRequest,
     DeleteTaskResponse,
+    DoTaskRequest,
+    DoTaskResponse,
     GetTaskRequest,
     GetTaskResponse,
 )
-from flyteidl.service.agent_pb2_grpc import AsyncAgentServiceServicer
+from flyteidl.service.agent_pb2_grpc import AsyncAgentServiceServicer, SyncAgentServiceServicer
 from prometheus_client import Counter, Summary
 
 from flytekit import logger
@@ -23,6 +25,7 @@
 create_operation = "create"
 get_operation = "get"
 delete_operation = "delete"
+do_operation = "do"
 
 # Follow the naming convention. https://prometheus.io/docs/practices/naming/
 request_success_count = Counter(
@@ -44,7 +47,7 @@
 def agent_exception_handler(func):
     async def wrapper(
         self,
-        request: typing.Union[CreateTaskRequest, GetTaskRequest, DeleteTaskRequest],
+        request: typing.Union[CreateTaskRequest, GetTaskRequest, DeleteTaskRequest, DoTaskRequest],
         context: grpc.ServicerContext,
         *args,
         **kwargs,
@@ -60,6 +63,11 @@ async def wrapper(
         elif isinstance(request, DeleteTaskRequest):
             task_type = request.task_type
             operation = delete_operation
+        elif isinstance(request, DoTaskRequest):
+            task_type = request.template.type
+            operation = do_operation
+            if request.inputs:
+                input_literal_size.labels(task_type=task_type).observe(request.inputs.ByteSize())
         else:
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
             context.set_details("Method not implemented!")
@@ -122,3 +130,25 @@ async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerCon
         if agent.asynchronous:
             return await agent.async_delete(context=context, resource_meta=request.resource_meta)
         return await asyncio.get_running_loop().run_in_executor(None, agent.delete, context, request.resource_meta)
+
+
+class SyncAgentService(SyncAgentServiceServicer):
+    @agent_exception_handler
+    async def DoTask(self, request: DoTaskRequest, context: grpc.ServicerContext) -> DoTaskResponse:
+        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 doing the job")
+        if agent.asynchronous:
+            return await agent.async_do(
+                context=context, inputs=inputs, output_prefix=request.output_prefix, task_template=tmp
+            )
+        return await asyncio.get_running_loop().run_in_executor(
+            None,
+            agent.do,
+            context,
+            request.output_prefix,
+            tmp,
+            inputs,
+        )