wrap agent outputs in a literalmap (#3143)

* enclose agent outputs in a literalmap Signed-off-by: Samhita Alla <[email protected]> * add type check Signed-off-by: Samhita Alla <[email protected]> * fix assert check Signed-off-by: Samhita Alla <[email protected]> --------- Signed-off-by: Samhita Alla <[email protected]>
flyteorg · Feb 19, 2025 · 806ff20 · 806ff20
1 parent e19fedc
commit 806ff20
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 19 deletions.
diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_agent.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_agent.py
@@ -122,20 +122,22 @@ async def do(
                 )
             )
             with context_manager.FlyteContextManager.with_context(builder) as new_ctx:
-                outputs = {
-                    "result": TypeEngine.to_literal(
-                        new_ctx,
-                        truncated_result if truncated_result else result,
-                        Annotated[dict, kwtypes(allow_pickle=True)],
-                        TypeEngine.to_literal_type(dict),
-                    ),
-                    "idempotence_token": TypeEngine.to_literal(
-                        new_ctx,
-                        idempotence_token,
-                        str,
-                        TypeEngine.to_literal_type(str),
-                    ),
-                }
+                outputs = LiteralMap(
+                    {
+                        "result": TypeEngine.to_literal(
+                            new_ctx,
+                            truncated_result if truncated_result else result,
+                            Annotated[dict, kwtypes(allow_pickle=True)],
+                            TypeEngine.to_literal_type(dict),
+                        ),
+                        "idempotence_token": TypeEngine.to_literal(
+                            new_ctx,
+                            idempotence_token,
+                            str,
+                            TypeEngine.to_literal_type(str),
+                        ),
+                    }
+                )
 
         return Resource(phase=TaskExecution.SUCCEEDED, outputs=outputs)
 

diff --git a/plugins/flytekit-aws-sagemaker/tests/test_boto3_agent.py b/plugins/flytekit-aws-sagemaker/tests/test_boto3_agent.py
@@ -68,7 +68,7 @@
 @mock.patch(
     "flytekitplugins.awssagemaker_inference.boto3_agent.Boto3AgentMixin._call",
 )
-async def test_agent(mock_boto_call, mock_return_value):
+async def test_agent(mock_boto_call, mock_return_value, request):
     mock_boto_call.return_value = mock_return_value[0]
 
     agent = AgentRegistry.get_agent("boto")
@@ -159,6 +159,9 @@ async def test_agent(mock_boto_call, mock_return_value):
 
     assert resource.phase == TaskExecution.SUCCEEDED
 
+    if request.node.callspec.indices["mock_return_value"] in (0, 1):
+        assert isinstance(resource.outputs, literals.LiteralMap)
+
     if mock_return_value[0][0]:
         outputs = literal_map_string_repr(resource.outputs)
         if "pickle_check" in mock_return_value[0][0]:

diff --git a/plugins/flytekit-openai/flytekitplugins/openai/batch/agent.py b/plugins/flytekit-openai/flytekitplugins/openai/batch/agent.py
@@ -105,7 +105,7 @@ async def get(
         result = retrieved_result.to_dict()
 
         ctx = FlyteContextManager.current_context()
-        outputs = {"result": TypeEngine.to_literal(ctx, result, Dict, TypeEngine.to_literal_type(Dict))}
+        outputs = LiteralMap({"result": TypeEngine.to_literal(ctx, result, Dict, TypeEngine.to_literal_type(Dict))})
 
         return Resource(phase=flyte_phase, outputs=outputs, message=message)
 

diff --git a/plugins/flytekit-openai/tests/openai_batch/test_agent.py b/plugins/flytekit-openai/tests/openai_batch/test_agent.py
@@ -1,9 +1,8 @@
+import json
 from datetime import timedelta
 from unittest import mock
 from unittest.mock import AsyncMock
-import json
-import msgpack
-import base64
+
 import pytest
 from flyteidl.core.execution_pb2 import TaskExecution
 from flytekitplugins.openai.batch.agent import BatchEndpointMetadata
@@ -157,6 +156,7 @@ async def test_openai_batch_agent(mock_retrieve, mock_create, mock_context):
     mock_retrieve.return_value = batch_retrieve_result
     resource = await agent.get(metadata)
     assert resource.phase == TaskExecution.SUCCEEDED
+    assert isinstance(resource.outputs, literals.LiteralMap)
 
     outputs = literal_map_string_repr(resource.outputs)