dao-ai 0.0.26__py3-none-any.whl → 0.0.27__py3-none-any.whl

dao_ai/config.py

@@ -397,6 +397,13 @@ class VectorSearchEndpoint(BaseModel):
     name: str
     type: VectorSearchEndpointType = VectorSearchEndpointType.STANDARD
 
+    @field_serializer("type")
+    def serialize_type(self, value: VectorSearchEndpointType) -> str:
+        """Ensure enum is serialized to string value."""
+        if isinstance(value, VectorSearchEndpointType):
+            return value.value
+        return str(value)
+
 
 class IndexModel(BaseModel, HasFullName, IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -998,22 +1005,120 @@ class TransportType(str, Enum):
 class McpFunctionModel(BaseFunctionModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     type: Literal[FunctionType.MCP] = FunctionType.MCP
-
     transport: TransportType = TransportType.STREAMABLE_HTTP
     command: Optional[str] = "python"
     url: Optional[AnyVariable] = None
-    connection: Optional[ConnectionModel] = None
     headers: dict[str, AnyVariable] = Field(default_factory=dict)
     args: list[str] = Field(default_factory=list)
     pat: Optional[AnyVariable] = None
     client_id: Optional[AnyVariable] = None
     client_secret: Optional[AnyVariable] = None
     workspace_host: Optional[AnyVariable] = None
+    connection: Optional[ConnectionModel] = None
+    functions: Optional[SchemaModel] = None
+    genie_room: Optional[GenieRoomModel] = None
+    sql: Optional[bool] = None
+    vector_search: Optional[VectorStoreModel] = None
 
     @property
     def full_name(self) -> str:
         return self.name
 
+    def _get_workspace_host(self) -> str:
+        """
+        Get the workspace host, either from config or from workspace client.
+
+        If connection is provided, uses its workspace client.
+        Otherwise, falls back to creating a new workspace client.
+
+        Returns:
+            str: The workspace host URL without trailing slash
+        """
+        from databricks.sdk import WorkspaceClient
+
+        # Try to get workspace_host from config
+        workspace_host: str | None = (
+            value_of(self.workspace_host) if self.workspace_host else None
+        )
+
+        # If no workspace_host in config, get it from workspace client
+        if not workspace_host:
+            # Use connection's workspace client if available
+            if self.connection:
+                workspace_host = self.connection.workspace_client.config.host
+            else:
+                # Create a default workspace client
+                w: WorkspaceClient = WorkspaceClient()
+                workspace_host = w.config.host
+
+        # Remove trailing slash
+        return workspace_host.rstrip("/")
+
+    @property
+    def mcp_url(self) -> str:
+        """
+        Get the MCP URL for this function.
+
+        Returns the URL based on the configured source:
+        - If url is set, returns it directly
+        - If connection is set, constructs URL from connection
+        - If genie_room is set, constructs Genie MCP URL
+        - If sql is set, constructs DBSQL MCP URL (serverless)
+        - If vector_search is set, constructs Vector Search MCP URL
+        - If functions is set, constructs UC Functions MCP URL
+
+        URL patterns (per https://docs.databricks.com/aws/en/generative-ai/mcp/managed-mcp):
+        - Genie: https://{host}/api/2.0/mcp/genie/{space_id}
+        - DBSQL: https://{host}/api/2.0/mcp/sql (serverless, workspace-level)
+        - Vector Search: https://{host}/api/2.0/mcp/vector-search/{catalog}/{schema}
+        - UC Functions: https://{host}/api/2.0/mcp/functions/{catalog}/{schema}
+        - Connection: https://{host}/api/2.0/mcp/external/{connection_name}
+        """
+        # Direct URL provided
+        if self.url:
+            return self.url
+
+        # Get workspace host (from config, connection, or default workspace client)
+        workspace_host: str = self._get_workspace_host()
+
+        # UC Connection
+        if self.connection:
+            connection_name: str = self.connection.name
+            return f"{workspace_host}/api/2.0/mcp/external/{connection_name}"
+
+        # Genie Room
+        if self.genie_room:
+            space_id: str = value_of(self.genie_room.space_id)
+            return f"{workspace_host}/api/2.0/mcp/genie/{space_id}"
+
+        # DBSQL MCP server (serverless, workspace-level)
+        if self.sql:
+            return f"{workspace_host}/api/2.0/mcp/sql"
+
+        # Vector Search
+        if self.vector_search:
+            if (
+                not self.vector_search.index
+                or not self.vector_search.index.schema_model
+            ):
+                raise ValueError(
+                    "vector_search must have an index with a schema (catalog/schema) configured"
+                )
+            catalog: str = self.vector_search.index.schema_model.catalog_name
+            schema: str = self.vector_search.index.schema_model.schema_name
+            return f"{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}"
+
+        # UC Functions MCP server
+        if self.functions:
+            catalog: str = self.functions.catalog_name
+            schema: str = self.functions.schema_name
+            return f"{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}"
+
+        raise ValueError(
+            "No URL source configured. Provide one of: url, connection, genie_room, "
+            "sql, vector_search, or functions"
+        )
+
     @field_serializer("transport")
     def serialize_transport(self, value) -> str:
         if isinstance(value, TransportType):
@@ -1021,32 +1126,56 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
         return str(value)
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self):
-        if self.transport == TransportType.STREAMABLE_HTTP and not (
-            self.url or self.connection
-        ):
-            raise ValueError(
-                "url or connection must be provided for STREAMABLE_HTTP transport"
-            )
-        if self.transport == TransportType.STDIO and not self.command:
-            raise ValueError("command must not be provided for STDIO transport")
-        if self.transport == TransportType.STDIO and not self.args:
-            raise ValueError("args must not be provided for STDIO transport")
+    def validate_mutually_exclusive(self) -> "McpFunctionModel":
+        """Validate that exactly one URL source is provided."""
+        # Count how many URL sources are provided
+        url_sources: list[tuple[str, Any]] = [
+            ("url", self.url),
+            ("connection", self.connection),
+            ("genie_room", self.genie_room),
+            ("sql", self.sql),
+            ("vector_search", self.vector_search),
+            ("functions", self.functions),
+        ]
+
+        provided_sources: list[str] = [
+            name for name, value in url_sources if value is not None
+        ]
+
+        if self.transport == TransportType.STREAMABLE_HTTP:
+            if len(provided_sources) == 0:
+                raise ValueError(
+                    "For STREAMABLE_HTTP transport, exactly one of the following must be provided: "
+                    "url, connection, genie_room, sql, vector_search, or functions"
+                )
+            if len(provided_sources) > 1:
+                raise ValueError(
+                    f"For STREAMABLE_HTTP transport, only one URL source can be provided. "
+                    f"Found: {', '.join(provided_sources)}. "
+                    f"Please provide only one of: url, connection, genie_room, sql, vector_search, or functions"
+                )
+
+        if self.transport == TransportType.STDIO:
+            if not self.command:
+                raise ValueError("command must be provided for STDIO transport")
+            if not self.args:
+                raise ValueError("args must be provided for STDIO transport")
+
         return self
 
     @model_validator(mode="after")
-    def update_url(self):
+    def update_url(self) -> "McpFunctionModel":
         self.url = value_of(self.url)
         return self
 
     @model_validator(mode="after")
-    def update_headers(self):
+    def update_headers(self) -> "McpFunctionModel":
         for key, value in self.headers.items():
             self.headers[key] = value_of(value)
         return self
 
     @model_validator(mode="after")
-    def validate_auth_methods(self):
+    def validate_auth_methods(self) -> "McpFunctionModel":
         oauth_fields: Sequence[Any] = [
             self.client_id,
             self.client_secret,
@@ -1062,10 +1191,7 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
                 "Please provide either OAuth credentials or user credentials."
             )
 
-        if (has_oauth or has_user_auth) and not self.workspace_host:
-            raise ValueError(
-                "Workspace host must be provided when using OAuth or user credentials."
-            )
+        # Note: workspace_host is optional - it will be derived from workspace client if not provided
 
         return self
 

dao_ai/graph.py

@@ -62,11 +62,19 @@ def _handoffs_for_agent(agent: AgentModel, config: AppConfig) -> Sequence[BaseTo
         logger.debug(
             f"Creating handoff tool from agent {agent.name} to {handoff_to_agent.name}"
         )
+
+        # Use handoff_prompt if provided, otherwise create default description
+        handoff_description = handoff_to_agent.handoff_prompt or (
+            handoff_to_agent.description
+            if handoff_to_agent.description
+            else "general assistance and questions"
+        )
+
         handoff_tools.append(
             swarm_handoff_tool(
                 agent_name=handoff_to_agent.name,
                 description=f"Ask {handoff_to_agent.name} for help with: "
-                + handoff_to_agent.handoff_prompt,
+                + handoff_description,
             )
         )
     return handoff_tools
@@ -87,10 +95,17 @@ def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
                 additional_tools=[],
             )
         )
