dao-ai 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

dao_ai/apps/__init__.py

@@ -0,0 +1,24 @@
+"""
+Databricks Apps deployment module for dao-ai.
+
+This subpackage contains all modules related to deploying dao-ai agents
+as Databricks Apps or Model Serving endpoints.
+
+Modules:
+    handlers: MLflow AgentServer request handlers (@invoke, @stream)
+    server: Entry point for Databricks Apps deployment
+    resources: Databricks App resource configuration generation
+    model_serving: Entry point for Databricks Model Serving deployment
+"""
+
+from dao_ai.apps.resources import (
+    generate_app_resources,
+    generate_app_yaml,
+    generate_sdk_resources,
+)
+
+__all__ = [
+    "generate_app_resources",
+    "generate_app_yaml",
+    "generate_sdk_resources",
+]

dao_ai/{app_server.py → apps/handlers.py}

@@ -1,20 +1,12 @@
 """
-App server module for running dao-ai agents as Databricks Apps.
+Agent request handlers for MLflow AgentServer.
 
-This module provides the entry point for deploying dao-ai agents as Databricks Apps
-using MLflow's AgentServer. It follows the same pattern as agent_as_code.py but
-uses the AgentServer for the Databricks Apps runtime.
+This module defines the invoke and stream handlers that are registered
+with the MLflow AgentServer. These handlers delegate to the ResponsesAgent
+created from the dao-ai configuration.
 
-Configuration Loading:
-    The config path is specified via the DAO_AI_CONFIG_PATH environment variable,
-    or defaults to model_config.yaml in the current directory.
-
-Usage:
-    # With environment variable
-    DAO_AI_CONFIG_PATH=/path/to/config.yaml python -m dao_ai.app_server
-
-    # With default model_config.yaml in current directory
-    python -m dao_ai.app_server
+The handlers use async methods (apredict, apredict_stream) to be compatible
+with both Databricks Model Serving and Databricks Apps environments.
 """
 
 import os
@@ -22,8 +14,7 @@ from typing import AsyncGenerator
 
 import mlflow
 from dotenv import load_dotenv
-from mlflow.genai.agent_server import AgentServer, invoke, stream
-from mlflow.pyfunc import ResponsesAgent
+from mlflow.genai.agent_server import invoke, stream
 from mlflow.types.responses import (
     ResponsesAgentRequest,
     ResponsesAgentResponse,
@@ -32,6 +23,7 @@ from mlflow.types.responses import (
 
 from dao_ai.config import AppConfig
 from dao_ai.logging import configure_logging
+from dao_ai.models import LanggraphResponsesAgent
 
 # Load environment variables from .env.local if it exists
 load_dotenv(dotenv_path=".env.local", override=True)
@@ -51,53 +43,42 @@ config: AppConfig = AppConfig.from_file(config_path)
 if config.app and config.app.log_level:
     configure_logging(level=config.app.log_level)
 
-# Create the ResponsesAgent
-_responses_agent: ResponsesAgent = config.as_responses_agent()
+# Create the ResponsesAgent - cast to LanggraphResponsesAgent to access async methods
+_responses_agent: LanggraphResponsesAgent = config.as_responses_agent()  # type: ignore[assignment]
 
 
 @invoke()
-def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
+async def non_streaming(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
     """
     Handle non-streaming requests by delegating to the ResponsesAgent.
 
+    Uses the async apredict() method for compatibility with both
+    Model Serving and Apps environments.
+
     Args:
         request: The incoming ResponsesAgentRequest
 
     Returns:
         ResponsesAgentResponse with the complete output
     """
-    return _responses_agent.predict(request)
+    return await _responses_agent.apredict(request)
 
 
 @stream()
-def streaming(
+async def streaming(
     request: ResponsesAgentRequest,
 ) -> AsyncGenerator[ResponsesAgentStreamEvent, None]:
     """
     Handle streaming requests by delegating to the ResponsesAgent.
 
+    Uses the async apredict_stream() method for compatibility with both
+    Model Serving and Apps environments.
+
     Args:
         request: The incoming ResponsesAgentRequest
 
     Yields:
         ResponsesAgentStreamEvent objects as they are generated
     """
-    # The predict_stream method returns a generator, convert to async generator
-    for event in _responses_agent.predict_stream(request):
+    async for event in _responses_agent.apredict_stream(request):
         yield event
-
-
-# Create the AgentServer instance
-agent_server = AgentServer("ResponsesAgent", enable_chat_proxy=True)
-
-# Define the app as a module level variable to enable multiple workers
-app = agent_server.app
-
-
-def main() -> None:
-    """Entry point for running the agent server."""
-    agent_server.run(app_import_string="dao_ai.app_server:app")
-
-
-if __name__ == "__main__":
-    main()

dao_ai/apps/model_serving.py

@@ -0,0 +1,29 @@
+# Apply nest_asyncio FIRST before any other imports
+# This allows dao-ai's async/sync patterns to work in Model Serving
+# where there may already be an event loop running (e.g., notebook context)
+import nest_asyncio
+
+nest_asyncio.apply()
+
+import mlflow  # noqa: E402
+from mlflow.models import ModelConfig  # noqa: E402
+from mlflow.pyfunc import ResponsesAgent  # noqa: E402
+
+from dao_ai.config import AppConfig  # noqa: E402
+from dao_ai.logging import configure_logging  # noqa: E402
+
+mlflow.set_registry_uri("databricks-uc")
+mlflow.set_tracking_uri("databricks")
+
+mlflow.langchain.autolog()
+
+model_config: ModelConfig = ModelConfig()
+config: AppConfig = AppConfig(**model_config.to_dict())
+
+log_level: str = config.app.log_level
+
+configure_logging(level=log_level)
+
+app: ResponsesAgent = config.as_responses_agent()
+
+mlflow.models.set_model(app)