chore: clean library import

Author: daltonnyx

AgentCrew/modules/a2a/common/client/client.py

@@ -1,35 +1,42 @@
+from __future__ import annotations
+
 import json
 
 from collections.abc import AsyncGenerator
-from typing import Any, Optional, Dict
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
 import httpx
 
-from httpx._types import TimeoutTypes
 from httpx_sse import aconnect_sse
 
 from a2a.types import (
     A2ARequest,
-    AgentCard,
-    CancelTaskRequest,
     CancelTaskResponse,
-    GetTaskPushNotificationConfigRequest,
     GetTaskPushNotificationConfigResponse,
-    GetTaskRequest,
     GetTaskResponse,
-    TaskQueryParams,
-    SendMessageRequest,
     SendMessageResponse,
-    SendStreamingMessageRequest,
     TaskPushNotificationConfig,
     SendStreamingMessageResponse,
-    SetTaskPushNotificationConfigRequest,
     SetTaskPushNotificationConfigResponse,
-    TaskIdParams,
-    MessageSendParams,
 )
 
+if TYPE_CHECKING:
+    from typing import Any, Dict, Optional
+    from httpx._types import TimeoutTypes
+    from a2a.types import (
+        AgentCard,
+        CancelTaskRequest,
+        GetTaskPushNotificationConfigRequest,
+        GetTaskRequest,
+        MessageSendParams,
+        SendMessageRequest,
+        SendStreamingMessageRequest,
+        SetTaskPushNotificationConfigRequest,
+        TaskIdParams,
+        TaskQueryParams,
+    )
+
 
 class A2AClient:
     def __init__(

AgentCrew/modules/a2a/common/server/task_manager.py

@@ -1,16 +1,15 @@
+from __future__ import annotations
+
 import asyncio
 
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterable
+from typing import TYPE_CHECKING
 
 from .utils import new_not_implemented_error
 from a2a.types import (
-    Artifact,
-    CancelTaskRequest,
     CancelTaskResponse,
-    GetTaskPushNotificationConfigRequest,
     GetTaskPushNotificationConfigResponse,
-    GetTaskRequest,
     GetTaskResponse,
     SetTaskPushNotificationConfigSuccessResponse,
     GetTaskPushNotificationConfigSuccessResponse,
@@ -20,21 +19,13 @@
     GetTaskSuccessResponse,
     JSONRPCErrorResponse,
     JSONRPCResponse,
-    PushNotificationConfig,
-    SendMessageRequest,
     SendMessageResponse,
-    SendStreamingMessageRequest,
     SendStreamingMessageResponse,
-    SetTaskPushNotificationConfigRequest,
     SetTaskPushNotificationConfigResponse,
     Task,
-    TaskIdParams,
     TaskNotCancelableError,
     TaskNotFoundError,
     TaskPushNotificationConfig,
-    TaskQueryParams,
-    TaskResubscriptionRequest,
-    MessageSendParams,
     TaskState,
     TaskStatus,
     TaskStatusUpdateEvent,
@@ -43,6 +34,22 @@
 
 from loguru import logger
 
+if TYPE_CHECKING:
+    from a2a.types import (
+        Artifact,
+        CancelTaskRequest,
+        GetTaskPushNotificationConfigRequest,
+        GetTaskRequest,
+        MessageSendParams,
+        PushNotificationConfig,
+        SendMessageRequest,
+        SendStreamingMessageRequest,
+        SetTaskPushNotificationConfigRequest,
+        TaskIdParams,
+        TaskQueryParams,
+        TaskResubscriptionRequest,
+    )
+
 
 class TaskManager(ABC):
     @abstractmethod

AgentCrew/modules/a2a/registry.py

@@ -1,12 +1,15 @@
-"""
-Registry for managing A2A agents and their capabilities.
-"""
+from __future__ import annotations
 
-from typing import Dict, List, Optional, Any
+from typing import TYPE_CHECKING
 from pydantic import BaseModel
-from AgentCrew.modules.agents import AgentManager, LocalAgent
-from a2a.types import AgentCard
 from .agent_cards import create_agent_card
+from AgentCrew.modules.agents import LocalAgent
+
+
+if TYPE_CHECKING:
+    from typing import Any, Dict, List, Optional
+    from AgentCrew.modules.agents import AgentManager
+    from a2a.types import AgentCard
 
 
 class AgentInfo(BaseModel):

AgentCrew/modules/a2a/task_manager.py

@@ -1,52 +1,57 @@
-"""
-Task management for A2A protocol implementation.
-"""
+from __future__ import annotations
 
 import asyncio
 from datetime import datetime
-from typing import Dict, AsyncIterable, Optional, Any, Union
-from AgentCrew.modules.agents import AgentManager, LocalAgent
+from typing import TYPE_CHECKING
 from AgentCrew.modules.agents.base import MessageType
 from loguru import logger
 import tempfile
 import os
 
 from a2a.types import (
     CancelTaskResponse,
+    JSONRPCError,
     GetTaskResponse,
     GetTaskSuccessResponse,
-    JSONRPCError,
     JSONRPCErrorResponse,
     SendMessageResponse,
-    SendStreamingMessageRequest,
     SendStreamingMessageResponse,
     SendStreamingMessageSuccessResponse,
     CancelTaskSuccessResponse,
-    SetTaskPushNotificationConfigRequest,
     SetTaskPushNotificationConfigResponse,
-    GetTaskPushNotificationConfigRequest,
     GetTaskPushNotificationConfigResponse,
     SendMessageSuccessResponse,
-    TaskResubscriptionRequest,
     Task,
     TaskStatus,
     TaskState,
-    SendMessageRequest,
-    GetTaskRequest,
-    CancelTaskRequest,
-    JSONRPCResponse,
     TaskStatusUpdateEvent,
     TaskArtifactUpdateEvent,
     TaskNotFoundError,
     TaskNotCancelableError,
 )
+
+from AgentCrew.modules.agents import LocalAgent
 from .adapters import (
     convert_a2a_message_to_agent,
     convert_agent_response_to_a2a_artifact,
     convert_agent_message_to_a2a,
 )
 from .common.server.task_manager import TaskManager
 
+if TYPE_CHECKING:
+    from typing import Any, AsyncIterable, Dict, Optional, Union
+    from AgentCrew.modules.agents import AgentManager
+    from a2a.types import (
+        CancelTaskRequest,
+        GetTaskPushNotificationConfigRequest,
+        GetTaskRequest,
+        SendMessageRequest,
+        SendStreamingMessageRequest,
+        SetTaskPushNotificationConfigRequest,
+        TaskResubscriptionRequest,
+        JSONRPCResponse,
+    )
+
 
 class AgentTaskManager(TaskManager):
     """Manages tasks for a specific agent"""
@@ -126,8 +131,6 @@ async def on_send_message(
                         f.write(part["file_data"])
                     file_part = self.file_handler.process_file(temp_file)
                     if not file_part:
-                        from AgentCrew.modules.agents.base import MessageType
-
                         file_part = self.agent.format_message(
                             MessageType.FileContent, {"file_uri": temp_file}
                         )

AgentCrew/modules/agents/base.py

@@ -1,7 +1,12 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import AsyncGenerator, Tuple, Dict, List, Optional, Any, Callable, Union
+from typing import TYPE_CHECKING
 from enum import Enum
 
+if TYPE_CHECKING:
+    from typing import AsyncGenerator, Tuple, Dict, List, Optional, Any, Callable, Union
+
 
 class MessageType(Enum):
     Assistant = 0

AgentCrew/modules/agents/local_agent.py

@@ -1,13 +1,18 @@
+from __future__ import annotations
+
 from datetime import datetime
 import os
 import time
 import copy
-from typing import Dict, Any, List, Optional, Callable, Literal, Union
+from typing import List, TYPE_CHECKING
 from AgentCrew.modules.llm import BaseLLMService
 
 from .base import BaseAgent, MessageType
 from loguru import logger
 
+if TYPE_CHECKING:
+    from typing import Dict, Any, Optional, Callable, Literal, Union
+
 SHRINK_CONTEXT_THRESHOLD = 90_000
 SHRINK_LENGTH_THRESHOLD = 15
 

AgentCrew/modules/agents/remote_agent.py

@@ -1,12 +1,13 @@
-from typing import Callable, Dict, Any, List, Optional, Tuple, Union
+from __future__ import annotations
+
+from typing import Dict, TYPE_CHECKING
 from uuid import uuid4
 
 from pydantic import ValidationError
 from AgentCrew.modules.a2a.adapters import (
     convert_agent_message_to_a2a,
 )
 
-# from AgentCrew.modules.llm.message import MessageTransformer
 from .base import BaseAgent, MessageType
 from AgentCrew.modules.a2a.common.client import A2ACardResolver, A2AClient
 from a2a.types import (
@@ -17,6 +18,9 @@
     TextPart,
 )
 
+if TYPE_CHECKING:
+    from typing import Callable, Any, List, Optional, Tuple, Union
+
 
 class RemoteAgent(BaseAgent):
     def __init__(

AgentCrew/modules/browser_automation/tool.py

@@ -1,22 +1,13 @@
-"""
-Browser automation tool definitions and handlers.
-
-Provides eight tools for browser automation:
-- browser_navigate: Navigate to URLs
-- browser_click: Click elements using UUID selectors
-- browser_scroll: Scroll page content
-- browser_get_content: Extract page content and interactive elements as markdown
-- browser_input: Input data into form fields using UUID selectors
-- browser_get_elements_by_text: Find elements by text content
-- browser_capture_screenshot: Capture page screenshots
-- browser_send_key: send keyboard events (arrow keys, function keys, etc.)
-"""
-
-from typing import Dict, Any, Callable, Union, List
-from .service import BrowserAutomationService
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 import difflib
 import time
 
+if TYPE_CHECKING:
+    from .service import BrowserAutomationService
+    from typing import Dict, Any, Callable, Union, List
+
 
 def get_browser_navigate_tool_definition(provider="claude") -> Dict[str, Any]:
     """Get tool definition for browser navigation."""

AgentCrew/modules/console/display_handlers.py

@@ -16,10 +16,10 @@
     RICH_STYLE_YELLOW,
     RICH_STYLE_BLUE,
     RICH_STYLE_RED,
+    RICH_STYLE_GREEN,
     RICH_STYLE_GRAY,
     RICH_STYLE_YELLOW_BOLD,
     RICH_STYLE_GREEN_BOLD,
-    RICH_STYLE_BLUE_BOLD,
     RICH_STYLE_FILE_ACCENT_BOLD,
     RICH_STYLE_WHITE,
     CODE_THEME,
@@ -64,6 +64,32 @@ def display_message(self, message: Text):
         """Display a generic message."""
         self.console.print(message)
 
+    def display_user_message(self, message: str):
+        header = Text(
+            "👤 YOU:",
+            style=RICH_STYLE_GREEN_BOLD,
+        )
+        user_panel = Panel(
+            Text(message),
+            title=header,
+            title_align="left",
+            border_style=RICH_STYLE_BLUE,
+        )
+        self.console.print(user_panel)
+
+    def display_assistant_message(self, agent_name: str, message: str):
+        header = Text(
+            f"🤖 {agent_name.upper()}:",
+            style=RICH_STYLE_GREEN_BOLD,
+        )
+        assistant_panel = Panel(
+            Markdown(message, code_theme=CODE_THEME),
+            title=header,
+            title_align="left",
+            border_style=RICH_STYLE_GREEN,
+        )
+        self.console.print(assistant_panel)
+
     def display_divider(self):
         """Display a divider line."""
         pass
@@ -163,19 +189,13 @@ def display_loaded_conversation(self, messages, message_handler):
         for msg in messages[last_consolidated_idx:]:
             role = msg.get("role")
             if role == "user":
-                self.console.print(Text("\n👤 YOU:", style=RICH_STYLE_BLUE_BOLD))
                 content = self._extract_message_content(msg)
-                self.console.print(content)
-                self.display_divider()
+                self.display_user_message(content)
             elif role == "assistant":
-                agent_name = message_handler.agent.name
-                self.console.print(
-                    Text(f"\n🤖 {agent_name.upper()}:", style=RICH_STYLE_GREEN_BOLD)
-                )
+                agent_name = msg.get("agent") or message_handler.agent.name
                 content = self._extract_message_content(msg)
                 # Format as markdown for better display
-                self.console.print(Markdown(content, code_theme=CODE_THEME))
-                self.display_divider()
+                self.display_assistant_message(agent_name, content)
                 if "tool_calls" in msg:
                     from .tool_display import ToolDisplayHandlers