+        # Use handoff_prompt if provided, otherwise create default description
+        handoff_description = registered_agent.handoff_prompt or (
+            registered_agent.description
+            if registered_agent.description
+            else f"General assistance with {registered_agent.name} related tasks"
+        )
+
         tools.append(
             supervisor_handoff_tool(
                 agent_name=registered_agent.name,
-                description=registered_agent.handoff_prompt,
+                description=handoff_description,
             )
         )
 

dao_ai/providers/databricks.py

@@ -331,7 +331,7 @@ class DatabricksProvider(ServiceProvider):
             logged_agent_info: ModelInfo = mlflow.pyfunc.log_model(
                 python_model=model_path.as_posix(),
                 code_paths=code_paths,
-                model_config=config.model_dump(by_alias=True),
+                model_config=config.model_dump(mode="json", by_alias=True),
                 name="agent",
                 pip_requirements=pip_requirements,
                 input_example=input_example,
@@ -1280,7 +1280,9 @@ class DatabricksProvider(ServiceProvider):
         # Get agent and prompt (prompt is guaranteed to be set by validator)
         agent_model: AgentModel = optimization.agent
         prompt: PromptModel = optimization.prompt  # type: ignore[assignment]
+        agent_model.prompt = prompt.uri
 
+        print(f"prompt={agent_model.prompt}")
         # Log the prompt URI scheme being used
         # Supports three schemes:
         # 1. Specific version: "prompts:/qa/1" (when version is specified)
@@ -1350,20 +1352,8 @@ class DatabricksProvider(ServiceProvider):
         # DO NOT overwrite with prompt_version.uri as that uses fallback logic
         logger.debug(f"Optimizing prompt: {prompt_uri}")
 
-        # Create the agent once outside the predict function to avoid repeated construction
-        # and to ensure dao_ai modules are available in the current environment
-        logger.debug(f"Creating ResponsesAgent for optimization from agent '{agent_model.name}'")
-        try:
-            agent: ResponsesAgent = agent_model.as_responses_agent()
-            logger.info(f"Successfully created ResponsesAgent for '{agent_model.name}'")
-        except Exception as e:
-            logger.error(f"Failed to create ResponsesAgent: {e}")
-            logger.error("Ensure dao_ai package is properly installed and accessible")
-            raise RuntimeError(
-                f"Failed to create ResponsesAgent for prompt optimization: {e}. "
-                "This may indicate dao_ai is not properly installed or accessible."
-            ) from e
-        
+        agent: ResponsesAgent = agent_model.as_responses_agent()
+
         # Create predict function that will be optimized
         def predict_fn(**inputs: dict[str, Any]) -> str:
             """Prediction function that uses the ResponsesAgent with ChatPayload.

dao_ai/tools/mcp.py

@@ -30,24 +30,15 @@ def create_mcp_tools(
     """
     logger.debug(f"create_mcp_tools: {function}")
 
+    # Get MCP URL - handles all convenience objects (connection, genie_room, warehouse, etc.)
+    mcp_url = function.mcp_url
+    logger.debug(f"Using MCP URL: {mcp_url}")
+
     # Check if using UC Connection or direct MCP connection
     if function.connection:
         # Use UC Connection approach with DatabricksOAuthClientProvider
         logger.debug(f"Using UC Connection for MCP: {function.connection.name}")
 
-        # Construct URL if not provided
-        if function.url:
-            mcp_url = function.url
-            logger.debug(f"Using provided MCP URL: {mcp_url}")
-        else:
-            # Construct URL from workspace host and connection name
-            # Pattern: https://{workspace_host}/api/2.0/mcp/external/{connection_name}
-            workspace_client = function.connection.workspace_client
-            workspace_host = workspace_client.config.host
-            connection_name = function.connection.name
-            mcp_url = f"{workspace_host}/api/2.0/mcp/external/{connection_name}"
-            logger.debug(f"Constructed MCP URL from connection: {mcp_url}")
-
         async def _list_tools_with_connection():
             """List available tools using DatabricksOAuthClientProvider."""
             workspace_client = function.connection.workspace_client
@@ -147,7 +138,7 @@ def create_mcp_tools(
                 logger.debug("Using existing authentication token")
 
             return {
-                "url": function.url,
+                "url": mcp_url,  # Use the resolved MCP URL
                 "transport": function.transport,
                 "headers": headers,
             }

{dao_ai-0.0.26.dist-info → dao_ai-0.0.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.26
+Version: 0.0.27
 Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
 Project-URL: Homepage, https://github.com/natefleming/dao-ai
 Project-URL: Documentation, https://natefleming.github.io/dao-ai
@@ -41,7 +41,7 @@ Requires-Dist: langgraph>=0.6.10
 Requires-Dist: langmem>=0.0.29
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: mcp>=1.17.0
-Requires-Dist: mlflow>=3.5.0
+Requires-Dist: mlflow>=3.5.1
 Requires-Dist: nest-asyncio>=1.6.0
 Requires-Dist: openevals>=0.0.19
 Requires-Dist: openpyxl>=3.1.5

{dao_ai-0.0.26.dist-info → dao_ai-0.0.27.dist-info}/RECORD

@@ -3,8 +3,8 @@ dao_ai/agent_as_code.py,sha256=kPSeDz2-1jRaed1TMs4LA3VECoyqe9_Ed2beRLB9gXQ,472
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/chat_models.py,sha256=uhwwOTeLyHWqoTTgHrs4n5iSyTwe4EQcLKnh3jRxPWI,8626
 dao_ai/cli.py,sha256=gq-nsapWxDA1M6Jua3vajBvIwf0Oa6YLcB58lEtMKUo,22503
-dao_ai/config.py,sha256=MdX7zGOXf9Fg03lo1vA7G9yoiIzYmJ6kqTJH4aiQ3VQ,65068
-dao_ai/graph.py,sha256=RrcYfdyj9lYXUZyx4Ax4VIlCUbOWviwr1t2qGDM69LQ,8318
+dao_ai/config.py,sha256=DRFj_1W5sfGH5f2tGQaeC733pTIDTqvbyAt14v8FQYs,70296
+dao_ai/graph.py,sha256=9kjJx0oFZKq5J9-Kpri4-0VCJILHYdYyhqQnj0_noxQ,8913
 dao_ai/guardrails.py,sha256=4TKArDONRy8RwHzOT1plZ1rhy3x9GF_aeGpPCRl6wYA,4016
 dao_ai/messages.py,sha256=xl_3-WcFqZKCFCiov8sZOPljTdM3gX3fCHhxq-xFg2U,7005
 dao_ai/models.py,sha256=8r8GIG3EGxtVyWsRNI56lVaBjiNrPkzh4HdwMZRq8iw,31689
@@ -22,20 +22,20 @@ dao_ai/memory/core.py,sha256=DnEjQO3S7hXr3CDDd7C2eE7fQUmcCS_8q9BXEgjPH3U,4271
 dao_ai/memory/postgres.py,sha256=vvI3osjx1EoU5GBA6SCUstTBKillcmLl12hVgDMjfJY,15346
 dao_ai/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dao_ai/providers/base.py,sha256=-fjKypCOk28h6vioPfMj9YZSw_3Kcbi2nMuAyY7vX9k,1383
-dao_ai/providers/databricks.py,sha256=jfQEEgr9mEIttde1182wPpYm3kHq4sRwXoeMNfkpEmA,66512
+dao_ai/providers/databricks.py,sha256=W_lXSMbPTULMAx-KW7zBJfP7LtkcPGRnEfGcSYuu708,65824
 dao_ai/tools/__init__.py,sha256=G5-5Yi6zpQOH53b5IzLdtsC6g0Ep6leI5GxgxOmgw7Q,1203
 dao_ai/tools/agent.py,sha256=WbQnyziiT12TLMrA7xK0VuOU029tdmUBXbUl-R1VZ0Q,1886
 dao_ai/tools/core.py,sha256=Kei33S8vrmvPOAyrFNekaWmV2jqZ-IPS1QDSvU7RZF0,1984
 dao_ai/tools/genie.py,sha256=8HSOCzSg6PlBzBYXMmNfUnl-LO03p3Ki3fxLPm_dhPg,15051
 dao_ai/tools/human_in_the_loop.py,sha256=yk35MO9eNETnYFH-sqlgR-G24TrEgXpJlnZUustsLkI,3681
-dao_ai/tools/mcp.py,sha256=gTY2xTInC5sy9ckZubwp1QIyZ6_qaUXGeJM_NlmPnGE,8892
+dao_ai/tools/mcp.py,sha256=5aQoRtx2z4xm6zgRslc78rSfEQe-mfhqov2NsiybYfc,8416
 dao_ai/tools/python.py,sha256=XcQiTMshZyLUTVR5peB3vqsoUoAAy8gol9_pcrhddfI,1831
 dao_ai/tools/slack.py,sha256=SCvyVcD9Pv_XXPXePE_fSU1Pd8VLTEkKDLvoGTZWy2Y,4775
 dao_ai/tools/time.py,sha256=Y-23qdnNHzwjvnfkWvYsE7PoWS1hfeKy44tA7sCnNac,8759
 dao_ai/tools/unity_catalog.py,sha256=uX_h52BuBAr4c9UeqSMI7DNz3BPRLeai5tBVW4sJqRI,13113
 dao_ai/tools/vector_search.py,sha256=EDYQs51zIPaAP0ma1D81wJT77GQ-v-cjb2XrFVWfWdg,2621
-dao_ai-0.0.26.dist-info/METADATA,sha256=V5Bvxs1b3FJR-rIjfJHqrNyv9r46rZMHeY9ze8xkMX0,42695
-dao_ai-0.0.26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dao_ai-0.0.26.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.26.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.26.dist-info/RECORD,,
+dao_ai-0.0.27.dist-info/METADATA,sha256=iKbZTl0tFi0S2XUir9uIGh6WyOTvM9yH-4FHCHOzsuE,42695
+dao_ai-0.0.27.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dao_ai-0.0.27.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.27.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.27.dist-info/RECORD,,