Fix output_prefix in do() method for ChatGPT Agent (#2457)

Signed-off-by: Future-Outlier <[email protected]>
Co-authored-by: pingsutw <[email protected]>

flytekit/extend/backend/base_agent.py

@@ -119,7 +119,9 @@ class SyncAgentBase(AgentBase):
     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, output_prefix: str, inputs: Optional[LiteralMap] = None, **kwargs
+    ) -> Resource:
         """
         This is the method that the agent will run.
         """
@@ -247,7 +249,9 @@ def execute(self: PythonTask, **kwargs) -> LiteralMap:
 
         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}")
 

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"