feat: simplify durable agent wf def (#204)

* feat: simplify durable agent wf def

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e ruff

Signed-off-by: Samantha Coyle <[email protected]>

* fix: updates for tests

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e ruff again

Signed-off-by: Samantha Coyle <[email protected]>

* feat: update to add tests

Signed-off-by: Samantha Coyle <[email protected]>

* fix: updates for tests and use concrete types not dictionaries

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e ruff

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e flake*

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e ruff again

Signed-off-by: Samantha Coyle <[email protected]>

* fix: updates for tests

Signed-off-by: Samantha Coyle <[email protected]>

* style: renaming for bilgin

Signed-off-by: Samantha Coyle <[email protected]>

* fix: use dict instead of pydantic due to storage layer incompatibility

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e ruff

Signed-off-by: Samantha Coyle <[email protected]>

---------

Signed-off-by: Samantha Coyle <[email protected]>

Author: sicoyle

dapr_agents/agents/base.py

@@ -134,20 +134,20 @@ def set_name_from_role(cls, values: dict):
         return values
 
     @model_validator(mode="after")
-    def validate_llm(cls, values):
+    def validate_llm(self):
         """Validate that LLM is properly configured."""
-        if hasattr(values, "llm"):
-            if values.llm is None:
+        if hasattr(self, "llm"):
+            if self.llm is None:
                 logger.warning("LLM client is None, some functionality may be limited.")
             else:
                 try:
                     # Validate LLM is properly configured by accessing it as this is required to be set.
-                    _ = values.llm
+                    _ = self.llm
                 except Exception as e:
                     logger.error(f"Failed to initialize LLM: {e}")
-                    values.llm = None
+                    self.llm = None
 
-        return values
+        return self
 
     def model_post_init(self, __context: Any) -> None:
         """

dapr_agents/agents/durableagent/agent.py

@@ -1,6 +1,7 @@
 from pydantic import BaseModel, Field
 from typing import List, Optional, Dict, Any
 from dapr_agents.types import MessageContent, ToolExecutionRecord
+from dapr_agents.types.workflow import DaprWorkflowStatus
 from datetime import datetime
 import uuid
 
@@ -60,6 +61,10 @@ class DurableAgentWorkflowEntry(BaseModel):
         default=None,
         description="OpenTelemetry trace context for workflow resumption.",
     )
+    status: str = Field(
+        default=DaprWorkflowStatus.RUNNING.value,
+        description="Current status of the workflow.",
+    )
 
 
 class DurableAgentWorkflowState(BaseModel):

dapr_agents/agents/durableagent/state.py

@@ -125,7 +125,7 @@ def create_resumed_workflow_context(
 
             # Create AGENT span with proper agent name for resumed workflow
             agent_display_name = agent_name or "DurableAgent"
-            span_name = f"{agent_display_name}.ToolCallingWorkflow"
+            span_name = f"{agent_display_name}.AgenticWorkflow"
             with tracer.start_as_current_span(span_name) as span:
                 # Set AGENT span attributes
                 from .constants import OPENINFERENCE_SPAN_KIND

dapr_agents/observability/context_storage.py

@@ -329,7 +329,7 @@ def _extract_workflow_name(self, args: Any, kwargs: Any) -> str:
         workflow_name = (
             workflow
             if isinstance(workflow, str)
-            else getattr(workflow, "__name__", "ToolCallingWorkflow")
+            else getattr(workflow, "__name__", "AgenticWorkflow")
         )
         return workflow_name
 

dapr_agents/observability/wrappers/workflow.py

@@ -192,9 +192,7 @@ def _build_task_attributes(
 
             # Add resource-level attributes for Phoenix UI grouping
             attributes["resource.workflow.instance_id"] = instance_id
-            attributes[
-                "resource.workflow.name"
-            ] = "ToolCallingWorkflow"  # Could be dynamic
+            attributes["resource.workflow.name"] = "AgenticWorkflow"  # Could be dynamic
 
             # Log the trace context for debugging (expected to be disconnected for Dapr Workflows)
             from opentelemetry import trace
@@ -323,7 +321,7 @@ def _create_context_for_resumed_workflow(
             logger.debug(f"Creating AGENT span for resumed workflow {instance_id}")
 
             agent_name = getattr(instance, "name", "DurableAgent")
-            workflow_name = instance_data.get("workflow_name", "ToolCallingWorkflow")
+            workflow_name = instance_data.get("workflow_name", "AgenticWorkflow")
             span_name = f"{agent_name}.{workflow_name}"
             attributes = {
                 "openinference.span.kind": "AGENT",
@@ -675,15 +673,13 @@ def _categorize_workflow_task(self, task_name: str) -> str:
         Returns:
             str: Semantic category for the task type
         """
-        if task_name in ["record_initial_entry", "get_workflow_entry_info"]:
+        if task_name in ["record_initial_entry"]:
             return "initialization"
-        elif task_name in ["append_assistant_message", "append_tool_message"]:
-            return "state_management"
         elif task_name in ["finalize_workflow", "finish_workflow"]:
             return "finalization"
         elif task_name in ["broadcast_message_to_agents", "send_response_back"]:
             return "communication"
-        elif task_name == "generate_response":
+        elif task_name == "call_llm":
             return "llm_generation"
         elif task_name == "run_tool":
             return "tool_execution"

dapr_agents/observability/wrappers/workflow_task.py

@@ -141,7 +141,8 @@ async def executor(**kwargs: Any) -> Any:
                 tool_args_model = create_pydantic_model_from_schema(
                     mcp_tool.inputSchema, f"{tool_name}Args"
                 )
-            except Exception:
+            except Exception as e:
+                logger.warning(f"Failed to create schema for tool '{tool_name}': {e}")
                 pass
 
         return cls(

dapr_agents/tool/base.py

@@ -8,15 +8,15 @@
 from datetime import datetime, timezone
 from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
 
+from pydantic import BaseModel
 from dapr.ext.workflow import (
     DaprWorkflowClient,
     WorkflowActivityContext,
     WorkflowRuntime,
 )
 from dapr.ext.workflow.workflow_state import WorkflowState
 from durabletask import task as dtask
-from pydantic import BaseModel, ConfigDict, Field
-from typing import ClassVar
+from pydantic import ConfigDict, Field
 
 from dapr_agents.agents.base import ChatClientBase
 from dapr_agents.types.workflow import DaprWorkflowStatus
@@ -83,12 +83,6 @@ def model_post_init(self, __context: Any) -> None:
 
         self.start_runtime()
 
-        # Discover and register tasks and workflows
-        discovered_tasks = self._discover_tasks()
-        self._register_tasks(discovered_tasks)
-        discovered_wfs = self._discover_workflows()
-        self._register_workflows(discovered_wfs)
-
         # Set up automatic signal handlers for graceful shutdown
         try:
             self.setup_signal_handlers()
@@ -365,6 +359,10 @@ def register_tasks_from_package(self, package_name: str) -> None:
     def _register_tasks(self, tasks: Dict[str, Callable]) -> None:
         """Register each discovered task with the Dapr runtime using direct registration."""
         for task_name, method in tasks.items():
+            # Don't reregister tasks that are already registered
+            if task_name in self.tasks:
+                continue
+
             llm = self._choose_llm_for(method)
             logger.debug(
                 f"Registering task '{task_name}' with llm={getattr(llm, '__class__', None)}"
@@ -451,6 +449,10 @@ def _discover_workflows(self) -> Dict[str, Callable]:
     def _register_workflows(self, wfs: Dict[str, Callable]) -> None:
         """Register each discovered workflow with the Dapr runtime."""
         for wf_name, method in wfs.items():
+            # Don't reregister workflows that are already registered
+            if wf_name in self.workflows:
+                continue
+
             # Use a closure helper to avoid late-binding capture issues.
             def make_wrapped(meth: Callable) -> Callable:
                 @functools.wraps(meth)
@@ -538,6 +540,17 @@ def start_runtime(self):
         else:
             logger.debug("Workflow runtime already running; skipping.")
 
+        self._ensure_activities_registered()
+
+    def _ensure_activities_registered(self):
+        """Ensure all workflow activities are registered with the Dapr runtime."""
+        # Discover and register tasks and workflows
+        discovered_tasks = self._discover_tasks()
+        self._register_tasks(discovered_tasks)
+        discovered_wfs = self._discover_workflows()
+        self._register_workflows(discovered_wfs)
+        logger.debug("Workflow activities registration completed.")
+
     def _sync_workflow_state_after_startup(self):
         """
         Sync database workflow state with actual Dapr workflow status after runtime startup.
@@ -555,25 +568,17 @@ def _sync_workflow_state_after_startup(self):
                 )
                 return
 
-            logger.debug("Syncing workflow state with Dapr after runtime startup...")
             self.load_state()
-
-            # Check if we have instances to sync
-            instances = (
-                getattr(self.state, "instances", {})
-                if hasattr(self.state, "instances")
-                else self.state.get("instances", {})
-            )
-            if not instances:
-                return
+            instances = self.state.get("instances", {})
 
             logger.debug(f"Found {len(instances)} workflow instances to sync")
 
             # Sync each instance with Dapr's actual status
             for instance_id, instance_data in instances.items():
                 try:
                     # Skip if already completed
-                    if instance_data.get("end_time") is not None:
+                    end_time = instance_data.get("end_time")
+                    if end_time is not None:
                         continue
 
                     # Get actual status from Dapr
@@ -595,6 +600,7 @@ def _sync_workflow_state_after_startup(self):
                                 timezone.utc
                             ).isoformat()
                             instance_data["status"] = runtime_status.lower()
+
                             logger.debug(
                                 f"Marked workflow {instance_id} as {runtime_status.lower()} in database"
                             )
@@ -615,6 +621,7 @@ def _sync_workflow_state_after_startup(self):
                             timezone.utc
                         ).isoformat()
                         instance_data["status"] = DaprWorkflowStatus.COMPLETED.value
+
                         logger.debug(
                             f"Workflow {instance_id} no longer in Dapr, marked as completed"
                         )
@@ -819,9 +826,7 @@ async def _create_agent_span_for_resumed_workflow(
                 # Get tracer and create AGENT span as child of the original trace
                 tracer = trace.get_tracer(__name__)
                 agent_name = getattr(self, "name", "DurableAgent")
-                workflow_name = instance_data.get(
-                    "workflow_name", "ToolCallingWorkflow"
-                )
+                workflow_name = instance_data.get("workflow_name", "AgenticWorkflow")
                 span_name = f"{agent_name}.{workflow_name}"
 
                 # Create the AGENT span that will show up in the trace

dapr_agents/workflow/base.py

@@ -37,7 +37,7 @@ def initialize_state(self) -> None:
 
             if not isinstance(self.state, dict):
                 raise TypeError(
-                    f"Invalid state type: {type(self.state)}. Expected dict or Pydantic model."
+                    f"Invalid state type: {type(self.state)}. Expected dict."
                 )
 
             logger.debug(f"Workflow state initialized with {len(self.state)} key(s).")

dapr_agents/workflow/mixins/state.py

@@ -35,6 +35,7 @@ dependencies = [
     "cloudevents>=1.11.0,<2.0.0",
     "numpy>=2.2.2,<3.0.0",
     "mcp>=1.7.1,<2.0.0",
+    "websockets>=15.0.0,<16.0.0",
     "python-dotenv>=1.1.1,<2.0.0",
     "posthog<6.0.0",
     "nltk>=3.8.0,<4.0.0",