fix: completion parsing

Author: guillaq

examples/01_basic_agent.py

@@ -88,10 +88,10 @@ async def main():
         return
 
     # Example 2: Using Tokyo
-    print("\nExample 2: Using Tokyo")
-    print("-" * 50)
-    run = await get_capital_info.run(CityInput(city="Tokyo"))
-    print(run)
+    # print("\nExample 2: Using Tokyo")
+    # print("-" * 50)
+    # run = await get_capital_info.run(CityInput(city="Tokyo"))
+    # print(run)
 
     # Fetch and display completions for the Tokyo example
     await display_completions(run)

tests/fixtures/completions2.json

@@ -0,0 +1,36 @@
+{
+  "completions": [
+    {
+      "messages": [
+        {
+          "role": "assistant",
+          "content": [
+            {
+              "type": "text",
+              "text": "<instructions>\\nFind the capital city of the country where the input city is located.\\n\\nGuidelines:\\n1. First identify the country where the input city is located\\n2. Then provide the capital city of that country\\n3. Include an interesting historical or cultural fact about the capital\\n4. Be accurate and precise with geographical information\\n5. If the input city is itself the capital, still provide the information\\n</instructions>\\n\\nInput will be provided in the user message using a JSON following the schema:\\n```json\\n{\\n  \"description\": \"Input model for the city-to-capital agent.\",\\n  \"properties\": {\\n    \"city\": {\\n      \"description\": \"The name of the city for which to find the country\\'s capital\",\\n      \"examples\": [\\n        \"Paris\",\\n        \"New York\",\\n        \"Tokyo\"\\n      ],\\n      \"type\": \"string\"\\n    }\\n  },\\n  \"required\": [\\n    \"city\"\\n  ],\\n  \"type\": \"object\"\\n}\\n```\\n\\nReturn a single JSON object enforcing the following schema:\\n```json\\n{\\n  \"description\": \"Output model containing information about the capital city.\",\\n  \"properties\": {\\n    \"country\": {\\n      \"description\": \"The country where the input city is located\",\\n      \"examples\": [\\n        \"France\",\\n        \"United States\",\\n        \"Japan\"\\n      ],\\n      \"type\": \"string\"\\n    },\\n    \"capital\": {\\n      \"description\": \"The capital city of the country\",\\n      \"examples\": [\\n        \"Paris\",\\n        \"Washington D.C.\",\\n        \"Tokyo\"\\n      ],\\n      \"type\": \"string\"\\n    },\\n    \"fun_fact\": {\\n      \"description\": \"An interesting fact about the capital city\",\\n      \"examples\": [\\n        \"Paris has been the capital of France since 508 CE\"\\n      ],\\n      \"type\": \"string\"\\n    }\\n  },\\n  \"required\": [\\n    \"capital\",\\n    \"country\",\\n    \"fun_fact\"\\n  ],\\n  \"type\": \"object\"\\n}\\n```"
+            }
+          ]
+        },
+        {
+          "role": "user",
+          "content": [
+            {
+              "type": "text",
+              "text": "Input is:\\n```json\\n{\\n  \"city\": \"Paris\"\\n}\\n```"
+            }
+          ]
+        }
+      ],
+      "response": "{\\n  \"country\": \"France\",\\n  \"capital\": \"Paris\",\\n  \"fun_fact\": \"Paris became the capital of France in 508 CE when Clovis I, the first King of the Franks, made the city his seat of power. It\\'s one of the oldest capital cities in Europe and was originally called Lutetia during the Roman period.\"\\n}",
+      "usage": {
+        "completion_token_count": 88,
+        "completion_cost_usd": 0.00132,
+        "prompt_token_count": 516,
+        "prompt_cost_usd": 0.0015480000000000001,
+        "model_context_window_size": 200000
+      },
+      "duration_seconds": 4.13,
+      "provider": "anthropic"
+    }
+  ]
+}

workflowai/core/client/agent_test.py

@@ -20,7 +20,7 @@
 from workflowai.core.client.client import (
     WorkflowAI,
 )
-from workflowai.core.domain.completion import Completion, CompletionUsage, Message
+from workflowai.core.domain.completion import Completion, CompletionUsage, Message, TextContent
 from workflowai.core.domain.errors import MaxTurnsReachedError, WorkflowAIError
 from workflowai.core.domain.run import Run
 from workflowai.core.domain.tool_call import ToolCallRequest
