orchestrator-core 4.5.1a1__py3-none-any.whl → 4.5.2__py3-none-any.whl

orchestrator/__init__.py

@@ -13,7 +13,7 @@
 
 """This is the orchestrator workflow engine."""
 
-__version__ = "4.5.1a1"
+__version__ = "4.5.2"
 
 
 from structlog import get_logger
@@ -25,18 +25,9 @@ logger.info("Starting the orchestrator", version=__version__)
 from orchestrator.llm_settings import llm_settings
 from orchestrator.settings import app_settings
 
-if llm_settings.LLM_ENABLED:
-    try:
-        from importlib import import_module
+if llm_settings.SEARCH_ENABLED or llm_settings.AGENT_ENABLED:
 
-        import_module("pydantic_ai")
-        from orchestrator.agentic_app import AgenticOrchestratorCore as OrchestratorCore
-
-    except ImportError:
-        logger.error(
-            "Unable to import 'pydantic_ai' module, please install the orchestrator with llm dependencies. `pip install orchestrator-core[llm]",
-        )
-        exit(1)
+    from orchestrator.agentic_app import LLMOrchestratorCore as OrchestratorCore
 else:
     from orchestrator.app import OrchestratorCore  # type: ignore[assignment]
 

orchestrator/agentic_app.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
 """The main application module.
 
-This module contains the main `AgenticOrchestratorCore` class for the `FastAPI` backend and
-provides the ability to run the CLI.
+This module contains the main `LLMOrchestratorCore` class for the `FastAPI` backend and
+provides the ability to run the CLI with LLM features (search and/or agent).
 """
 # Copyright 2019-2025 SURF
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,65 +16,84 @@ provides the ability to run the CLI.
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import typer
-from pydantic_ai.models.openai import OpenAIModel
-from pydantic_ai.toolsets import FunctionToolset
 from structlog import get_logger
 
 from orchestrator.app import OrchestratorCore
 from orchestrator.cli.main import app as cli_app
 from orchestrator.llm_settings import LLMSettings, llm_settings
 
+if TYPE_CHECKING:
+    from pydantic_ai.models.openai import OpenAIModel
+    from pydantic_ai.toolsets import FunctionToolset
+
 logger = get_logger(__name__)
 
 
-class AgenticOrchestratorCore(OrchestratorCore):
+class LLMOrchestratorCore(OrchestratorCore):
     def __init__(
         self,
         *args: Any,
-        llm_model: OpenAIModel | str = "gpt-4o-mini",
         llm_settings: LLMSettings = llm_settings,
-        agent_tools: list[FunctionToolset] | None = None,
+        agent_model: "OpenAIModel | str | None" = None,
+        agent_tools: "list[FunctionToolset] | None" = None,
         **kwargs: Any,
     ) -> None:
-        """Initialize the `AgenticOrchestratorCore` class.
+        """Initialize the `LLMOrchestratorCore` class.
 
-        This class takes the same arguments as the `OrchestratorCore` class.
+        This class extends `OrchestratorCore` with LLM features (search and agent).
+        It runs the search migration and mounts the agent endpoint based on feature flags.
 
         Args:
             *args: All the normal arguments passed to the `OrchestratorCore` class.
-            llm_model: An OpenAI model class or string, not limited to OpenAI models (gpt-4o-mini etc)
             llm_settings: A class of settings for the LLM
