refactor: Remove unnecessary TYPE_CHECKING guard for StateModelBundle import and clarify bundle override logic

Signed-off-by: Roberto Rodriguez <[email protected]>

Author: Cyb3rWard0g

dapr_agents/agents/components.py

@@ -4,7 +4,7 @@
 import random
 import time
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Sequence
+from typing import Any, Callable, Dict, Optional, Sequence
 
 from dapr.clients.grpc._state import Concurrency, Consistency, StateOptions
 from pydantic import BaseModel, ValidationError
@@ -14,14 +14,12 @@
     AgentRegistryConfig,
     AgentStateConfig,
     DEFAULT_AGENT_WORKFLOW_BUNDLE,
+    StateModelBundle,
 )
 from dapr_agents.agents.schemas import AgentWorkflowEntry
 from dapr_agents.storage.daprstores.stateservice import StateStoreError
 from dapr_agents.types.workflow import DaprWorkflowStatus
 
-if TYPE_CHECKING:
-    from dapr_agents.agents.configs import StateModelBundle
-
 logger = logging.getLogger(__name__)
 
 
@@ -46,7 +44,7 @@ def __init__(
         registry: Optional[AgentRegistryConfig] = None,
         base_metadata: Optional[Dict[str, Any]] = None,
         max_etag_attempts: int = 10,
-        default_bundle: Optional["StateModelBundle"] = None,
+        default_bundle: Optional[StateModelBundle] = None,
     ) -> None:
         """
         Initialize component wiring.
@@ -84,6 +82,9 @@ def __init__(
 
         bundle = None
         if state is not None:
+            # Allow default_bundle to override the state's bundle. This enables
+            # orchestrators and agents to share the same AgentStateConfig instance
+            # while each using their own specialized state model schemas.
             if default_bundle is not None:
                 state.ensure_bundle(default_bundle)
             try:

dapr_agents/agents/orchestrators/base.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 from collections.abc import Coroutine
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import Any, Dict, Optional
 
 import dapr.ext.workflow as wf
 
@@ -13,12 +13,10 @@
     AgentPubSubConfig,
     AgentRegistryConfig,
     AgentStateConfig,
+    StateModelBundle,
 )
 from dapr_agents.agents.utils.text_printer import ColorTextFormatter
 
-if TYPE_CHECKING:
-    from dapr_agents.agents.configs import StateModelBundle
-
 logger = logging.getLogger(__name__)
 
 
@@ -44,7 +42,7 @@ def __init__(
         agent_metadata: Optional[Dict[str, Any]] = None,
         runtime: Optional[wf.WorkflowRuntime] = None,
         workflow_client: Optional[wf.DaprWorkflowClient] = None,
-        default_bundle: Optional["StateModelBundle"] = None,
+        default_bundle: Optional[StateModelBundle] = None,
     ) -> None:
         """
         Initialize the orchestrator base.