Chatgpt Agent

flytekit/__init__.py

@@ -237,6 +237,7 @@
 from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar
 from flytekit.models.types import LiteralType
 from flytekit.sensor.sensor_engine import SensorEngine
+from flytekit.dispatcher.dispatcher_engine import DispatcherEngine
 from flytekit.types import directory, file, iterator
 from flytekit.types.structured.structured_dataset import (
     StructuredDataset,

flytekit/dispatcher/__init__.py

@@ -0,0 +1,3 @@
+from .base_dispatcher import BaseDispatcher
+from .chatgpt_dispatcher import ChatGPTDispatcher
+from .dispatcher_engine import DispatcherEngine

flytekit/dispatcher/base_dispatcher.py

@@ -0,0 +1,64 @@
+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 AsyncAgentExecutorMixin
+
+T = TypeVar("T")
+DISPATCHER_MODULE = "dispatcher_module"
+DISPATCHER_NAME = "dispatcher_name"
+DISPATCHER_CONFIG_PKL = "dispatcher_config_pkl"
+INPUTS = "inputs"
+
+
+class BaseDispatcher(AsyncAgentExecutorMixin, PythonTask):
+    """
+    TODO: Write the docstring
+    """
+
+    def __init__(
+        self,
+        name: str,
+        dispatcher_config: Optional[T] = None,
+        task_type: str = "dispatcher",
+        return_type: Optional[T] = 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})
+
+        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=None,
+            interface=Interface(inputs=inputs, outputs=outputs),
+            **kwargs,
+        )
+        self._dispatcher_config = dispatcher_config
+
+    @abstractmethod
+    async def async_do(self, **kwargs) -> DoTaskResponse:
+        raise NotImplementedError
+
+    def get_custom(self, settings: SerializationSettings = None) -> Dict[str, Any]:
+        cfg = {
+            DISPATCHER_MODULE: type(self).__module__,
+            DISPATCHER_NAME: type(self).__name__,
+        }
+        if self._dispatcher_config is not None:
+            cfg[DISPATCHER_CONFIG_PKL] = jsonpickle.encode(self._dispatcher_config)
+        return cfg

flytekit/dispatcher/chatgpt_dispatcher.py

@@ -0,0 +1,70 @@
+import json
+from typing import Any, Dict
+
+import aiohttp
+from flyteidl.admin.agent_pb2 import SUCCEEDED, DoTaskResponse, Resource
+
+from flytekit import FlyteContextManager
+from flytekit.core import utils
+from flytekit.core.type_engine import TypeEngine
+from flytekit.extend.backend.base_agent import get_secret
+from flytekit.models.literals import LiteralMap
+from flytekit.dispatcher.base_dispatcher import BaseDispatcher
+
+
+class ChatGPTDispatcher(BaseDispatcher):
+    """
+    TODO: Write the docstring
+    """
+
+    _openai_organization: str = None
+    _chatgpt_conf: Dict[str, Any] = None
+
+    # TODO,  Add Value Error
+    def __init__(self, name: str, config: Dict[str, Any], **kwargs):
+        super().__init__(name=name, dispatcher_config=config, return_type=str, **kwargs)
+        self._openai_organization = config["openai_organization"]
+        self._chatgpt_conf = config["chatgpt_conf"]
+
+    async def do(
+        self,
+        message: str = None,
+    ) -> DoTaskResponse:
+        self._chatgpt_conf["messages"] = [{"role": "user", "content": message}]
+        openai_url = "https://api.openai.com/v1/chat/completions"
+        data = json.dumps(self._chatgpt_conf)
+
+        message = "TEST SYNC PLUGIN"
+
+        # async with aiohttp.ClientSession() as session:
+        #     async with session.post(
+        #         openai_url, headers=get_header(openai_organization=self._openai_organization), data=data
+        #     ) as resp:
+        #         if resp.status != 200:
+        #             raise Exception(f"Failed to execute chatgpt job with error: {resp.reason}")
+        #         response = await resp.json()
+
+        # message = response["choices"][0]["message"]["content"]
+
+        ctx = FlyteContextManager.current_context()
+        outputs = LiteralMap(
+            {
+                "o0": TypeEngine.to_literal(
+                    ctx,
+                    message,
+                    type(message),
+                    TypeEngine.to_literal_type(type(message)),
+                )
+            }
+        ).to_flyte_idl()
+
+        return DoTaskResponse(resource=Resource(state=SUCCEEDED, outputs=outputs))
+
+
+def get_header(openai_organization: str):
+    token = get_secret(secret_key="OPENAI_ACCESS_TOKEN")
+    return {
+        "OpenAI-Organization": openai_organization,
+        "Authorization": f"Bearer {token}",
+        "content-type": "application/json",
+    }

