Fix output_prefix in do() method for ChatGPT Agent

flytekit/extend/backend/base_agent.py

@@ -119,7 +119,9 @@
     name = "Base Sync Agent"
 
     @abstractmethod
-    def do(self, task_template: TaskTemplate, inputs: Optional[LiteralMap], output_prefix: str, **kwargs) -> Resource:
+    def do(
+        self, task_template: TaskTemplate, inputs: Optional[LiteralMap], output_prefix: Optional[str], **kwargs
+    ) -> Resource:
         """
         This is the method that the agent will run.
         """
@@ -247,7 +249,9 @@
 
         agent = AgentRegistry.get_agent(task_template.type, task_template.task_type_version)
 
-        resource = asyncio.run(self._do(agent, task_template, output_prefix, kwargs))
+        resource = asyncio.run(
+            self._do(agent=agent, template=task_template, output_prefix=output_prefix, inputs=kwargs)
+        )
         if resource.phase != TaskExecution.SUCCEEDED:
             raise FlyteUserException(f"Failed to run the task {self.name} with error: {resource.message}")
 
@@ -259,8 +263,8 @@
         self: PythonTask,
         agent: SyncAgentBase,
         template: TaskTemplate,
-        output_prefix: str,
         inputs: Dict[str, Any] = None,
+        output_prefix: Optional[str] = None,
     ) -> Resource:
         try:
             ctx = FlyteContext.current_context()

plugins/flytekit-huggingface/setup.py

@@ -6,7 +6,7 @@
 
 plugin_requires = [
     "flytekit>=1.3.0b2,<2.0.0",
-    "datasets>=2.4.0",
+    "datasets>=2.4.0,<2.19.2",
 ]
 
 __version__ = "0.0.0+develop"

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

@@ -27,6 +27,7 @@ async def do(
         self,
         task_template: TaskTemplate,
         inputs: Optional[LiteralMap] = None,
+        **kwargs,
     ) -> Resource:
         ctx = FlyteContextManager.current_context()
         input_python_value = TypeEngine.literal_map_to_kwargs(ctx, inputs, {"message": str})

plugins/flytekit-openai/tests/chatgpt/test_chatgpt.py

@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from unittest import mock
 
 from flytekitplugins.openai import ChatGPTTask
 
@@ -7,6 +8,14 @@
 from flytekit.models.types import SimpleType
 
 
+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
+
+
 def test_chatgpt_task():
     chatgpt_task = ChatGPTTask(
         name="chatgpt",
@@ -40,3 +49,16 @@ def test_chatgpt_task():
 
     assert chatgpt_task_spec.template.interface.inputs["message"].type.simple == SimpleType.STRING
     assert chatgpt_task_spec.template.interface.outputs["o0"].type.simple == SimpleType.STRING
+
+    with mock.patch("openai.resources.chat.completions.AsyncCompletions.create", new=mock_acreate):
+        chatgpt_task = ChatGPTTask(
+            name="chatgpt",
+            openai_organization="TEST ORGANIZATION ID",
+            chatgpt_config={
+                "model": "gpt-3.5-turbo",
+                "temperature": 0.7,
+            },
+        )
+
+        response = chatgpt_task(message="hi")
+        assert response == "mocked_message"

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     "flyteidl>=1.11.0b1",
     "fsspec>=2023.3.0",
     "gcsfs>=2023.3.0",
-    "googleapis-common-protos>=1.57",
+    "googleapis-common-protos>=1.57,!=1.63.1",
     "grpcio",
     "grpcio-status",
     "importlib-metadata",