chore: cleanup

Author: damassi

src/agent_chat_cli/utils/__init__.py

@@ -1,13 +1,15 @@
 """Utility modules for Agent Chat CLI."""
 
-from agent_chat_cli.utils.agent_loop import AgentLoop, AgentMessage, MessageType
+from agent_chat_cli.utils.agent_loop import AgentLoop, AgentMessage
+from agent_chat_cli.utils.enums import AgentMessageType, ContentType
 from agent_chat_cli.utils.tool_info import ToolInfo, get_tool_info
 from agent_chat_cli.utils.format_tool_input import format_tool_input
 
 __all__ = [
     "AgentLoop",
     "AgentMessage",
-    "MessageType",
+    "AgentMessageType",
+    "ContentType",
     "ToolInfo",
     "get_tool_info",
     "format_tool_input",

src/agent_chat_cli/utils/agent_loop.py

@@ -1,5 +1,4 @@
 import asyncio
-from enum import Enum
 from typing import Callable, Awaitable, Any
 from dataclasses import dataclass
 
@@ -10,24 +9,12 @@
 from claude_agent_sdk.types import AssistantMessage, TextBlock, ToolUseBlock
 
 from agent_chat_cli.utils.config import load_config
-
-
-class MessageType(Enum):
-    ASSISTANT = "assistant"
-    INIT = "init"
-    RESULT = "result"
-    STREAM_EVENT = "stream_event"
-    SYSTEM = "system"
-
-
-class ContentType(Enum):
-    TEXT = "text"
-    TOOL_USE = "tool_use"
+from agent_chat_cli.utils.enums import AgentMessageType, ContentType
 
 
 @dataclass
 class AgentMessage:
-    type: MessageType
+    type: AgentMessageType
     data: Any
 
 
@@ -59,22 +46,22 @@ async def start(self) -> None:
             async for message in self.client.receive_response():
                 await self._handle_message(message)
 
-            await self.on_message(AgentMessage(type=MessageType.RESULT, data=None))
+            await self.on_message(AgentMessage(type=AgentMessageType.RESULT, data=None))
 
     async def _handle_message(self, message: Any) -> None:
         if hasattr(message, "event"):
             event = message.event  # type: ignore[attr-defined]
 
-            if event.get("type") == "content_block_delta":
+            if event.get("type") == ContentType.CONTENT_BLOCK_DELTA.value:
                 delta = event.get("delta", {})
 
-                if delta.get("type") == "text_delta":
+                if delta.get("type") == ContentType.TEXT_DELTA.value:
                     text_chunk = delta.get("text", "")
 
                     if text_chunk:
                         await self.on_message(
                             AgentMessage(
-                                type=MessageType.STREAM_EVENT,
+                                type=AgentMessageType.STREAM_EVENT,
                                 data={"text": text_chunk},
                             )
                         )
@@ -99,7 +86,7 @@ async def _handle_message(self, message: Any) -> None:
 
             await self.on_message(
                 AgentMessage(
-                    type=MessageType.ASSISTANT,
+                    type=AgentMessageType.ASSISTANT,
                     data={"content": content},
                 )
             )

src/agent_chat_cli/utils/enums.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+
+class AgentMessageType(Enum):
+    ASSISTANT = "assistant"
+    INIT = "init"
+    RESULT = "result"
+    STREAM_EVENT = "stream_event"
+    SYSTEM = "system"
+
+
+class ContentType(Enum):
+    TEXT = "text"
+    TOOL_USE = "tool_use"
+    CONTENT_BLOCK_DELTA = "content_block_delta"
+    TEXT_DELTA = "text_delta"

src/agent_chat_cli/utils/message_bus.py

@@ -10,7 +10,8 @@
     AgentMessage as AgentMessageWidget,
     ToolMessage,
 )
-from agent_chat_cli.utils.agent_loop import AgentMessage, MessageType, ContentType
+from agent_chat_cli.utils.agent_loop import AgentMessage
+from agent_chat_cli.utils.enums import AgentMessageType, ContentType
 
 if TYPE_CHECKING:
     from textual.app import App
@@ -30,11 +31,11 @@ async def _scroll_to_bottom(self) -> None:
 
     async def handle_agent_message(self, message: AgentMessage) -> None:
         match message.type:
-            case MessageType.STREAM_EVENT:
+            case AgentMessageType.STREAM_EVENT:
                 await self._handle_stream_event(message)
-            case MessageType.ASSISTANT:
+            case AgentMessageType.ASSISTANT:
                 await self._handle_assistant(message)
-            case MessageType.RESULT:
+            case AgentMessageType.RESULT:
                 await self._handle_result()
 
     async def _handle_stream_event(self, message: AgentMessage) -> None: