Fix completion object parsing

workflowai/core/domain/completion.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Annotated, Any, Literal, Optional, Union
 
 from pydantic import BaseModel, Field
 
@@ -18,22 +18,67 @@ class CompletionUsage(BaseModel):
     model_context_window_size: Optional[int] = None
 
 
+class TextContent(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+class DocumentURL(BaseModel):
+    url: str
+
+
+class DocumentContent(BaseModel):
+    type: Literal["document_url"] = "document_url"
+    source: DocumentURL
+
+
+class ImageURL(BaseModel):
+    url: str
+
+
+class AudioURL(BaseModel):
+    url: str
+
+
+class ImageContent(BaseModel):
+    type: Literal["image_url"] = "image_url"
+    image_url: ImageURL
+
+
+class AudioContent(BaseModel):
+    type: Literal["audio_url"] = "audio_url"
+    audio_url: AudioURL
+
+
+class ToolCallRequest(BaseModel):
+    type: Literal["tool_call_request"] = "tool_call_request"
+    id: Union[str, None] = None
+    tool_name: str
+    tool_input_dict: Union[dict[str, Any], None] = None
+
+
+class ToolCallResult(BaseModel):
+    type: Literal["tool_call_result"] = "tool_call_result"
+    id: Union[str, None] = None
+    tool_name: Union[str, None] = None
+    tool_input_dict: Union[dict[str, Any], None] = None
+    result: Union[Any, None] = None
+    error: Union[str, None] = None
+
+
+MessageContent = Annotated[Union[TextContent, DocumentContent, ImageContent, AudioContent], Field(discriminator="type")]
+
+
 class Message(BaseModel):
     """A message in a completion."""
 
     role: str = ""
-    content: str = ""
+    content: Union[str, MessageContent] = Field(default="")
 
 
 class Completion(BaseModel):
     """A completion from the model."""
 
     messages: list[Message] = Field(default_factory=list)
-    response: Optional[str] = None
+    response: Optional[str] = Field(default=None)
     usage: CompletionUsage = Field(default_factory=CompletionUsage)
-
-
-class CompletionsResponse(BaseModel):
-    """Response from the completions API endpoint."""
-
-    completions: list[Completion]

workflowai/core/domain/completion_test.py

@@ -0,0 +1,124 @@
+import pytest
+from pydantic import ValidationError
+
+from workflowai.core.domain.completion import AudioContent, DocumentContent, ImageContent, Message, TextContent
+
+
+class TestMessage:
+    def test_basic_text(self):
+        # Test basic text message validation
+        json_str = '{"role": "user", "content": "Hello, world!"}'
+        message = Message.model_validate_json(json_str)
+        assert message.role == "user"
+        assert message.content == "Hello, world!"
+
+    def test_with_text_content(self):
+        # Test message with TextContent
+        json_str = """
+        {
+            "role": "assistant",
+            "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"
+
+    def test_with_document_content(self):
+        # Test message with DocumentContent
+        json_str = """
+        {
+            "role": "user",
+            "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"
+
+    def test_with_image_content(self):
+        # Test message with ImageContent
+        json_str = """
+        {
+            "role": "user",
+            "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"
+
+    def test_with_audio_content(self):
+        # Test message with AudioContent
+        json_str = """
+        {
+            "role": "user",
+            "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"
+
+    def test_empty_role(self):
+        # Test message with empty role
+        json_str = '{"role": "", "content": "Test message"}'
+        message = Message.model_validate_json(json_str)
+        assert message.role == ""
+        assert message.content == "Test message"
+
+    def test_missing_role(self):
+        # Test message with missing role
+        json_str = '{"content": "Test message"}'
+        message = Message.model_validate_json(json_str)
+        assert message.role == ""  # Default value
+        assert message.content == "Test message"
+
+    def test_invalid_content_type(self):
+        # Test message with invalid content type
+        json_str = """
+        {
+            "role": "user",
+            "content": {
+                "type": "invalid_type",
+                "text": "This should fail"
+            }
+        }
+        """
+        with pytest.raises(ValidationError):
+            Message.model_validate_json(json_str)
+
+    def test_missing_content(self):
+        # Test message with missing content
+        json_str = '{"role": "user"}'
+        message = Message.model_validate_json(json_str)
+        assert message.role == "user"
+        assert message.content == ""  # Default value
+
+    def test_invalid_json(self):
+        # Test with invalid JSON string
+        json_str = '{"role": "user", "content": "Test message"'  # Missing closing brace
+        with pytest.raises(ValidationError):
+            Message.model_validate_json(json_str)