dao-ai 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

dao_ai/cli.py

@@ -47,6 +47,57 @@ def get_default_user_id() -> str:
         return local_user
 
 
+def detect_cloud_provider(profile: Optional[str] = None) -> Optional[str]:
+    """
+    Detect the cloud provider from the Databricks workspace URL.
+    
+    The cloud provider is determined by the workspace URL pattern:
+    - Azure: *.azuredatabricks.net
+    - AWS: *.cloud.databricks.com (without gcp subdomain)
+    - GCP: *.gcp.databricks.com
+    
+    Args:
+        profile: Optional Databricks CLI profile name
+        
+    Returns:
+        Cloud provider string ('azure', 'aws', 'gcp') or None if detection fails
+    """
+    try:
+        from databricks.sdk import WorkspaceClient
+        
+        # Create workspace client with optional profile
+        if profile:
+            w = WorkspaceClient(profile=profile)
+        else:
+            w = WorkspaceClient()
+        
+        # Get the workspace URL from config
+        host = w.config.host
+        if not host:
+            logger.warning("Could not determine workspace URL for cloud detection")
+            return None
+            
+        host_lower = host.lower()
+        
+        if "azuredatabricks.net" in host_lower:
+            logger.debug(f"Detected Azure cloud from workspace URL: {host}")
+            return "azure"
+        elif ".gcp.databricks.com" in host_lower:
+            logger.debug(f"Detected GCP cloud from workspace URL: {host}")
+            return "gcp"
+        elif ".cloud.databricks.com" in host_lower or "databricks.com" in host_lower:
+            # AWS uses *.cloud.databricks.com or regional patterns
+            logger.debug(f"Detected AWS cloud from workspace URL: {host}")
+            return "aws"
+        else:
+            logger.warning(f"Could not determine cloud provider from URL: {host}")
+            return None
+            
+    except Exception as e:
+        logger.warning(f"Could not detect cloud provider: {e}")
+        return None
+
+
 env_path: str = find_dotenv()
 if env_path:
     logger.info(f"Loading environment variables from: {env_path}")
@@ -220,6 +271,13 @@ Examples:
         "-t",
         "--target",
         type=str,
