fix: literal type

Author: rushilpatel0

src/codegen/agents/data.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Literal, Optional, Union
 from dataclasses import dataclass, field
 from datetime import datetime
 
@@ -13,12 +13,12 @@ class BaseMessage:
 @dataclass
 class UserMessage(BaseMessage):
     """Represents a message from the user."""
-    type: str = field(default="user")
+    type: Literal["user"] = field(default="user")
 
 @dataclass
 class SystemMessageData(BaseMessage):
     """Represents a system message."""
-    type: str = field(default="system")
+    type: Literal["system"] = field(default="system")
 
 @dataclass
 class ToolCall:
@@ -30,25 +30,25 @@ class ToolCall:
 @dataclass
 class AssistantMessage(BaseMessage):
     """Represents a message from the assistant."""
-    type: str = field(default="assistant")
+    type: Literal["assistant"] = field(default="assistant")
     tool_calls: List[ToolCall] = field(default_factory=list)
 
 @dataclass
 class ToolMessageData(BaseMessage):
     """Represents a tool response message."""
-    type: str = field(default="tool")
+    type: Literal["tool"] = field(default="tool")
     tool_name: Optional[str] = None
     tool_response: Optional[str] = None
     tool_id: Optional[str] = None
 
 @dataclass
 class FunctionMessageData(BaseMessage):
     """Represents a function message."""
-    type: str = field(default="function")
+    type: Literal["function"] = field(default="function")
 
 @dataclass
 class UnknownMessage(BaseMessage):
     """Represents an unknown message type."""
-    type: str = field(default="unknown") 
+    type: Literal["unknown"] = field(default="unknown") 
 
 type AgentRunMessage = Union[UserMessage, SystemMessageData, AssistantMessage, ToolMessageData, FunctionMessageData, UnknownMessage]