+            agent_model: Override the agent model (defaults to llm_settings.AGENT_MODEL)
             agent_tools: A list of tools that can be used by the agent
             **kwargs: Additional arguments passed to the `OrchestratorCore` class.
 
         Returns:
             None
         """
-        self.llm_model = llm_model
-        self.agent_tools = agent_tools
         self.llm_settings = llm_settings
+        self.agent_model = agent_model or llm_settings.AGENT_MODEL
+        self.agent_tools = agent_tools
 
         super().__init__(*args, **kwargs)
 
-        logger.info("Mounting the agent")
-        self.register_llm_integration()
-
-    def register_llm_integration(self) -> None:
-        """Mount the Agent endpoint.
-
-        This helper mounts the agent endpoint on the application.
-
-        Returns:
-            None
-
-        """
-        from orchestrator.search.agent import build_agent_app
-
-        agent_app = build_agent_app(self.llm_model, self.agent_tools)
-        self.mount("/agent", agent_app)
+        # Run search migration if search or agent is enabled
+        if self.llm_settings.SEARCH_ENABLED or self.llm_settings.AGENT_ENABLED:
+            logger.info("Running search migration")
+            try:
+                from orchestrator.db import db
+                from orchestrator.search.llm_migration import run_migration
+
+                with db.engine.begin() as connection:
+                    run_migration(connection)
+            except ImportError as e:
+                logger.error(
+                    "Unable to run search migration. Please install search dependencies: "
+                    "`pip install orchestrator-core[search]`",
+                    error=str(e),
+                )
+                raise
+
+        # Mount agent endpoint if agent is enabled
+        if self.llm_settings.AGENT_ENABLED:
+            logger.info("Initializing agent features", model=self.agent_model)
+            try:
+                from orchestrator.search.agent import build_agent_router
+
+                agent_app = build_agent_router(self.agent_model, self.agent_tools)
+                self.mount("/agent", agent_app)
+            except ImportError as e:
+                logger.error(
+                    "Unable to initialize agent features. Please install agent dependencies: "
+                    "`pip install orchestrator-core[agent]`",
+                    error=str(e),
+                )
+                raise
 
 
 main_typer_app = typer.Typer()

orchestrator/api/api_v1/api.py

@@ -76,20 +76,22 @@ api_router.include_router(user.router, prefix="/user", tags=["Core", "User"], de
 api_router.include_router(
     settings.router, prefix="/settings", tags=["Core", "Settings"], dependencies=[Depends(authorize)]
 )
-api_router.include_router(settings.ws_router, prefix="/settings", tags=["Core", "Settings"])
+api_router.include_router(
+    settings.ws_router, prefix="/settings", tags=["Core", "Settings"]
+)  # Auth on the websocket is handled in the Websocket Manager
 api_router.include_router(health.router, prefix="/health", tags=["Core"])
 api_router.include_router(
     translations.router,
     prefix="/translations",
     tags=["Core", "Translations"],
 )
-api_router.include_router(ws.router, prefix="/ws", tags=["Core", "Events"])
+api_router.include_router(
+    ws.router, prefix="/ws", tags=["Core", "Events"]
+)  # Auth on the websocket is handled in the Websocket Manager
 
-if llm_settings.LLM_ENABLED:
+if llm_settings.SEARCH_ENABLED:
     from orchestrator.api.api_v1.endpoints import search
 
     api_router.include_router(
-        search.router,
-        prefix="/search",
-        tags=["Core", "Search"],
+        search.router, prefix="/search", tags=["Core", "Search"], dependencies=[Depends(authorize)]
     )

orchestrator/api/api_v1/endpoints/processes.py

@@ -244,6 +244,8 @@ def continue_awaiting_process_endpoint(
         continue_awaiting_process(process, token=token, input_data=json_data, broadcast_func=broadcast_func)
     except AssertionError as e:
         raise_status(HTTPStatus.NOT_FOUND, str(e))
+    except ValueError as e:
+        raise_status(HTTPStatus.BAD_REQUEST, str(e))
 
 
 @router.post(

orchestrator/api/api_v1/endpoints/search.py

@@ -1,3 +1,16 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any, Literal, overload
 
 from fastapi import APIRouter, HTTPException, Query, status
@@ -11,6 +24,7 @@ from orchestrator.db import (
     db,
 )
 from orchestrator.domain.base import SubscriptionModel
+from orchestrator.domain.context_cache import cache_subscription_models
 from orchestrator.schemas.search import (
     PageInfoSchema,
     PathsResponse,
@@ -179,17 +193,22 @@ async def search_subscriptions(
     page_info = PageInfoSchema(has_next_page=has_next_page, next_page_cursor=next_page_cursor)
 
     search_info_map = {res.entity_id: res for res in search_response.results}
-    results_data = []
-    for sub_id, search_info in search_info_map.items():
-        subscription_model = SubscriptionModel.from_subscription(sub_id)
-        sub_data = subscription_model.model_dump(exclude_unset=False)
-        search_result_item = SubscriptionSearchResult(
-            subscription=format_special_types(sub_data),
+
+    with cache_subscription_models():
+        subscriptions_data = {
+            sub_id: SubscriptionModel.from_subscription(sub_id).model_dump(exclude_unset=False)
+            for sub_id in search_info_map
+        }
+
+    results_data = [
+        SubscriptionSearchResult(
+            subscription=format_special_types(subscriptions_data[sub_id]),
             score=search_info.score,
             perfect_match=search_info.perfect_match,
             matching_field=search_info.matching_field,
         )
-        results_data.append(search_result_item)
+        for sub_id, search_info in search_info_map.items()
+    ]
 
     return SearchResultsSchema(data=results_data, page_info=page_info, search_metadata=search_response.metadata)
 

orchestrator/cli/main.py

@@ -25,8 +25,8 @@ app.add_typer(scheduler.app, name="scheduler", help="Access all the scheduler fu
 app.add_typer(database.app, name="db", help="Interact with the application database")
 app.add_typer(generate.app, name="generate", help="Generate products, workflows and other artifacts")
 
-if llm_settings.LLM_ENABLED:
-    from orchestrator.cli import index_llm, resize_embedding, search_explore, speedtest
+if llm_settings.SEARCH_ENABLED:
+    from orchestrator.cli.search import index_llm, resize_embedding, search_explore, speedtest
 
     app.add_typer(index_llm.app, name="index", help="(Re-)Index the search table.")
     app.add_typer(search_explore.app, name="search", help="Try out different search types.")

orchestrator/cli/search/__init__.py

@@ -0,0 +1,32 @@
+# Copyright 2019-2020 SURF.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typer
+
+from orchestrator.cli.search import index_llm, resize_embedding, search_explore, speedtest
+
+
+def register_commands(app: typer.Typer) -> None:
+    """Register all LLM/search related commands to the main app."""
+    app.add_typer(index_llm.app, name="index", help="(Re-)Index the search table.")
+    app.add_typer(search_explore.app, name="search", help="Try out different search types.")
+    app.add_typer(
+        resize_embedding.app,
+        name="embedding",
+        help="Resize the vector dimension of the embedding column in the search table.",
+    )
+    app.add_typer(
+        speedtest.app,
+        name="speedtest",
+        help="Search performance testing and analysis.",
+    )

orchestrator/devtools/populator.py

@@ -371,6 +371,22 @@ class Populator:
         self.log.info("Started modify workflow")
         return self._start_workflow(workflow_name, subscription_id=subscription_id, **kwargs)
 
+    def start_reconcile_workflow(self, workflow_name: str, subscription_id: UUIDstr | UUID, **kwargs: Any) -> UUIDstr:
+        """Start a reconcile workflow for the provided name and subscription_id.
+
+        Args:
+            workflow_name: workflow name
+            subscription_id: uuid of the subscription you want to modify
+            kwargs: values to be used as form input
+
+        Returns: the process_id of the workflow process
+
+        """
+        subscription_id = str(subscription_id)
+        self.log = self.log.bind(subscription_id=subscription_id)
+        self.log.info("Started reconcile workflow")
+        return self._start_workflow(workflow_name, subscription_id=subscription_id, **kwargs)
+
     def start_verify_workflow(self, workflow_name: str, subscription_id: UUIDstr | UUID) -> UUIDstr:
         subscription_id = str(subscription_id)
         self.log = self.log.bind(subscription_id=subscription_id)

orchestrator/domain/base.py

@@ -614,9 +614,7 @@ class ProductBlockModel(DomainModel):
         product_blocks_in_model = cls._get_depends_on_product_block_types()
         product_blocks_types_in_model = get_depends_on_product_block_type_list(product_blocks_in_model)
 
-        product_blocks_in_model = set(
-            flatten(map(attrgetter("__names__"), product_blocks_types_in_model))
-        )  # type: ignore
+        product_blocks_in_model = set(flatten(map(attrgetter("__names__"), product_blocks_types_in_model)))  # type: ignore
 
         missing_product_blocks_in_db = product_blocks_in_model - product_blocks_in_db  # type: ignore
         missing_product_blocks_in_model = product_blocks_in_db - product_blocks_in_model  # type: ignore
@@ -1084,9 +1082,7 @@ class SubscriptionModel(DomainModel):
         product_blocks_in_model = cls._get_depends_on_product_block_types()
         product_blocks_types_in_model = get_depends_on_product_block_type_list(product_blocks_in_model)
 
-        product_blocks_in_model = set(
-            flatten(map(attrgetter("__names__"), product_blocks_types_in_model))
-        )  # type: ignore
+        product_blocks_in_model = set(flatten(map(attrgetter("__names__"), product_blocks_types_in_model)))  # type: ignore
 
         missing_product_blocks_in_db = product_blocks_in_model - product_blocks_in_db  # type: ignore
         missing_product_blocks_in_model = product_blocks_in_db - product_blocks_in_model  # type: ignore
@@ -1294,7 +1290,6 @@ class SubscriptionModel(DomainModel):
     # Some common functions shared by from_other_product and from_subscription
     @classmethod
     def _get_subscription(cls: type[S], subscription_id: UUID | UUIDstr) -> SubscriptionTable | None:
-
         if not isinstance(subscription_id, UUID | UUIDstr):
             raise TypeError(f"subscription_id is of type {type(subscription_id)} instead of UUID | UUIDstr")
 

orchestrator/domain/lifecycle.py

@@ -16,11 +16,16 @@ from typing import TYPE_CHECKING, TypeVar
 import strawberry
 import structlog
 
+from orchestrator.settings import LifecycleValidationMode, app_settings
 from orchestrator.types import SubscriptionLifecycle
 from pydantic_forms.types import strEnum
 
 if TYPE_CHECKING:
-    from orchestrator.domain.base import DomainModel
+    from orchestrator.domain.base import DomainModel, SubscriptionModel
+else:
+    SubscriptionModel = None
+    DomainModel = None
+T = TypeVar("T", bound=SubscriptionModel)
 
 logger = structlog.get_logger(__name__)
 
@@ -71,9 +76,21 @@ def validate_lifecycle_status(
         )
 
 
-if TYPE_CHECKING:
-    from orchestrator.domain.base import DomainModel, SubscriptionModel
-else:
-    SubscriptionModel = None
-    DomainModel = None
-T = TypeVar("T", bound=SubscriptionModel)
+def validate_subscription_model_product_type(
+    subscription: SubscriptionModel,
+    validation_mode: LifecycleValidationMode = app_settings.LIFECYCLE_VALIDATION_MODE,
+) -> None:
+    """Validate that a subscription model has been instantiated with the correct product type class for its lifecycle status."""
+
+    actual_class = subscription.__class__
+    expected_class = lookup_specialized_type(actual_class, subscription.status)
+
+    if actual_class != expected_class:
+        msg = f"Subscription of type {actual_class} should use {expected_class} for lifecycle status '{subscription.status}'"
+        if validation_mode == LifecycleValidationMode.STRICT:
+            logger.error(msg)
+            raise ValueError(msg)
+        if validation_mode == LifecycleValidationMode.LOOSE:
+            logger.warning(msg)
+        elif validation_mode == LifecycleValidationMode.IGNORED:
+            pass

orchestrator/llm_settings.py

@@ -18,7 +18,10 @@ logger = get_logger(__name__)
 
 
 class LLMSettings(BaseSettings):
-    LLM_ENABLED: bool = False  # Default to false
+    # Feature flags for LLM functionality
+    SEARCH_ENABLED: bool = False  # Enable search/indexing with embeddings
+    AGENT_ENABLED: bool = False  # Enable agentic functionality
+
     # Pydantic-ai Agent settings
     AGENT_MODEL: str = "gpt-4o-mini"  # See pydantic-ai docs for supported models.
     AGENT_MODEL_VERSION: str = "2025-01-01-preview"
@@ -30,16 +33,19 @@ class LLMSettings(BaseSettings):
         0.1, description="Safety margin as a percentage (e.g., 0.1 for 10%) for token budgeting.", ge=0, le=1
     )
 
-    # The following settings are only needed for local models.
+    # The following settings are only needed for local models or system constraints.
     # By default, they are set conservative assuming a small model like All-MiniLM-L6-V2.
     OPENAI_BASE_URL: str | None = None
     EMBEDDING_FALLBACK_MAX_TOKENS: int | None = 512
-    EMBEDDING_MAX_BATCH_SIZE: int | None = 32
+    EMBEDDING_MAX_BATCH_SIZE: int | None = None
 
     # General LiteLLM settings
     LLM_MAX_RETRIES: int = 3
     LLM_TIMEOUT: int = 30
 
+    # Toggle creation of extensions
+    LLM_FORCE_EXTENTION_MIGRATION: bool = False
+
     @field_validator("EMBEDDING_MODEL")
     def validate_embedding_model_format(cls, v: str) -> str:
         """Validate that embedding model is in 'vendor/model' format."""

orchestrator/log_config.py

@@ -44,5 +44,6 @@ LOGGER_OVERRIDES = dict(
         logger_config("orchestrator.graphql.autoregistration"),
         logger_config("sqlalchemy.engine", default_level="WARNING"),
         logger_config("uvicorn"),
+        logger_config("LiteLLM", default_level="WARNING"),
     ]
 )

orchestrator/migrations/helpers.py

@@ -155,7 +155,7 @@ def create_workflow(conn: sa.engine.Connection, workflow: dict) -> None:
         conn: DB connection as available in migration main file.
         workflow: Dict with data for a new workflow.
             name: Name of the workflow.
-            target: Target of the workflow ("CREATE", "MODIFY", "TERMINATE", "SYSTEM")
+            target: Target of the workflow ("CREATE", "MODIFY", "RECONCILE", "TERMINATE", "SYSTEM")
             description: Description of the workflow.
             product_type: Product type to add the workflow to.
 
@@ -166,12 +166,16 @@ def create_workflow(conn: sa.engine.Connection, workflow: dict) -> None:
             "is_task": False,
             "description": "workflow description",
             "product_type": "product_type",
+            "product_tag": "product_tag",
         }
         >>> create_workflow(conn, workflow)
     """
     if not workflow.get("is_task", False):
         workflow["is_task"] = False
 
