dao-ai 0.0.23__py3-none-any.whl → 0.0.24__py3-none-any.whl

dao_ai/config.py

@@ -1012,8 +1012,12 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
 
     @model_validator(mode="after")
     def validate_mutually_exclusive(self):
-        if self.transport == TransportType.STREAMABLE_HTTP and not self.url:
-            raise ValueError("url must be provided for STREAMABLE_HTTP transport")
+        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:
@@ -1162,6 +1166,40 @@ class MemoryModel(BaseModel):
 FunctionHook: TypeAlias = PythonFunctionModel | FactoryFunctionModel | str
 
 
+class PromptModel(BaseModel, HasFullName):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
+    name: str
+    description: Optional[str] = None
+    default_template: Optional[str] = None
+    alias: Optional[str] = None
+    version: Optional[int] = None
+    tags: Optional[dict[str, Any]] = Field(default_factory=dict)
+
+    @property
+    def template(self) -> str:
+        from dao_ai.providers.databricks import DatabricksProvider
+
+        provider: DatabricksProvider = DatabricksProvider()
+        prompt: str = provider.get_prompt(self)
+        return prompt
+
+    @property
+    def full_name(self) -> str:
+        if self.schema_model:
+            name: str = ""
+            if self.name:
+                name = f".{self.name}"
+            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}{name}"
+        return self.name
+
+    @model_validator(mode="after")
+    def validate_mutually_exclusive(self):
+        if self.alias and self.version:
+            raise ValueError("Cannot specify both alias and version")
+        return self
+
+
 class AgentModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
@@ -1169,7 +1207,7 @@ class AgentModel(BaseModel):
     model: LLMModel
     tools: list[ToolModel] = Field(default_factory=list)
     guardrails: list[GuardrailModel] = Field(default_factory=list)
-    prompt: Optional[str] = None
+    prompt: Optional[str | PromptModel] = None
     handoff_prompt: Optional[str] = None
     create_agent_hook: Optional[FunctionHook] = None
     pre_agent_hook: Optional[FunctionHook] = None
@@ -1495,6 +1533,7 @@ class AppConfig(BaseModel):
     tools: dict[str, ToolModel] = Field(default_factory=dict)
     guardrails: dict[str, GuardrailModel] = Field(default_factory=dict)
     memory: Optional[MemoryModel] = None
+    prompts: dict[str, PromptModel] = Field(default_factory=dict)
     agents: dict[str, AgentModel] = Field(default_factory=dict)
     app: Optional[AppModel] = None
     evaluation: Optional[EvaluationModel] = None