@@ -1131,6 +1131,24 @@ async def test_fetch_completions(self, agent: Agent[HelloTaskInput, HelloTaskOut
             ),
         ]
 
+    async def test_fetch_completions_content_list(
+        self,
+        agent: Agent[HelloTaskInput, HelloTaskOutput],
+        httpx_mock: HTTPXMock,
+    ):
+        """Test that fetch_completions correctly fetches and returns completions
+        when the message content is a list of objects"""
+        # Mock the HTTP response instead of the API client method
+        httpx_mock.add_response(
+            url="http://localhost:8000/v1/_/agents/123/runs/1/completions",
+            json=fixtures_json("completions2.json"),
+        )
+
+        completions = await agent.fetch_completions("1")
+        assert len(completions) == 1
+        assert isinstance(completions[0].messages[0].content, list)
+        assert isinstance(completions[0].messages[0].content[0], TextContent)
+
 
 class TestStream:
     async def test_stream(self, httpx_mock: HTTPXMock, agent: Agent[HelloTaskInput, HelloTaskOutput]):

workflowai/core/domain/completion.py

@@ -73,7 +73,7 @@ class Message(BaseModel):
     """A message in a completion."""
 
     role: str = ""
-    content: Union[str, MessageContent] = Field(default="")
+    content: Union[str, list[MessageContent]] = Field(default="")
 
 
 class Completion(BaseModel):

workflowai/core/domain/completion_test.py

@@ -17,70 +17,74 @@ def test_with_text_content(self):
         json_str = """
         {
             "role": "assistant",
-            "content": {
+            "content": [{
                 "type": "text",
                 "text": "This is a test message"
-            }
+            }]
         }
         """
         message = Message.model_validate_json(json_str)
         assert message.role == "assistant"
-        assert isinstance(message.content, TextContent)
-        assert message.content.text == "This is a test message"
+        assert isinstance(message.content, list)
+        assert isinstance(message.content[0], TextContent)
+        assert message.content[0].text == "This is a test message"
 
     def test_with_document_content(self):
         # Test message with DocumentContent
         json_str = """
         {
             "role": "user",
-            "content": {
+            "content": [{
                 "type": "document_url",
                 "source": {
                     "url": "https://example.com/doc.pdf"
                 }
-            }
+            }]
         }
         """
         message = Message.model_validate_json(json_str)
         assert message.role == "user"
-        assert isinstance(message.content, DocumentContent)
-        assert message.content.source.url == "https://example.com/doc.pdf"
+        assert isinstance(message.content, list)
+        assert isinstance(message.content[0], DocumentContent)
+        assert message.content[0].source.url == "https://example.com/doc.pdf"
 
     def test_with_image_content(self):
         # Test message with ImageContent
         json_str = """
         {
             "role": "user",
-            "content": {
+            "content": [{
                 "type": "image_url",
                 "image_url": {
                     "url": "https://example.com/image.jpg"
                 }
-            }
+            }]
         }
         """
         message = Message.model_validate_json(json_str)
         assert message.role == "user"
-        assert isinstance(message.content, ImageContent)
-        assert message.content.image_url.url == "https://example.com/image.jpg"
+        assert isinstance(message.content, list)
+        assert isinstance(message.content[0], ImageContent)
+        assert message.content[0].image_url.url == "https://example.com/image.jpg"
 
     def test_with_audio_content(self):
         # Test message with AudioContent
         json_str = """
         {
             "role": "user",
-            "content": {
+            "content": [{
                 "type": "audio_url",
                 "audio_url": {
                     "url": "https://example.com/audio.mp3"
                 }
-            }
+            }]
         }
         """
         message = Message.model_validate_json(json_str)
         assert message.role == "user"
-        assert isinstance(message.content, AudioContent)
-        assert message.content.audio_url.url == "https://example.com/audio.mp3"
+        assert isinstance(message.content, list)
+        assert isinstance(message.content[0], AudioContent)
+        assert message.content[0].audio_url.url == "https://example.com/audio.mp3"
 
     def test_empty_role(self):
         # Test message with empty role