Fix imports for durabletask (#3345)

python/packages/azurefunctions/agent_framework_azurefunctions/__init__.py

@@ -2,8 +2,6 @@
 
 import importlib.metadata
 
-from agent_framework_durabletask import AgentCallbackContext, AgentResponseCallbackProtocol, DurableAIAgent
-
 from ._app import AgentFunctionApp
 
 try:
@@ -12,9 +10,6 @@ except importlib.metadata.PackageNotFoundError:
     __version__ = "0.0.0"  # Fallback for development mode
 
 __all__ = [
-    "AgentCallbackContext",
     "AgentFunctionApp",
-    "AgentResponseCallbackProtocol",
-    "DurableAIAgent",
     "__version__",
 ]

python/packages/core/agent_framework/azure/__init__.py

@@ -4,9 +4,9 @@ import importlib
 from typing import Any
 
 _IMPORTS: dict[str, tuple[str, str]] = {
-    "AgentCallbackContext": ("agent_framework_azurefunctions", "agent-framework-azurefunctions"),
+    "AgentCallbackContext": ("agent_framework_durabletask", "agent-framework-durabletask"),
     "AgentFunctionApp": ("agent_framework_azurefunctions", "agent-framework-azurefunctions"),
-    "AgentResponseCallbackProtocol": ("agent_framework_azurefunctions", "agent-framework-azurefunctions"),
+    "AgentResponseCallbackProtocol": ("agent_framework_durabletask", "agent-framework-durabletask"),
     "AzureAIAgentClient": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
     "AzureAIAgentOptions": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
     "AzureAIClient": ("agent_framework_azure_ai", "agent-framework-azure-ai"),
@@ -23,7 +23,10 @@ _IMPORTS: dict[str, tuple[str, str]] = {
     "AzureOpenAIResponsesOptions": ("agent_framework.azure._responses_client", "agent-framework-core"),
     "AzureOpenAISettings": ("agent_framework.azure._shared", "agent-framework-core"),
     "AzureUserSecurityContext": ("agent_framework.azure._chat_client", "agent-framework-core"),
-    "DurableAIAgent": ("agent_framework_azurefunctions", "agent-framework-azurefunctions"),
+    "DurableAIAgent": ("agent_framework_durabletask", "agent-framework-durabletask"),
+    "DurableAIAgentClient": ("agent_framework_durabletask", "agent-framework-durabletask"),
+    "DurableAIAgentOrchestrationContext": ("agent_framework_durabletask", "agent-framework-durabletask"),
+    "DurableAIAgentWorker": ("agent_framework_durabletask", "agent-framework-durabletask"),
     "get_entra_auth_token": ("agent_framework.azure._entra_id_authentication", "agent-framework-core"),
 }
 

python/packages/core/agent_framework/azure/__init__.pyi

@@ -8,11 +8,14 @@ from agent_framework_azure_ai import (
     AzureAISettings,
 )
 from agent_framework_azure_ai_search import AzureAISearchContextProvider, AzureAISearchSettings
-from agent_framework_azurefunctions import (
+from agent_framework_azurefunctions import AgentFunctionApp
+from agent_framework_durabletask import (
     AgentCallbackContext,
-    AgentFunctionApp,
     AgentResponseCallbackProtocol,
     DurableAIAgent,
+    DurableAIAgentClient,
+    DurableAIAgentOrchestrationContext,
+    DurableAIAgentWorker,
 )
 
 from agent_framework.azure._assistants_client import AzureOpenAIAssistantsClient
@@ -37,5 +40,8 @@ __all__ = [
     "AzureOpenAIResponsesClient",
     "AzureOpenAISettings",
     "DurableAIAgent",
+    "DurableAIAgentClient",
+    "DurableAIAgentOrchestrationContext",
+    "DurableAIAgentWorker",
     "get_entra_auth_token",
 ]

python/packages/durabletask/README.md

@@ -16,7 +16,7 @@ The durable task integration lets you host Microsoft Agent Framework agents usin
 
 ```python
 from durabletask import TaskHubGrpcWorker
-from agent_framework_durabletask import AgentWorker
+from agent_framework.azure import DurableAIAgentWorker
 
 # Create the worker
 with DurableTaskSchedulerWorker(...) as worker:

python/packages/durabletask/agent_framework_durabletask/__init__.py

@@ -2,6 +2,8 @@
 
 """Durable Task integration for Microsoft Agent Framework."""
 
+import importlib.metadata
+
 from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
 from ._client import DurableAIAgentClient
 from ._constants import (
@@ -49,6 +51,11 @@ from ._response_utils import ensure_response_format, load_agent_response
 from ._shim import DurableAIAgent
 from ._worker import DurableAIAgentWorker
 
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"  # Fallback for development mode
+
 __all__ = [
     "DEFAULT_MAX_POLL_RETRIES",
     "DEFAULT_POLL_INTERVAL_SECONDS",
@@ -96,6 +103,7 @@ __all__ = [
     "DurableAgentThread",
     "DurableStateFields",
     "RunRequest",
+    "__version__",
     "ensure_response_format",
     "load_agent_response",
 ]

python/packages/durabletask/agent_framework_durabletask/_client.py

@@ -27,7 +27,7 @@ class DurableAIAgentClient(DurableAgentProvider[AgentResponse]):
     Example:
         ```python
         from durabletask import TaskHubGrpcClient
-        from agent_framework_durabletask import DurableAIAgentClient
+        from agent_framework.azure import DurableAIAgentClient
 
         # Create the underlying client
         client = TaskHubGrpcClient(host_address="localhost:4001")

python/packages/durabletask/agent_framework_durabletask/_orchestration_context.py

@@ -27,7 +27,7 @@ class DurableAIAgentOrchestrationContext(DurableAgentProvider[DurableAgentTask])
     Example:
         ```python
         from durabletask import Orchestration
-        from agent_framework_durabletask import DurableAIAgentOrchestrationContext
+        from agent_framework.azure import DurableAIAgentOrchestrationContext
 
 
         def my_orchestration(context: OrchestrationContext):

python/packages/durabletask/agent_framework_durabletask/_worker.py

@@ -30,7 +30,7 @@ class DurableAIAgentWorker:
         ```python
         from durabletask import TaskHubGrpcWorker
         from agent_framework import ChatAgent
-        from agent_framework_durabletask import DurableAIAgentWorker
+        from agent_framework.azure import DurableAIAgentWorker
 
         # Create the underlying worker
         worker = TaskHubGrpcWorker(host_address="localhost:4001")