dao_ai/prompts.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Sequence
+from typing import Any, Callable, Optional, Sequence
 
 from langchain.prompts import PromptTemplate
 from langchain_core.messages import (
@@ -8,18 +8,26 @@ from langchain_core.messages import (
 from langchain_core.runnables import RunnableConfig
 from loguru import logger
 
+from dao_ai.config import PromptModel
 from dao_ai.state import SharedState
 
 
-def make_prompt(base_system_prompt: str) -> Callable[[dict, RunnableConfig], list]:
+def make_prompt(
+    base_system_prompt: Optional[str | PromptModel],
+) -> Callable[[dict, RunnableConfig], list]:
     logger.debug(f"make_prompt: {base_system_prompt}")
 
     def prompt(state: SharedState, config: RunnableConfig) -> list:
         system_prompt: str = ""
         if base_system_prompt:
-            prompt_template: PromptTemplate = PromptTemplate.from_template(
-                base_system_prompt
-            )
+            # Extract template string from PromptModel or use string directly
+            template_str: str
+            if isinstance(base_system_prompt, PromptModel):
+                template_str = base_system_prompt.template
+            else:
+                template_str = base_system_prompt
+
+            prompt_template: PromptTemplate = PromptTemplate.from_template(template_str)
 
             params: dict[str, Any] = {
                 input_variable: "" for input_variable in prompt_template.input_variables

dao_ai/providers/databricks.py

@@ -30,6 +30,7 @@ from databricks.vector_search.index import VectorSearchIndex
 from loguru import logger
 from mlflow import MlflowClient
 from mlflow.entities import Experiment
+from mlflow.entities.model_registry import PromptVersion
 from mlflow.entities.model_registry.model_version import ModelVersion
 from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy
 from mlflow.models.model import ModelInfo
@@ -52,6 +53,7 @@ from dao_ai.config import (
     IndexModel,
     IsDatabricksResource,
     LLMModel,
+    PromptModel,
     SchemaModel,
     TableModel,
     UnityCatalogFunctionSqlModel,
@@ -1023,3 +1025,79 @@ class DatabricksProvider(ServiceProvider):
                 f"Error creating instance role '{role_name}' for database {instance_name}: {e}"
             )
             raise
+
+    def get_prompt(self, prompt_model: PromptModel) -> str:
+        """Load prompt from MLflow Prompt Registry or fall back to default_template."""
+        prompt_name: str = prompt_model.full_name
+
+        # Build prompt URI based on alias, version, or default to latest
+        if prompt_model.alias:
+            prompt_uri = f"prompts:/{prompt_name}@{prompt_model.alias}"
+        elif prompt_model.version:
+            prompt_uri = f"prompts:/{prompt_name}/{prompt_model.version}"
+        else:
+            prompt_uri = f"prompts:/{prompt_name}@latest"
+
+        try:
+            from mlflow.genai.prompts import Prompt
+
+            prompt_obj: Prompt = mlflow.genai.load_prompt(prompt_uri)
+            return prompt_obj.to_single_brace_format()
+
+        except Exception as e:
+            logger.warning(f"Failed to load prompt '{prompt_name}' from registry: {e}")
+
+            if prompt_model.default_template:
+                logger.info(f"Using default_template for '{prompt_name}'")
+                self._sync_default_template_to_registry(
+                    prompt_name, prompt_model.default_template, prompt_model.description
+                )
+                return prompt_model.default_template
+
+            raise ValueError(
+                f"Prompt '{prompt_name}' not found in registry and no default_template provided"
+            ) from e
+
+    def _sync_default_template_to_registry(
+        self, prompt_name: str, default_template: str, description: str | None = None
+    ) -> None:
+        """Register default_template to prompt registry under 'default' alias if changed."""
+        try:
+            # Check if default alias already has the same template
+            try:
+                logger.debug(f"Loading prompt '{prompt_name}' from registry...")
+                existing: PromptVersion = mlflow.genai.load_prompt(
+                    f"prompts:/{prompt_name}@default"
+                )
+                if (
+                    existing.to_single_brace_format().strip()
+                    == default_template.strip()
+                ):
+                    logger.debug(f"Prompt '{prompt_name}' is already up-to-date")
+                    return  # Already up-to-date
+            except Exception:
+                logger.debug(
+                    f"Default alias for prompt '{prompt_name}' doesn't exist yet"
+                )
+
+            # Register new version and set as default alias
+            commit_message = description or "Auto-synced from default_template"
+            prompt_version: PromptVersion = mlflow.genai.register_prompt(
+                name=prompt_name,
+                template=default_template,
+                commit_message=commit_message,
+            )
+
+            logger.debug(f"Setting default alias for prompt '{prompt_name}'")
+            mlflow.genai.set_prompt_alias(
+                name=prompt_name,
+                alias="default",
+                version=prompt_version.version,
+            )
+
+            logger.info(
+                f"Synced prompt '{prompt_name}' v{prompt_version.version} to registry"
+            )
+
+        except Exception as e:
+            logger.warning(f"Failed to sync '{prompt_name}' to registry: {e}")

dao_ai/tools/mcp.py

@@ -35,14 +35,26 @@ def create_mcp_tools(
     if function.connection:
         # Use UC Connection approach with DatabricksOAuthClientProvider
         logger.debug(f"Using UC Connection for MCP: {function.connection.name}")
-        logger.debug(f"MCP URL: {function.url}")
+
+        # 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 _get_tools_with_connection():
             """Get tools using DatabricksOAuthClientProvider."""
             workspace_client = function.connection.workspace_client
 
             async with streamablehttp_client(
-                function.url, auth=DatabricksOAuthClientProvider(workspace_client)
+                mcp_url, auth=DatabricksOAuthClientProvider(workspace_client)
             ) as (read_stream, write_stream, _):
                 async with ClientSession(read_stream, write_stream) as session:
                     # Initialize and list tools

{dao_ai-0.0.23.dist-info → dao_ai-0.0.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.23
+Version: 0.0.24
 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
@@ -24,8 +24,8 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
 Requires-Python: >=3.12
-Requires-Dist: databricks-agents>=1.6.1
-Requires-Dist: databricks-langchain>=0.8.0
+Requires-Dist: databricks-agents>=1.7.0
+Requires-Dist: databricks-langchain>=0.8.1
 Requires-Dist: databricks-mcp>=0.3.0
 Requires-Dist: databricks-sdk[openai]>=0.67.0
 Requires-Dist: duckduckgo-search>=8.0.2
@@ -33,19 +33,19 @@ Requires-Dist: grandalf>=0.8
 Requires-Dist: langchain-mcp-adapters>=0.1.10
 Requires-Dist: langchain-tavily>=0.2.11
 Requires-Dist: langchain>=0.3.27
-Requires-Dist: langgraph-checkpoint-postgres>=2.0.23
+Requires-Dist: langgraph-checkpoint-postgres>=2.0.25
 Requires-Dist: langgraph-supervisor>=0.0.29
 Requires-Dist: langgraph-swarm>=0.0.14
-Requires-Dist: langgraph>=0.6.7
+Requires-Dist: langgraph>=0.6.10
 Requires-Dist: langmem>=0.0.29
 Requires-Dist: loguru>=0.7.3
-Requires-Dist: mcp>=1.14.1
+Requires-Dist: mcp>=1.17.0
 Requires-Dist: mlflow>=3.4.0
 Requires-Dist: nest-asyncio>=1.6.0
 Requires-Dist: openevals>=0.0.19
 Requires-Dist: openpyxl>=3.1.5
 Requires-Dist: psycopg[binary,pool]>=3.2.9
-Requires-Dist: pydantic>=2.11.3
+Requires-Dist: pydantic>=2.12.0
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: rich>=14.0.0

{dao_ai-0.0.23.dist-info → dao_ai-0.0.24.dist-info}/RECORD

@@ -3,13 +3,13 @@ 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=Aez2TQW3Q8Ho1IaIkRggt0NevDxAAVPjXkePC5GPJF0,20429
-dao_ai/config.py,sha256=XHU6xkRAoTeiZYH5ns_fLwcR6EaxRAGeMwRSoW3n0S8,55431
+dao_ai/config.py,sha256=j9SAdf7UHSoS2pLos-oypJNDPS48A2rRq55OEW1wsMI,56755
 dao_ai/graph.py,sha256=APYc2y3cig4P52X4sOHSFSZNK8j5EtEPJLFwWeJ3KQQ,7956
 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
 dao_ai/nodes.py,sha256=SSuFNTXOdFaKg_aX-yUkQO7fM9wvNGu14lPXKDapU1U,8461
-dao_ai/prompts.py,sha256=vpmIbWs_szXUgNNDs5Gh2LcxKZti5pHDKSfoClUcgX0,1289
+dao_ai/prompts.py,sha256=7Hcstmv514P0s9s-TVoIlbkDV2XXOphGCW6gcPeyUYE,1628
 dao_ai/state.py,sha256=_lF9krAYYjvFDMUwZzVKOn0ZnXKcOrbjWKdre0C5B54,1137
 dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dao_ai/utils.py,sha256=yXgqHrYdO5qDxgxUs2G5XJeLFgwg8D0BIJvbFkqSbhs,4519
@@ -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=CFZ2RojcTjiJ1aGwNI3_0qCGf339w2o5h9CRDKNesLs,39825
+dao_ai/providers/databricks.py,sha256=1BPYQxi4-Z4I1ygZYlKV8ycdxZTtWNXplySToayHCEI,43096
 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=RAAG97boEDJKlX7X_XUz-l-nH5DdqtHUG_I2zw1lWNk,6844
+dao_ai/tools/mcp.py,sha256=CYv59yn-LIY11atUgNtN2W6vR7C6Qyo7-rvPcVJnXVk,7461
 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.23.dist-info/METADATA,sha256=6GfCnhhQN9t4x1LX8mUHsOTfr4mgdGR1Xx070pjIm_g,42638
-dao_ai-0.0.23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dao_ai-0.0.23.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.23.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.23.dist-info/RECORD,,
+dao_ai-0.0.24.dist-info/METADATA,sha256=69qYBqdZg8tZF1ni90RBrSG9OAmE4jXdf5lec6U0TL8,42639
+dao_ai-0.0.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dao_ai-0.0.24.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.24.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.24.dist-info/RECORD,,