+        help="Bundle target name (default: auto-generated from app name and cloud)",
+    )
+    bundle_parser.add_argument(
+        "--cloud",
+        type=str,
+        choices=["azure", "aws", "gcp"],
+        help="Cloud provider (auto-detected from workspace URL if not specified)",
     )
     bundle_parser.add_argument(
         "--dry-run",
@@ -676,7 +734,7 @@ def generate_bundle_from_template(config_path: Path, app_name: str) -> Path:
     4. Returns the path to the generated file
 
     The generated databricks.yaml is overwritten on each deployment and is not tracked in git.
-    Schema reference remains pointing to ./schemas/bundle_config_schema.json.
+    The template contains cloud-specific targets (azure, aws, gcp) with appropriate node types.
 
     Args:
         config_path: Path to the app config file
@@ -713,15 +771,17 @@ def run_databricks_command(
     profile: Optional[str] = None,
     config: Optional[str] = None,
     target: Optional[str] = None,
+    cloud: Optional[str] = None,
     dry_run: bool = False,
 ) -> None:
-    """Execute a databricks CLI command with optional profile and target.
+    """Execute a databricks CLI command with optional profile, target, and cloud.
     
     Args:
         command: The databricks CLI command to execute (e.g., ["bundle", "deploy"])
         profile: Optional Databricks CLI profile name
         config: Optional path to the configuration file
-        target: Optional bundle target name
+        target: Optional bundle target name (if not provided, auto-generated from app name and cloud)
+        cloud: Optional cloud provider ('azure', 'aws', 'gcp'). Auto-detected if not specified.
         dry_run: If True, print the command without executing
     """
     config_path = Path(config) if config else None
@@ -730,22 +790,30 @@ def run_databricks_command(
         logger.error(f"Configuration file {config_path} does not exist.")
         sys.exit(1)
 
-    # Load app config and generate bundle from template
+    # Load app config
     app_config: AppConfig = AppConfig.from_file(config_path) if config_path else None
     normalized_name: str = normalize_name(app_config.app.name) if app_config else None
 
+    # Auto-detect cloud provider if not specified
+    if not cloud:
+        cloud = detect_cloud_provider(profile)
+        if cloud:
+            logger.info(f"Auto-detected cloud provider: {cloud}")
+        else:
+            logger.warning("Could not detect cloud provider. Defaulting to 'azure'.")
+            cloud = "azure"
+
     # Generate app-specific bundle from template (overwrites databricks.yaml temporarily)
     if config_path and app_config:
         generate_bundle_from_template(config_path, normalized_name)
 
-    # Use app name as target if not explicitly provided
-    # This ensures each app gets its own Terraform state in .databricks/bundle/<app-name>/
-    if not target and normalized_name:
-        target = normalized_name
-        logger.debug(f"Using app-specific target: {target}")
+    # Use cloud as target (azure, aws, gcp) - can be overridden with explicit --target
+    if not target:
+        target = cloud
+        logger.debug(f"Using cloud-based target: {target}")
 
-    # Build databricks command (no -c flag needed, uses databricks.yaml in current dir)
-    # Note: --profile is a global flag, but --target is a subcommand flag for 'bundle'
+    # Build databricks command
+    # --profile is a global flag, --target is a subcommand flag for 'bundle'
     cmd = ["databricks"]
     if profile:
         cmd.extend(["--profile", profile])
@@ -810,30 +878,38 @@ def handle_bundle_command(options: Namespace) -> None:
     profile: Optional[str] = options.profile
     config: Optional[str] = options.config
     target: Optional[str] = options.target
+    cloud: Optional[str] = options.cloud
     dry_run: bool = options.dry_run
 
     if options.deploy:
         logger.info("Deploying DAO AI asset bundle...")
         run_databricks_command(
-            ["bundle", "deploy"], profile, config, target, dry_run=dry_run
+            ["bundle", "deploy"],
+            profile=profile,
+            config=config,
+            target=target,
+            cloud=cloud,
+            dry_run=dry_run,
         )
     if options.run:
         logger.info("Running DAO AI system with current configuration...")
         # Use static job resource key that matches databricks.yaml (resources.jobs.deploy_job)
         run_databricks_command(
             ["bundle", "run", "deploy_job"],
-            profile,
-            config,
-            target,
+            profile=profile,
+            config=config,
+            target=target,
+            cloud=cloud,
             dry_run=dry_run,
         )
     if options.destroy:
         logger.info("Destroying DAO AI system with current configuration...")
         run_databricks_command(
             ["bundle", "destroy", "--auto-approve"],
-            profile,
-            config,
-            target,
+            profile=profile,
+            config=config,
+            target=target,
+            cloud=cloud,
             dry_run=dry_run,
         )
     else:

dao_ai/config.py

@@ -601,6 +601,8 @@ class VectorSearchEndpoint(BaseModel):
 
 
 class IndexModel(IsDatabricksResource, HasFullName):
+    """Model representing a Databricks Vector Search index."""
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
     name: str
@@ -624,6 +626,22 @@ class IndexModel(IsDatabricksResource, HasFullName):
             )
         ]
 
+    def exists(self) -> bool:
+        """Check if this vector search index exists.
+
+        Returns:
+            True if the index exists, False otherwise.
+        """
+        try:
+            self.workspace_client.vector_search_indexes.get_index(self.full_name)
+            return True
+        except NotFound:
+            logger.debug(f"Index not found: {self.full_name}")
+            return False
+        except Exception as e:
+            logger.warning(f"Error checking index existence for {self.full_name}: {e}")
+            return False
+
 
 class FunctionModel(IsDatabricksResource, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -1159,18 +1177,60 @@ class VectorStoreModel(IsDatabricksResource):
         return self.index.as_resources()
 
     def as_index(self, vsc: VectorSearchClient | None = None) -> VectorSearchIndex:
-        from dao_ai.providers.base import ServiceProvider
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: ServiceProvider = DatabricksProvider(vsc=vsc)
+        provider: DatabricksProvider = DatabricksProvider(vsc=vsc)
         index: VectorSearchIndex = provider.get_vector_index(self)
         return index
 
     def create(self, vsc: VectorSearchClient | None = None) -> None:
-        from dao_ai.providers.base import ServiceProvider
+        """
+        Create or validate the vector search index.
+
+        Behavior depends on configuration mode:
+        - **Provisioning Mode** (source_table provided): Creates the index
+        - **Use Existing Mode** (only index provided): Validates the index exists
+
+        Args:
+            vsc: Optional VectorSearchClient instance
+
+        Raises:
+            ValueError: If configuration is invalid or index doesn't exist
+        """
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: ServiceProvider = DatabricksProvider(vsc=vsc)
+        provider: DatabricksProvider = DatabricksProvider(vsc=vsc)
+
+        if self.source_table is not None:
+            self._create_new_index(provider)
+        else:
+            self._validate_existing_index(provider)
+
+    def _validate_existing_index(self, provider: Any) -> None:
+        """Validate that an existing index is accessible."""
+        if self.index is None:
+            raise ValueError("index is required for 'use existing' mode")
+
+        if self.index.exists():
+            logger.info(
+                "Vector search index exists and ready",
+                index_name=self.index.full_name,
+            )
+        else:
+            raise ValueError(
+                f"Index '{self.index.full_name}' does not exist. "
+                "Provide 'source_table' to provision it."
+            )
+
+    def _create_new_index(self, provider: Any) -> None:
+        """Create a new vector search index from source table."""
+        if self.embedding_source_column is None:
+            raise ValueError("embedding_source_column is required for provisioning")
+        if self.endpoint is None:
+            raise ValueError("endpoint is required for provisioning")
+        if self.index is None:
+            raise ValueError("index is required for provisioning")
+
         provider.create_vector_store(self)
 
 
@@ -1333,32 +1393,12 @@ class DatabaseModel(IsDatabricksResource):
 
     @model_validator(mode="after")
     def update_host(self) -> Self:
-        if self.host is not None:
+        # Lakebase uses instance_name directly via databricks_langchain - host not needed
+        if self.is_lakebase:
             return self
 
-        # If instance_name is provided (Lakebase), try to fetch host from existing instance
-        # This may fail for OBO/ambient auth during model logging (before deployment)
-        if self.is_lakebase:
-            try:
-                existing_instance: DatabaseInstance = (
-                    self.workspace_client.database.get_database_instance(
-                        name=self.instance_name
-                    )
-                )
-                self.host = existing_instance.read_write_dns
-            except Exception as e:
-                # For Lakebase with OBO/ambient auth, we can't fetch at config time
-                # The host will need to be provided explicitly or fetched at runtime
-                if self.on_behalf_of_user:
-                    logger.debug(
-                        f"Could not fetch host for database {self.instance_name} "
-                        f"(Lakebase with OBO mode - will be resolved at runtime): {e}"
-                    )
-                else:
-                    raise ValueError(
-                        f"Could not fetch host for database {self.instance_name}. "
-                        f"Please provide the 'host' explicitly or ensure the instance exists: {e}"
-                    )
+        # For standard PostgreSQL, host must be provided by the user
+        # (enforced by validate_connection_type)
         return self
 
     @model_validator(mode="after")

dao_ai/providers/databricks.py

@@ -625,22 +625,17 @@ class DatabricksProvider(ServiceProvider):
                 df.write.mode("overwrite").saveAsTable(table)
 
     def create_vector_store(self, vector_store: VectorStoreModel) -> None:
-        # Validate that this is a provisioning-mode config
-        if vector_store.source_table is None:
-            raise ValueError(
-                "Cannot create vector store: source_table is required for provisioning. "
-                "This VectorStoreModel appears to be configured for 'use existing index' mode. "
-                "To provision a new vector store, provide source_table and embedding_source_column."
-            )
-        if vector_store.embedding_source_column is None:
-            raise ValueError(
-                "Cannot create vector store: embedding_source_column is required for provisioning."
-            )
-        if vector_store.endpoint is None:
-            raise ValueError(
-                "Cannot create vector store: endpoint is required for provisioning."
-            )
+        """
+        Create a vector search index from a source table.
 
+        This method expects a VectorStoreModel in provisioning mode with all
+        required fields validated. Use VectorStoreModel.create() which handles
+        mode detection and validation.
+
+        Args:
+            vector_store: VectorStoreModel configured for provisioning
+        """
+        # Ensure endpoint exists
         if not endpoint_exists(self.vsc, vector_store.endpoint.name):
             self.vsc.create_endpoint_and_wait(
                 name=vector_store.endpoint.name,

dao_ai/vector_search.py

@@ -72,3 +72,40 @@ def index_exists(
             raise e
     # If we reach here, the index doesn't exist
     return False
+
+
+def find_index(
+    vsc: VectorSearchClient, index_full_name: str
+) -> tuple[bool, str | None]:
+    """
+    Find a Vector Search index across all endpoints.
+
+    Searches all available endpoints to find where the index is located.
+
+    Args:
+        vsc: Databricks Vector Search client instance
+        index_full_name: Fully qualified name of the index (catalog.schema.index)
+
+    Returns:
+        Tuple of (exists: bool, endpoint_name: str | None)
+        - (True, endpoint_name) if index is found
+        - (False, None) if index is not found on any endpoint
+    """
+    try:
+        endpoints = vsc.list_endpoints().get("endpoints", [])
+    except Exception as e:
+        if "REQUEST_LIMIT_EXCEEDED" in str(e):
+            print("WARN: couldn't list endpoints due to REQUEST_LIMIT_EXCEEDED error.")
+            return (False, None)
+        raise e
+
+    for endpoint in endpoints:
+        endpoint_name: str = endpoint["name"]
+        try:
+            vsc.get_index(endpoint_name, index_full_name).describe()
+            return (True, endpoint_name)
+        except Exception:
+            # Index not on this endpoint, try next
+            continue
+
+    return (False, None)

{dao_ai-0.1.3.dist-info → dao_ai-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.1.3
+Version: 0.1.4
 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
@@ -293,6 +293,16 @@ This single command:
 3. Deploys it to Databricks
 4. Creates a serving endpoint
 
+**Deploying to a specific workspace:**
+
+```bash
+# Deploy to AWS workspace
+dao-ai bundle --deploy --run -c config/my_agent.yaml --profile aws-field-eng
+
+# Deploy to Azure workspace
+dao-ai bundle --deploy --run -c config/my_agent.yaml --profile azure-retail
+```
+
 **Step 5: Interact with your agent**
 
 Once deployed, you can chat with your agent using Python:
@@ -422,10 +432,34 @@ dao-ai graph -c config/my_config.yaml -o workflow.png
 # Deploy with Databricks Asset Bundles
 dao-ai bundle --deploy --run -c config/my_config.yaml
 
+# Deploy to a specific workspace (multi-cloud support)
+dao-ai bundle --deploy -c config/my_config.yaml --profile aws-field-eng
+dao-ai bundle --deploy -c config/my_config.yaml --profile azure-retail
+
 # Interactive chat with agent
 dao-ai chat -c config/my_config.yaml
 ```
 
+### Multi-Cloud Deployment
+
+DAO AI supports deploying to Azure, AWS, and GCP workspaces with automatic cloud detection:
+
+```bash
+# Deploy to AWS workspace
+dao-ai bundle --deploy -c config/my_config.yaml --profile aws-prod
+
+# Deploy to Azure workspace
+dao-ai bundle --deploy -c config/my_config.yaml --profile azure-prod
+
+# Deploy to GCP workspace
+dao-ai bundle --deploy -c config/my_config.yaml --profile gcp-prod
+```
+
+The CLI automatically:
+- Detects the cloud provider from your profile's workspace URL
+- Selects appropriate compute node types for each cloud
+- Creates isolated deployment state per profile
+
 👉 **Learn more:** [CLI Reference Documentation](docs/cli-reference.md)
 
 ---

{dao_ai-0.1.3.dist-info → dao_ai-0.1.4.dist-info}/RECORD

@@ -1,8 +1,8 @@
 dao_ai/__init__.py,sha256=18P98ExEgUaJ1Byw440Ct1ty59v6nxyWtc5S6Uq2m9Q,1062
 dao_ai/agent_as_code.py,sha256=xIlLDpPVfmDVzLvbdY_V_CrC4Jvj2ItCWJ-NzdrszTo,538
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
-dao_ai/cli.py,sha256=1TS4I6LQq6ixxd0JQA1bt3DHRizcSiynZkSBcZA3-XE,33149
-dao_ai/config.py,sha256=FmXXwFMWqSxl2ofxPIL2ZLRbx7kx-u6VjFx4wdQnFVk,119755
+dao_ai/cli.py,sha256=htmnPXBpKZuXlWRhIl-91zndZO8-nm-REzTBZ49Y-T8,35849
+dao_ai/config.py,sha256=GWBmrbiixMG0ZszLk_XTRKRIS0QqOk_TIQhauK--MIY,120863
 dao_ai/graph.py,sha256=1-uQlo7iXZQTT3uU8aYu0N5rnhw5_g_2YLwVsAs6M-U,1119
 dao_ai/logging.py,sha256=lYy4BmucCHvwW7aI3YQkQXKJtMvtTnPDu9Hnd7_O4oc,1556
 dao_ai/messages.py,sha256=4ZBzO4iFdktGSLrmhHzFjzMIt2tpaL-aQLHOQJysGnY,6959
@@ -13,7 +13,7 @@ dao_ai/prompts.py,sha256=G0ng5f2PkzfgdKrSl03Rnd6riZn5APedof0GAzsWQI8,4792
 dao_ai/state.py,sha256=0wbbzfQmldkCu26gdTE5j0Rl-_pfilza-YIHPbSWlvI,6394
 dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dao_ai/utils.py,sha256=_Urd7Nj2VzrgPKf3NS4E6vt0lWRhEUddBqWN9BksqeE,11543
-dao_ai/vector_search.py,sha256=jlaFS_iizJ55wblgzZmswMM3UOL-qOp2BGJc0JqXYSg,2839
+dao_ai/vector_search.py,sha256=8d3xROg9zSIYNXjRRl6rSexsJTlufjRl5Fy1ZA8daKA,4019
 dao_ai/genie/__init__.py,sha256=vdEyGhrt6L8GlK75SyYvTnl8QpHKDCJC5hJKLg4DesQ,1063
 dao_ai/genie/core.py,sha256=HPKbocvhnnw_PkQwfoq5bpgQmL9lZyyS6_goTJL8yiY,1073
 dao_ai/genie/cache/__init__.py,sha256=JfgCJl1NYQ1aZvZ4kly4T6uQK6ZCJ6PX_htuq7nJF50,1203
@@ -42,7 +42,7 @@ dao_ai/orchestration/supervisor.py,sha256=0aQLL1a3oBjG9LYTuV1y73rB0mqkIkax7MwtZ7
 dao_ai/orchestration/swarm.py,sha256=8tp1eGmsQqqWpaDcjPoJckddPWohZdmmN0RGRJ_xzOA,9198
 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=63jUkvFvdQoJGUlSEJgCsq6hPtj6tviE-AW_TPJvuvc,61850
+dao_ai/providers/databricks.py,sha256=XxYkyoDYkwGV_Xg1IJBpGOl4d7U5HiFP4RtjjSLgenI,61437
 dao_ai/tools/__init__.py,sha256=SRd7W2DOCXKbWWy8lclRtJiCskz7SDAm94qaFF47urQ,1664
 dao_ai/tools/agent.py,sha256=plIWALywRjaDSnot13nYehBsrHRpBUpsVZakoGeajOE,1858
 dao_ai/tools/core.py,sha256=bRIN3BZhRQX8-Kpu3HPomliodyskCqjxynQmYbk6Vjs,3783
@@ -57,8 +57,8 @@ dao_ai/tools/sql.py,sha256=tKd1gjpLuKdQDyfmyYYtMiNRHDW6MGRbdEVaeqyB8Ok,7632
 dao_ai/tools/time.py,sha256=tufJniwivq29y0LIffbgeBTIDE6VgrLpmVf8Qr90qjw,9224
 dao_ai/tools/unity_catalog.py,sha256=AjQfW7bvV8NurqDLIyntYRv2eJuTwNdbvex1L5CRjOk,15534
 dao_ai/tools/vector_search.py,sha256=oe2uBwl2TfeJIXPpwiS6Rmz7wcHczSxNyqS9P3hE6co,14542
-dao_ai-0.1.3.dist-info/METADATA,sha256=0WeeUk6J2uQRdk9Hm1FkCk75RG8psnh3WrSt4A8MfNY,15615
-dao_ai-0.1.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.1.3.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.1.3.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.1.3.dist-info/RECORD,,
+dao_ai-0.1.4.dist-info/METADATA,sha256=d19ZTQ88c21cR9B9W14BZX-zk2YUcMJk1NRioMjeCuw,16685
+dao_ai-0.1.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.1.4.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.1.4.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.1.4.dist-info/RECORD,,