flytekit/dispatcher/dispatcher_engine.py

@@ -0,0 +1,46 @@
+import importlib
+import typing
+from typing import Optional
+
+import grpc
+import jsonpickle
+from flyteidl.admin.agent_pb2 import DoTaskResponse
+
+from flytekit import FlyteContextManager
+from flytekit.core.type_engine import TypeEngine
+from flytekit.extend.backend.base_agent import AgentRegistry, DispatcherAgent
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+from flytekit.dispatcher.base_dispatcher import INPUTS, DISPATCHER_CONFIG_PKL, DISPATCHER_MODULE, DISPATCHER_NAME
+
+T = typing.TypeVar("T")
+
+
+class DispatcherEngine(DispatcherAgent):
+    async def async_do(
+        self,
+        context: grpc.ServicerContext,
+        task_template: TaskTemplate,
+        inputs: Optional[LiteralMap] = None,
+        output_prefix: Optional[str] = None,
+    ) -> DoTaskResponse:
+        python_interface_inputs = {
+            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
+        else:
+            raise ValueError("Dispatcher needs a input!")
+
+        meta = task_template.custom
+
+        dispatcher_module = importlib.import_module(name=meta[DISPATCHER_MODULE])
+        dispatcher_def = getattr(dispatcher_module, meta[DISPATCHER_NAME])
+        dispatcher_config = jsonpickle.decode(meta[DISPATCHER_CONFIG_PKL]) if meta.get(DISPATCHER_CONFIG_PKL) else None
+        inputs = meta.get(INPUTS, {})
+        return await dispatcher_def("dispatcher", config=dispatcher_config).do(**inputs)
+
+
+AgentRegistry.register(DispatcherEngine())

flytekit/extend/backend/agent_service.py

@@ -6,6 +6,8 @@
     CreateTaskResponse,
     DeleteTaskRequest,
     DeleteTaskResponse,
+    DoTaskRequest,
+    DoTaskResponse,
     GetTaskRequest,
     GetTaskResponse,
 )
@@ -21,6 +23,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(
@@ -114,7 +117,7 @@ async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerCon
                         logger.error(f"failed to run async delete with error {e}")
                         raise
                 try:
-                    res = asyncio.to_thread(agent.delete, context=context, resource_meta=request.resource_meta)
+                    res = await asyncio.to_thread(agent.delete, context=context, resource_meta=request.resource_meta)
                     request_success_count.labels(task_type=request.task_type, operation=delete_operation).inc()
                     return res
                 except Exception as e:
@@ -124,3 +127,36 @@ async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerCon
             context.set_code(grpc.StatusCode.INTERNAL)
             context.set_details(f"failed to delete task with error {e}")
             request_failure_count.labels(task_type=request.task_type, operation=delete_operation).inc()
+
+    async def DoTask(self, request: DoTaskRequest, context: grpc.ServicerContext) -> DoTaskResponse:
+        try:
+            with request_latency.labels(task_type=request.template.type, operation=do_operation).time():
+                tmp = TaskTemplate.from_flyte_idl(request.template)
+                inputs = LiteralMap.from_flyte_idl(request.inputs) if request.inputs else None
+                input_literal_size.labels(task_type=tmp.type).observe(request.inputs.ByteSize())
+                agent = AgentRegistry.get_agent(tmp.type)
+                logger.info(f"{agent.task_type} agent start doing the job")
+                if agent.asynchronous:
+                    try:
+                        res = await agent.async_do(
+                            context=context, inputs=inputs, task_template=tmp
+                        )
+                        request_success_count.labels(task_type=tmp.type, operation=do_operation).inc()
+                        return res
+                    except Exception as e:
+                        logger.error(f"failed to run async do with error {e}")
+                        raise e
+                try:
+                    res = await asyncio.to_thread(
+                        agent.do, context=context, inputs=inputs, task_template=tmp
+                    )
+                    request_success_count.labels(task_type=tmp.type, operation=do_operation).inc()
+                    return res
+                except Exception as e:
+                    logger.error(f"failed to run sync do with error {e}")
+                    raise
+        except Exception as e:
+            logger.error(f"failed to do task with error {e}")
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(f"failed to do task with error {e}")
+            request_failure_count.labels(task_type=request.task_type, operation=do_operation).inc()

flytekit/extend/backend/base_agent.py

@@ -16,14 +16,19 @@
     SUCCEEDED,
     CreateTaskResponse,
     DeleteTaskResponse,
+    DoTaskResponse,
     GetTaskResponse,
+    Resource,
     State,
 )
