Python: Move azurefunctions to azure for import (#2141)

* Move import to Azure

* fix mypy

* Update python/packages/azurefunctions/README.md

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Add missing types

* Address comments

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

python/packages/azurefunctions/README.md

@@ -16,6 +16,12 @@ The durable agent extension lets you host Microsoft Agent Framework agents on Az
 
 See the durable functions integration sample in the repository to learn how to:
 
+```python
+from agent_framework.azure import AgentFunctionApp
+
+_app = AgentFunctionApp()
+```
+
 - Register agents with `AgentFunctionApp`
 - Post messages using the generated `/api/agents/{agent_name}/run` endpoint
 

python/packages/azurefunctions/agent_framework_azurefunctions/__init__.py

@@ -7,12 +7,11 @@ enabling durable, stateful AI agents deployed as Azure Function Apps.
 
 from ._app import AgentFunctionApp
 from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
-from ._orchestration import DurableAIAgent, get_agent
+from ._orchestration import DurableAIAgent
 
 __all__ = [
     "AgentCallbackContext",
     "AgentFunctionApp",
     "AgentResponseCallbackProtocol",
     "DurableAIAgent",
-    "get_agent",
 ]

python/packages/azurefunctions/agent_framework_azurefunctions/_app.py

@@ -9,7 +9,7 @@ with Azure Durable Entities, enabling stateful and durable AI agent execution.
 import json
 import re
 from collections.abc import Callable, Mapping
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 import azure.durable_functions as df
 import azure.functions as func
@@ -19,6 +19,7 @@ from ._callbacks import AgentResponseCallbackProtocol
 from ._entities import create_agent_entity
 from ._errors import IncomingRequestError
 from ._models import AgentSessionId, RunRequest
+from ._orchestration import AgentOrchestrationContextType, DurableAIAgent
 from ._state import AgentState
 
 logger = get_logger("agent_framework.azurefunctions")
@@ -30,18 +31,46 @@ WAIT_FOR_RESPONSE_FIELD: str = "wait_for_response"
 WAIT_FOR_RESPONSE_HEADER: str = "x-ms-wait-for-response"
 
 
-class AgentFunctionApp(df.DFApp):
+EntityHandler = Callable[[df.DurableEntityContext], None]
+HandlerT = TypeVar("HandlerT", bound=Callable[..., Any])
+
+if TYPE_CHECKING:
+
+    class DFAppBase:
+        def __init__(self, http_auth_level: func.AuthLevel = func.AuthLevel.FUNCTION) -> None: ...
+
+        def function_name(self, name: str) -> Callable[[HandlerT], HandlerT]: ...
+
+        def route(self, route: str, methods: list[str]) -> Callable[[HandlerT], HandlerT]: ...
+
+        def durable_client_input(self, client_name: str) -> Callable[[HandlerT], HandlerT]: ...
+
+        def entity_trigger(self, context_name: str, entity_name: str) -> Callable[[EntityHandler], EntityHandler]: ...
+
+        def orchestration_trigger(self, context_name: str) -> Callable[[HandlerT], HandlerT]: ...
+
+        def activity_trigger(self, input_name: str) -> Callable[[HandlerT], HandlerT]: ...
+
+else:
+    DFAppBase = df.DFApp  # type: ignore[assignment]
+
+
+class AgentFunctionApp(DFAppBase):
     """Main application class for creating durable agent function apps using Durable Entities.
 
     This class uses Durable Entities pattern for agent execution, providing:
+
     - Stateful agent conversations
     - Conversation history management
     - Signal-based operation invocation
     - Better state management than orchestrations
 
-    Usage:
-        ```python
-    from agent_framework.azurefunctions import AgentFunctionApp
+    Example:
+    -------
+
+    .. code-block:: python
+
+        from agent_framework.azure import AgentFunctionApp
         from agent_framework.azure import AzureOpenAIAssistantsClient
 
         # Create agents with unique names
@@ -64,9 +93,18 @@ class AgentFunctionApp(df.DFApp):
         app = AgentFunctionApp()
         app.add_agent(weather_agent)
         app.add_agent(math_agent)
-        ```
+
+
+        @app.orchestration_trigger(context_name="context")
+        def my_orchestration(context):
+            writer = app.get_agent(context, "WeatherAgent")
+            thread = writer.get_new_thread()
+            forecast_task = writer.run("What's the forecast?", thread=thread)
+            forecast = yield forecast_task
+            return forecast
 
     This creates:
+
     - HTTP trigger endpoint for each agent's requests (if enabled)
     - Durable entity for each agent's state management and execution
     - Full access to all Azure Functions capabilities
@@ -197,6 +235,30 @@ class AgentFunctionApp(df.DFApp):
 
         logger.debug(f"[AgentFunctionApp] Agent '{name}' added successfully")
 
+    def get_agent(
+        self,
+        context: AgentOrchestrationContextType,
+        agent_name: str,
+    ) -> DurableAIAgent:
+        """Return a DurableAIAgent proxy for a registered agent.
+
+        Args:
+            context: Durable Functions orchestration context invoking the agent.
+            agent_name: Name of the agent registered on this app.
+
+        Raises:
+            ValueError: If the requested agent has not been registered.
+
+        Returns:
+            DurableAIAgent wrapper bound to the orchestration context.
+        """
+        normalized_name = str(agent_name)
+
+        if normalized_name not in self.agents:
+            raise ValueError(f"Agent '{normalized_name}' is not registered with this app.")
+
+        return DurableAIAgent(context, normalized_name)
+
     def _setup_agent_functions(
         self,
         agent: AgentProtocol,
@@ -232,9 +294,13 @@ class AgentFunctionApp(df.DFApp):
         """
         run_function_name = self._build_function_name(agent_name, "http")
 
-        @self.function_name(run_function_name)
-        @self.route(route=f"agents/{agent_name}/run", methods=["POST"])
-        @self.durable_client_input(client_name="client")
+        function_name_decorator = self.function_name(run_function_name)
+        route_decorator = self.route(route=f"agents/{agent_name}/run", methods=["POST"])
+        durable_client_decorator = self.durable_client_input(client_name="client")
+
+        @function_name_decorator
+        @route_decorator
+        @durable_client_decorator
         async def http_start(req: func.HttpRequest, client: df.DurableOrchestrationClient) -> func.HttpResponse:
             """HTTP trigger that calls a durable entity to execute the agent and returns the result.
 
@@ -379,8 +445,9 @@ class AgentFunctionApp(df.DFApp):
 
     def _setup_health_route(self) -> None:
         """Register the optional health check route."""
+        health_route = self.route(route="health", methods=["GET"])
 
-        @self.route(route="health", methods=["GET"])
+        @health_route
         def health_check(req: func.HttpRequest) -> func.HttpResponse:
             """Built-in health check endpoint."""
             agent_info = [
@@ -643,8 +710,7 @@ class AgentFunctionApp(df.DFApp):
         headers: dict[str, str] = {}
         raw_headers = req.headers
         if isinstance(raw_headers, Mapping):
-            headers_mapping = cast(Mapping[Any, Any], raw_headers)
-            for key, value in headers_mapping.items():
+            for key, value in raw_headers.items():
                 if value is not None:
                     headers[str(key).lower()] = str(value)
         return headers
@@ -708,8 +774,7 @@ class AgentFunctionApp(df.DFApp):
         header_value = None
         raw_headers = req.headers
         if isinstance(raw_headers, Mapping):
-            headers_mapping = cast(Mapping[Any, Any], raw_headers)
-            for key, value in headers_mapping.items():
+            for key, value in raw_headers.items():
                 if str(key).lower() == WAIT_FOR_RESPONSE_HEADER:
                     header_value = value
                     break

python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py

@@ -10,7 +10,7 @@ allows for long-running agent conversations.
 import asyncio
 import inspect
 import json
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterable, Callable
 from typing import Any, cast
 
 import azure.durable_functions as df
@@ -340,7 +340,7 @@ class AgentEntity:
 def create_agent_entity(
     agent: AgentProtocol,
     callback: AgentResponseCallbackProtocol | None = None,
-):
+) -> Callable[[df.DurableEntityContext], None]:
     """Factory function to create an agent entity class.
 
     Args:

python/packages/azurefunctions/agent_framework_azurefunctions/_models.py

@@ -374,7 +374,7 @@ class AgentResponse:
 
     def to_dict(self) -> dict[str, Any]:
         """Convert to dictionary for JSON serialization."""
-        result = {
+        result: dict[str, Any] = {
             "message": self.message,
             "thread_id": self.thread_id,
             "status": self.status,

python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py

@@ -37,7 +37,7 @@ class DurableAIAgent(AgentProtocol):
     yielded in orchestrations to wait for the entity call to complete.
 
     Example usage in orchestration:
-        writer = get_agent(context, "WriterAgent")
+        writer = app.get_agent(context, "WriterAgent")
         thread = writer.get_new_thread()  # NOT yielded - returns immediately
 
         response = yield writer.run(  # Yielded - waits for entity call
@@ -104,7 +104,7 @@ class DurableAIAgent(AgentProtocol):
         Example:
             @app.orchestration_trigger(context_name="context")
             def my_orchestration(context):
-                agent = get_agent(context, "MyAgent")
+                agent = app.get_agent(context, "MyAgent")
                 thread = agent.get_new_thread()
                 result = yield agent.run("Hello", thread=thread)
         """
@@ -209,27 +209,3 @@ class DurableAIAgent(AgentProtocol):
                 return "\n".join(cast(list[str], messages))
             return self._messages_to_string(cast(list[ChatMessage], messages))
         return str(messages)
-
-
-def get_agent(context: AgentOrchestrationContextType, agent_name: str) -> DurableAIAgent:
-    """Return a :class:`DurableAIAgent` proxy scoped to ``agent_name``.
-
-    Usage::
-
-    from agent_framework.azurefunctions import get_agent
-
-
-        @app.orchestration_trigger(context_name="context")
-        def my_orchestration(context: DurableOrchestrationContext):
-            writer = get_agent(context, "WriterAgent")
-            thread = writer.get_new_thread()
-            response = yield writer.run("Write a haiku", thread=thread)
-
-    Args:
-        context: The orchestration context provided by Durable Functions.
-        agent_name: Name of the durable agent entity to call.
-
-    Returns:
-        DurableAIAgent wrapper for the specified agent.
-    """
-    return DurableAIAgent(context, agent_name)

python/packages/azurefunctions/agent_framework_azurefunctions/_state.py

@@ -25,7 +25,7 @@ class AgentState:
     - Message counting
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize empty agent state."""
         self.conversation_history: list[ChatMessage] = []
         self.last_response: str | None = None

python/packages/azurefunctions/tests/test_orchestration.py

@@ -8,10 +8,19 @@ from unittest.mock import Mock
 import pytest
 from agent_framework import AgentThread
 
-from agent_framework_azurefunctions import DurableAIAgent, get_agent
+from agent_framework_azurefunctions import AgentFunctionApp, DurableAIAgent
 from agent_framework_azurefunctions._models import AgentSessionId, DurableAgentThread
 
 
+def _app_with_registered_agents(*agent_names: str) -> AgentFunctionApp:
+    app = AgentFunctionApp(enable_health_check=False, enable_http_endpoints=False)
+    for name in agent_names:
+        agent = Mock()
+        agent.name = name
+        app.add_agent(agent)
+    return app
+
+
 class TestDurableAIAgent:
     """Test suite for DurableAIAgent wrapper."""
 
@@ -266,20 +275,28 @@ class TestDurableAIAgent:
         assert str(entity_id) == "@dafx-writeragent@test-guid-789"
 
 
-class TestGetAgentHelper:
-    """Test suite for the get_agent helper function."""
+class TestAgentFunctionAppGetAgent:
+    """Test suite for AgentFunctionApp.get_agent."""
 
-    def test_get_agent_function(self) -> None:
-        """Test get_agent function creates DurableAIAgent."""
+    def test_get_agent_method(self) -> None:
+        """Test get_agent method creates DurableAIAgent for registered agent."""
+        app = _app_with_registered_agents("MyAgent")
         mock_context = Mock()
         mock_context.instance_id = "test-instance-100"
 
-        agent = get_agent(mock_context, "MyAgent")
+        agent = app.get_agent(mock_context, "MyAgent")
 
         assert isinstance(agent, DurableAIAgent)
         assert agent.agent_name == "MyAgent"
         assert agent.context == mock_context
 
+    def test_get_agent_raises_for_unregistered_agent(self) -> None:
+        """Test get_agent raises ValueError when agent is not registered."""
+        app = _app_with_registered_agents("KnownAgent")
+
+        with pytest.raises(ValueError, match=r"Agent 'MissingAgent' is not registered with this app\."):
+            app.get_agent(Mock(), "MissingAgent")
+
 
 class TestOrchestrationIntegration:
     """Integration tests for orchestration scenarios."""
@@ -307,8 +324,8 @@ class TestOrchestrationIntegration:
 
         mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
 
-        # Create agent
-        agent = get_agent(mock_context, "WriterAgent")
+        app = _app_with_registered_agents("WriterAgent")
+        agent = app.get_agent(mock_context, "WriterAgent")
 
         # Create thread
         thread = agent.get_new_thread()
@@ -347,9 +364,9 @@ class TestOrchestrationIntegration:
 
         mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
 
-        # Create multiple agents
-        writer = get_agent(mock_context, "WriterAgent")
-        editor = get_agent(mock_context, "EditorAgent")
+        app = _app_with_registered_agents("WriterAgent", "EditorAgent")
+        writer = app.get_agent(mock_context, "WriterAgent")
+        editor = app.get_agent(mock_context, "EditorAgent")
 
         writer_thread = writer.get_new_thread()
         editor_thread = editor.get_new_thread()

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

@@ -5,12 +5,16 @@ import importlib
 from typing import Any
 
 _IMPORTS: dict[str, tuple[str, str]] = {
+    "AgentCallbackContext": ("agent_framework_azurefunctions", "azurefunctions"),
+    "AgentFunctionApp": ("agent_framework_azurefunctions", "azurefunctions"),
+    "AgentResponseCallbackProtocol": ("agent_framework_azurefunctions", "azurefunctions"),
     "AzureAIAgentClient": ("agent_framework_azure_ai", "azure-ai"),
     "AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", "core"),
     "AzureOpenAIChatClient": ("agent_framework.azure._chat_client", "core"),
     "AzureAISettings": ("agent_framework_azure_ai", "azure-ai"),
     "AzureOpenAISettings": ("agent_framework.azure._shared", "core"),
     "AzureOpenAIResponsesClient": ("agent_framework.azure._responses_client", "core"),
+    "DurableAIAgent": ("agent_framework_azurefunctions", "azurefunctions"),
     "get_entra_auth_token": ("agent_framework.azure._entra_id_authentication", "core"),
 }
 

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

@@ -1,6 +1,12 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from agent_framework_azure_ai import AzureAIAgentClient, AzureAISettings
+from agent_framework_azurefunctions import (
+    AgentCallbackContext,
+    AgentFunctionApp,
+    AgentResponseCallbackProtocol,
+    DurableAIAgent,
+)
 
 from agent_framework.azure._assistants_client import AzureOpenAIAssistantsClient
 from agent_framework.azure._chat_client import AzureOpenAIChatClient
@@ -9,11 +15,15 @@ from agent_framework.azure._responses_client import AzureOpenAIResponsesClient
 from agent_framework.azure._shared import AzureOpenAISettings
 
 __all__ = [
+    "AgentCallbackContext",
+    "AgentFunctionApp",
+    "AgentResponseCallbackProtocol",
     "AzureAIAgentClient",
     "AzureAISettings",
     "AzureOpenAIAssistantsClient",
     "AzureOpenAIChatClient",
     "AzureOpenAIResponsesClient",
     "AzureOpenAISettings",
+    "DurableAIAgent",
     "get_entra_auth_token",
 ]

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

@@ -1,29 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-import importlib
-from typing import Any
-
-PACKAGE_NAME = "agent_framework_azurefunctions"
-PACKAGE_EXTRA = "azurefunctions"
-_IMPORTS = [
-    "AgentCallbackContext",
-    "AgentFunctionApp",
-    "AgentResponseCallbackProtocol",
-    "DurableAIAgent",
-    "get_agent",
-]
-
-
-def __getattr__(name: str) -> Any:
-    if name in _IMPORTS:
-        try:
-            return getattr(importlib.import_module(PACKAGE_NAME), name)
-        except ModuleNotFoundError as exc:
-            raise ModuleNotFoundError(
-                f"The '{PACKAGE_EXTRA}' extra is not installed, please do `pip install agent-framework-{PACKAGE_EXTRA}`"
-            ) from exc
-    raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
-
-
-def __dir__() -> list[str]:
-    return _IMPORTS

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

@@ -1,17 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from agent_framework_azurefunctions import (
-    AgentCallbackContext,
-    AgentFunctionApp,
-    AgentResponseCallbackProtocol,
-    DurableAIAgent,
-    get_agent,
-)
-
-__all__ = [
-    "AgentCallbackContext",
-    "AgentFunctionApp",
-    "AgentResponseCallbackProtocol",
-    "DurableAIAgent",
-    "get_agent",
-]

python/pyproject.toml

@@ -257,6 +257,7 @@ pytest --import-mode=importlib
 --cov-report=term-missing:skip-covered
 --ignore-glob=packages/lab/**
 --ignore-glob=packages/devui/**
+-rs
 -n logical --dist loadfile --dist worksteal
 packages/**/tests
 """
@@ -266,6 +267,7 @@ cmd = """
 pytest --import-mode=importlib
 --ignore-glob=packages/lab/**
 --ignore-glob=packages/devui/**
+-rs
 -n logical --dist loadfile --dist worksteal
 packages/**/tests
 """

python/samples/getting_started/azure_functions/01_single_agent/function_app.py

@@ -8,8 +8,7 @@ Prerequisites: set `AZURE_OPENAI_ENDPOINT` and `AZURE_OPENAI_CHAT_DEPLOYMENT_NAM
 
 from typing import Any
 
-from agent_framework.azure import AzureOpenAIChatClient
-from agent_framework.azurefunctions import AgentFunctionApp
+from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient
 from azure.identity import AzureCliCredential
 # 1. Instantiate the agent with the chosen deployment and instructions.
 def _create_agent() -> Any:

python/samples/getting_started/azure_functions/02_multi_agent/function_app.py

@@ -11,8 +11,7 @@ Prerequisites: set `AZURE_OPENAI_ENDPOINT` and `AZURE_OPENAI_CHAT_DEPLOYMENT_NAM
 import logging
 from typing import Any
 
-from agent_framework.azure import AzureOpenAIChatClient
-from agent_framework.azurefunctions import AgentFunctionApp
+from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient
 from azure.identity import AzureCliCredential
 
 logger = logging.getLogger(__name__)

python/samples/getting_started/azure_functions/03_callbacks/function_app.py

@@ -16,11 +16,14 @@ from typing import Any, DefaultDict
 
 import azure.functions as func
 from agent_framework import AgentRunResponseUpdate
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.azure import (
+    AgentCallbackContext,
+    AgentFunctionApp,
+    AgentResponseCallbackProtocol,
+    AzureOpenAIChatClient,
+)
 from azure.identity import AzureCliCredential
 
-from agent_framework.azurefunctions import AgentFunctionApp, AgentCallbackContext, AgentResponseCallbackProtocol
-
 logger = logging.getLogger(__name__)
 
 

python/samples/getting_started/azure_functions/04_single_agent_orchestration_chaining/function_app.py

@@ -14,10 +14,9 @@ from typing import Any
 
 import azure.durable_functions as df
 import azure.functions as func
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient
 from azure.durable_functions import DurableOrchestrationContext
 from azure.identity import AzureCliCredential
-from agent_framework.azurefunctions import AgentFunctionApp, get_agent
 
 logger = logging.getLogger(__name__)
 
@@ -49,7 +48,7 @@ app = AgentFunctionApp(agents=[_create_writer_agent()], enable_health_check=True
 def single_agent_orchestration(context: DurableOrchestrationContext):
     """Run the writer agent twice on the same thread to mirror chaining behaviour."""
 
-    writer = get_agent(context, WRITER_AGENT_NAME)
+    writer = app.get_agent(context, WRITER_AGENT_NAME)
     writer_thread = writer.get_new_thread()
 
     initial = yield writer.run(

python/samples/getting_started/azure_functions/05_multi_agent_orchestration_concurrency/function_app.py

@@ -14,10 +14,9 @@ from typing import Any
 
 import azure.durable_functions as df
 import azure.functions as func
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient
 from azure.durable_functions import DurableOrchestrationContext
 from azure.identity import AzureCliCredential
-from agent_framework.azurefunctions import AgentFunctionApp, get_agent
 
 logger = logging.getLogger(__name__)
 
@@ -59,8 +58,8 @@ def multi_agent_concurrent_orchestration(context: DurableOrchestrationContext):
     if not prompt or not str(prompt).strip():
         raise ValueError("Prompt is required")
 
-    physicist = get_agent(context, PHYSICIST_AGENT_NAME)
-    chemist = get_agent(context, CHEMIST_AGENT_NAME)
+    physicist = app.get_agent(context, PHYSICIST_AGENT_NAME)
+    chemist = app.get_agent(context, CHEMIST_AGENT_NAME)
 
     physicist_thread = physicist.get_new_thread()
     chemist_thread = chemist.get_new_thread()

python/samples/getting_started/azure_functions/06_multi_agent_orchestration_conditionals/function_app.py

@@ -11,15 +11,14 @@ Functions host."""
 
 import json
 import logging
-from typing import Any, cast
 from collections.abc import Mapping
+from typing import Any, cast
 
 import azure.durable_functions as df
 import azure.functions as func
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient
 from azure.durable_functions import DurableOrchestrationContext
 from azure.identity import AzureCliCredential
-from agent_framework.azurefunctions import AgentFunctionApp, get_agent
 from pydantic import BaseModel, ValidationError
 
 logger = logging.getLogger(__name__)
@@ -85,8 +84,8 @@ def spam_detection_orchestration(context: DurableOrchestrationContext):
     except ValidationError as exc:
         raise ValueError(f"Invalid email payload: {exc}") from exc
 
-    spam_agent = get_agent(context, SPAM_AGENT_NAME)
-    email_agent = get_agent(context, EMAIL_AGENT_NAME)
+    spam_agent = app.get_agent(context, SPAM_AGENT_NAME)
+    email_agent = app.get_agent(context, EMAIL_AGENT_NAME)
 
     spam_thread = spam_agent.get_new_thread()
 

python/samples/getting_started/azure_functions/07_single_agent_orchestration_hitl/function_app.py

@@ -10,16 +10,15 @@ either `AZURE_OPENAI_API_KEY` or sign in with Azure CLI before running `func sta
 
 import json
 import logging
+from collections.abc import Mapping
 from datetime import timedelta
 from typing import Any
-from collections.abc import Mapping
 
 import azure.durable_functions as df
 import azure.functions as func
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient
 from azure.durable_functions import DurableOrchestrationContext
 from azure.identity import AzureCliCredential
-from agent_framework.azurefunctions import AgentFunctionApp, get_agent
 from pydantic import BaseModel, ValidationError
 
 logger = logging.getLogger(__name__)
@@ -92,7 +91,7 @@ def content_generation_hitl_orchestration(context: DurableOrchestrationContext):
     except ValidationError as exc:
         raise ValueError(f"Invalid content generation input: {exc}") from exc
 
-    writer = get_agent(context, WRITER_AGENT_NAME)
+    writer = app.get_agent(context, WRITER_AGENT_NAME)
     writer_thread = writer.get_new_thread()
 
     context.set_custom_status(f"Starting content generation for topic: {payload.topic}")