+    if not workflow.get("product_tag"):
+        workflow["product_tag"] = None
+
     if has_table_column(table_name="workflows", column_name="is_task", conn=conn):
         query = """
                 WITH new_workflow AS (
@@ -186,6 +190,7 @@ def create_workflow(conn: sa.engine.Connection, workflow: dict) -> None:
                 FROM products AS p
                          CROSS JOIN new_workflow AS nw
                 WHERE p.product_type = :product_type
+                AND (:product_tag IS NULL OR p.tag = :product_tag)
                 ON CONFLICT DO NOTHING
                 """
     else:
@@ -203,6 +208,7 @@ def create_workflow(conn: sa.engine.Connection, workflow: dict) -> None:
                 FROM products AS p
                          CROSS JOIN new_workflow AS nw
                 WHERE p.product_type = :product_type
+                AND (:product_tag IS NULL OR p.tag = :product_tag)
                 ON CONFLICT DO NOTHING
                 """
 

orchestrator/schemas/search.py

@@ -1,3 +1,16 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from datetime import datetime
 from typing import Any, Generic, TypeVar
 from uuid import UUID

orchestrator/schemas/workflow.py

@@ -60,6 +60,7 @@ class SubscriptionWorkflowListsSchema(OrchestratorBaseModel):
     modify: list[WorkflowListItemSchema]
     terminate: list[WorkflowListItemSchema]
     system: list[WorkflowListItemSchema]
+    reconcile: list[WorkflowListItemSchema]
     validate_: list[WorkflowListItemSchema] = Field(default_factory=list, alias="validate")
 
 

orchestrator/search/agent/__init__.py

@@ -1,8 +1,21 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 # This module requires: pydantic-ai==0.7.0, ag-ui-protocol>=0.1.8
 
 
-from orchestrator.search.agent.agent import build_agent_app
+from orchestrator.search.agent.agent import build_agent_router
 
 __all__ = [
-    "build_agent_app",
+    "build_agent_router",
 ]

orchestrator/search/agent/agent.py

@@ -1,13 +1,26 @@
+# Copyright 2019-2025 SURF, GÉANT.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from typing import Any
 
 import structlog
-from fastapi import FastAPI, HTTPException
-from pydantic_ai.ag_ui import StateDeps
+from fastapi import APIRouter, HTTPException, Request
+from pydantic_ai.ag_ui import StateDeps, handle_ag_ui_request
 from pydantic_ai.agent import Agent
 from pydantic_ai.models.openai import OpenAIModel
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.toolsets import FunctionToolset
-from starlette.types import ASGIApp
+from starlette.responses import Response
 
 from orchestrator.search.agent.prompts import get_base_instructions, get_dynamic_instructions
 from orchestrator.search.agent.state import SearchState
@@ -16,17 +29,9 @@ from orchestrator.search.agent.tools import search_toolset
 logger = structlog.get_logger(__name__)
 
 
-def _disabled_agent_app(reason: str) -> FastAPI:
-    app = FastAPI(title="Agent disabled")
-
-    @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
-    async def _disabled(path: str) -> None:
-        raise HTTPException(status_code=503, detail=f"Agent disabled: {reason}")
-
-    return app
-
+def build_agent_router(model: str | OpenAIModel, toolsets: list[FunctionToolset[Any]] | None = None) -> APIRouter:
+    router = APIRouter()
 
-def build_agent_app(model: str | OpenAIModel, toolsets: list[FunctionToolset[Any]] | None = None) -> ASGIApp:
     try:
         toolsets = toolsets + [search_toolset] if toolsets else [search_toolset]
 
@@ -41,7 +46,17 @@ def build_agent_app(model: str | OpenAIModel, toolsets: list[FunctionToolset[Any
         agent.instructions(get_base_instructions)
         agent.instructions(get_dynamic_instructions)
 
-        return agent.to_ag_ui(deps=StateDeps(SearchState()))
+        @router.post("/")
+        async def agent_endpoint(request: Request) -> Response:
+            return await handle_ag_ui_request(agent, request, deps=StateDeps(SearchState()))
+
+        return router
     except Exception as e:
         logger.error("Agent init failed; serving disabled stub.", error=str(e))
-        return _disabled_agent_app(str(e))
+        error_msg = f"Agent disabled: {str(e)}"
+
+        @router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
+        async def _disabled(path: str) -> None:
+            raise HTTPException(status_code=503, detail=error_msg)
+
+        return router