+from flyteidl.core import literals_pb2
 from flyteidl.core.tasks_pb2 import TaskTemplate
 from rich.progress import Progress
 
+import flytekit
 from flytekit import FlyteContext, logger
 from flytekit.configuration import ImageConfig, SerializationSettings
+from flytekit.core import utils
 from flytekit.core.base_task import PythonTask
 from flytekit.core.type_engine import TypeEngine
 from flytekit.models.literals import LiteralMap
@@ -79,6 +84,18 @@ def delete(self, context: grpc.ServicerContext, resource_meta: bytes) -> DeleteT
         """
         raise NotImplementedError
 
+    def do(
+        self,
+        context: grpc.ServicerContext,
+        output_prefix: str,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+    ) -> DoTaskResponse:
+        """
+        Return the result of executing a task. It should return error code if the task creation failed.
+        """
+        raise NotImplementedError
+
     async def async_create(
         self,
         context: grpc.ServicerContext,
@@ -105,6 +122,45 @@ async def async_delete(self, context: grpc.ServicerContext, resource_meta: bytes
         """
         raise NotImplementedError
 
+    async def async_do(
+        self,
+        context: grpc.ServicerContext,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+        output_prefix: typing.Optional[str] = None,
+    ) -> DoTaskResponse:
+        """
+        Return the result of executing a task. It should return error code if the task creation failed.
+        """
+        raise NotImplementedError
+
+
+class DispatcherAgent(AgentBase):
+    def __init__(self):
+        super().__init__(task_type="dispatcher")
+
+    async def _do(
+        self,
+        entity: PythonTask,
+        task_template: TaskTemplate,
+        agent: AgentBase,
+        inputs: typing.Dict[str, typing.Any] = None,
+    ) -> DoTaskResponse:
+        ctx = FlyteContext.current_context()
+        grpc_ctx = _get_grpc_context()
+
+        literals = {}
+        for k, v in inputs.items():
+            literals[k] = TypeEngine.to_literal(ctx, v, type(v), entity.interface.inputs[k].type)
+        inputs = LiteralMap(literals) if literals else None
+        output_prefix = ctx.file_access.get_random_local_directory()
+
+        progress = Progress(transient=True)
+        task = progress.add_task(f"[cyan]Running Task {entity.name}...", total=None)
+        with progress:
+            progress.start_task(task)
+            return await agent.async_do(context=grpc_ctx, output_prefix=output_prefix, task_template=task_template, inputs=inputs)
+
 
 class AgentRegistry(object):
     """
@@ -149,6 +205,10 @@ def is_terminal_state(state: State) -> bool:
     return state in [SUCCEEDED, RETRYABLE_FAILURE, PERMANENT_FAILURE]
 
 
+def get_secret(secret_key: str) -> str:
+    return flytekit.current_context().secrets.get("flyteagent", secret_key)
+
+
 class AsyncAgentExecutorMixin:
     """
     This mixin class is used to run the agent task locally, and it's only used for local execution.
@@ -166,8 +226,11 @@ def execute(self, **kwargs) -> typing.Any:
         task_template = get_serializable(OrderedDict(), SerializationSettings(ImageConfig()), self._entity).template
         self._agent = AgentRegistry.get_agent(task_template.type)
 
-        res = asyncio.run(self._create(task_template, kwargs))
-        res = asyncio.run(self._get(resource_meta=res.resource_meta))
+        if isinstance(self._agent, DispatcherAgent):
+            res = asyncio.run(self._agent._do(self._entity, task_template, self._agent, kwargs))
+        else:
+            res = asyncio.run(self._create(task_template, kwargs))
+            res = asyncio.run(self._get(resource_meta=res.resource_meta))
 
         if res.resource.state != SUCCEEDED:
             raise Exception(f"Failed to run the task {self._entity.name}")
@@ -226,7 +289,7 @@ def signal_handler(self, resource_meta: bytes, signum: int, frame: FrameType) ->
             sys.exit(1)
 
 
-def _get_grpc_context():
+def _get_grpc_context() -> grpc.ServicerContext:
     from unittest.mock import MagicMock
 
     grpc_ctx = MagicMock(spec=grpc.ServicerContext)

flytekit/sensor/base_sensor.py

@@ -35,7 +35,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