fix: add type alias

Author: rushilpatel0

src/codegen/agents/data.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Union
 from dataclasses import dataclass, field
 from datetime import datetime
 
@@ -49,4 +49,6 @@ class FunctionMessageData(BaseMessage):
 @dataclass
 class UnknownMessage(BaseMessage):
     """Represents an unknown message type."""
-    type: str = field(default="unknown") 
+    type: str = field(default="unknown") 
+
+type AgentRunMessage = Union[UserMessage, SystemMessageData, AssistantMessage, ToolMessageData, FunctionMessageData, UnknownMessage]

src/codegen/agents/loggers.py

@@ -1,13 +1,11 @@
-import json
-from typing import Dict, List, Any, Union, Protocol
-from dataclasses import asdict
-from .data import BaseMessage
+from typing import Protocol
+from .data import AgentRunMessage
 
 # Define the interface for ExternalLogger
 class ExternalLogger(Protocol):
     """Protocol defining the interface for external loggers."""
 
-    def log(self, data: Union[Dict[str, Any], BaseMessage]) -> None:
+    def log(self, data: AgentRunMessage) -> None:
         """
         Log structured data to an external system.