dao-ai 0.1.2__py3-none-any.whl → 0.1.5__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",
@@ -549,13 +607,6 @@ def handle_chat_command(options: Namespace) -> None:
                         # Find the last AI message
                         for msg in reversed(latest_messages):
                             if isinstance(msg, AIMessage):
-                                logger.debug(f"AI message content: {msg.content}")
-                                logger.debug(
-                                    f"AI message has tool_calls: {hasattr(msg, 'tool_calls')}"
-                                )
-                                if hasattr(msg, "tool_calls"):
-                                    logger.debug(f"Tool calls: {msg.tool_calls}")
-
                                 if hasattr(msg, "content") and msg.content:
                                     response_content = msg.content
                                     print(response_content, end="", flush=True)
@@ -676,7 +727,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,39 +764,59 @@ 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 (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
 
     if config_path and not config_path.exists():
         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)
+    # Build databricks command
+    # --profile is a global flag, --target is a subcommand flag for 'bundle'
     cmd = ["databricks"]
     if profile:
         cmd.extend(["--profile", profile])
 
+    cmd.extend(command)
+
+    # --target must come after the bundle subcommand (it's a subcommand-specific flag)
     if target:
         cmd.extend(["--target", target])
 
-    cmd.extend(command)
-
     # Add config_path variable for notebooks
     if config_path and app_config:
         # Calculate relative path from notebooks directory to config file
@@ -800,30 +871,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")
@@ -1009,27 +1027,92 @@ class VolumePathModel(BaseModel, HasFullName):
 
 
 class VectorStoreModel(IsDatabricksResource):
+    """
+    Configuration model for a Databricks Vector Search store.
+
+    Supports two modes:
+    1. **Use Existing Index**: Provide only `index` (fully qualified name).
+       Used for querying an existing vector search index at runtime.
+    2. **Provisioning Mode**: Provide `source_table` + `embedding_source_column`.
+       Used for creating a new vector search index.
+
+    Examples:
+        Minimal configuration (use existing index):
+        ```yaml
+        vector_stores:
+          products_search:
+            index:
+              name: catalog.schema.my_index
+        ```
+
+        Full provisioning configuration:
+        ```yaml
+        vector_stores:
+          products_search:
+            source_table:
+              schema: *my_schema
+              name: products
+            embedding_source_column: description
+            endpoint:
+              name: my_endpoint
+        ```
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    embedding_model: Optional[LLMModel] = None
+
+    # RUNTIME: Only index is truly required for querying existing indexes
     index: Optional[IndexModel] = None
+
+    # PROVISIONING ONLY: Required when creating a new index
+    source_table: Optional[TableModel] = None
+    embedding_source_column: Optional[str] = None
+    embedding_model: Optional[LLMModel] = None
     endpoint: Optional[VectorSearchEndpoint] = None
-    source_table: TableModel
+
+    # OPTIONAL: For both modes
     source_path: Optional[VolumePathModel] = None
     checkpoint_path: Optional[VolumePathModel] = None
     primary_key: Optional[str] = None
     columns: Optional[list[str]] = Field(default_factory=list)
     doc_uri: Optional[str] = None
-    embedding_source_column: str
+
+    @model_validator(mode="after")
+    def validate_configuration_mode(self) -> Self:
+        """
+        Validate that configuration is valid for either:
+        - Use existing mode: index is provided
+        - Provisioning mode: source_table + embedding_source_column provided
+        """
+        has_index = self.index is not None
+        has_source_table = self.source_table is not None
+        has_embedding_col = self.embedding_source_column is not None
+
+        # Must have at least index OR source_table
+        if not has_index and not has_source_table:
+            raise ValueError(
+                "Either 'index' (for existing indexes) or 'source_table' "
+                "(for provisioning) must be provided"
+            )
+
+        # If provisioning mode, need embedding_source_column
+        if has_source_table and not has_embedding_col:
+            raise ValueError(
+                "embedding_source_column is required when source_table is provided (provisioning mode)"
+            )
+
+        return self
 
     @model_validator(mode="after")
     def set_default_embedding_model(self) -> Self:
-        if not self.embedding_model:
+        # Only set default embedding model in provisioning mode
+        if self.source_table is not None and not self.embedding_model:
             self.embedding_model = LLMModel(name="databricks-gte-large-en")
         return self
 
     @model_validator(mode="after")
     def set_default_primary_key(self) -> Self:
-        if self.primary_key is None:
+        # Only auto-discover primary key in provisioning mode
+        if self.primary_key is None and self.source_table is not None:
             from dao_ai.providers.databricks import DatabricksProvider
 
             provider: DatabricksProvider = DatabricksProvider()
@@ -1050,14 +1133,16 @@ class VectorStoreModel(IsDatabricksResource):
 
     @model_validator(mode="after")
     def set_default_index(self) -> Self:
-        if self.index is None:
+        # Only generate index from source_table in provisioning mode
+        if self.index is None and self.source_table is not None:
             name: str = f"{self.source_table.name}_index"
             self.index = IndexModel(schema=self.source_table.schema_model, name=name)
         return self
 
     @model_validator(mode="after")
     def set_default_endpoint(self) -> Self:
-        if self.endpoint is None:
+        # Only find/create endpoint in provisioning mode
+        if self.endpoint is None and self.source_table is not None:
             from dao_ai.providers.databricks import (
                 DatabricksProvider,
                 with_available_indexes,
@@ -1092,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)
 
 
@@ -1266,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")
@@ -1549,11 +1656,13 @@ class RerankParametersModel(BaseModel):
             top_n: 5  # Return top 5 after reranking
         ```
 
-    Available models (from fastest to most accurate):
-    - "ms-marco-TinyBERT-L-2-v2" (fastest, smallest)
-    - "ms-marco-MiniLM-L-6-v2"
-    - "ms-marco-MiniLM-L-12-v2" (default, good balance)
-    - "rank-T5-flan" (most accurate, slower)
+    Available models (see https://github.com/PrithivirajDamodaran/FlashRank):
+    - "ms-marco-TinyBERT-L-2-v2" (~4MB, fastest)
+    - "ms-marco-MiniLM-L-12-v2" (~34MB, best cross-encoder, default)
+    - "rank-T5-flan" (~110MB, best non cross-encoder)
+    - "ms-marco-MultiBERT-L-12" (~150MB, multilingual 100+ languages)
+    - "ce-esci-MiniLM-L12-v2" (e-commerce optimized, Amazon ESCI)
+    - "miniReranker_arabic_v1" (Arabic language)
     """
 
     model_config = ConfigDict(use_enum_values=True, extra="forbid")

dao_ai/middleware/__init__.py

@@ -3,8 +3,15 @@
 
 # Re-export LangChain built-in middleware
 from langchain.agents.middleware import (
+    ClearToolUsesEdit,
+    ContextEditingMiddleware,
     HumanInTheLoopMiddleware,
+    ModelCallLimitMiddleware,
+    ModelRetryMiddleware,
+    PIIMiddleware,
     SummarizationMiddleware,
+    ToolCallLimitMiddleware,
+    ToolRetryMiddleware,
     after_agent,
     after_model,
     before_agent,
@@ -37,6 +44,10 @@ from dao_ai.middleware.base import (
     ModelRequest,
     ModelResponse,
 )
+from dao_ai.middleware.context_editing import (
+    create_clear_tool_uses_edit,
+    create_context_editing_middleware,
+)
 from dao_ai.middleware.core import create_factory_middleware
 from dao_ai.middleware.guardrails import (
     ContentFilterMiddleware,
@@ -62,10 +73,15 @@ from dao_ai.middleware.message_validation import (
     create_thread_id_validation_middleware,
     create_user_id_validation_middleware,
 )
+from dao_ai.middleware.model_call_limit import create_model_call_limit_middleware
+from dao_ai.middleware.model_retry import create_model_retry_middleware
+from dao_ai.middleware.pii import create_pii_middleware
 from dao_ai.middleware.summarization import (
     LoggingSummarizationMiddleware,
     create_summarization_middleware,
 )
+from dao_ai.middleware.tool_call_limit import create_tool_call_limit_middleware
+from dao_ai.middleware.tool_retry import create_tool_retry_middleware
 
 __all__ = [
     # Base class (from LangChain)
@@ -85,6 +101,13 @@ __all__ = [
     "SummarizationMiddleware",
     "LoggingSummarizationMiddleware",
     "HumanInTheLoopMiddleware",
+    "ToolCallLimitMiddleware",
+    "ModelCallLimitMiddleware",
+    "ToolRetryMiddleware",
+    "ModelRetryMiddleware",
+    "ContextEditingMiddleware",
+    "ClearToolUsesEdit",
+    "PIIMiddleware",
     # Core factory function
     "create_factory_middleware",
     # DAO AI middleware implementations
@@ -122,4 +145,14 @@ __all__ = [
     "create_assert_middleware",
     "create_suggest_middleware",
     "create_refine_middleware",
+    # Limit and retry middleware factory functions
+    "create_tool_call_limit_middleware",
+    "create_model_call_limit_middleware",
+    "create_tool_retry_middleware",
+    "create_model_retry_middleware",
+    # Context editing middleware factory functions
+    "create_context_editing_middleware",
+    "create_clear_tool_uses_edit",
+    # PII middleware factory functions
+    "create_pii_middleware",
 ]

dao_ai/middleware/assertions.py

@@ -688,7 +688,7 @@ def create_assert_middleware(
         name: Name for function constraints
 
     Returns:
-        AssertMiddleware configured with the constraint
+        List containing AssertMiddleware configured with the constraint
 
     Example:
         # Using a Constraint class
@@ -737,7 +737,7 @@ def create_suggest_middleware(
         name: Name for function constraints
 
     Returns:
-        SuggestMiddleware configured with the constraint
+        List containing SuggestMiddleware configured with the constraint
 
     Example:
         def is_professional(response: str, ctx: dict) -> ConstraintResult:
@@ -783,7 +783,7 @@ def create_refine_middleware(
         select_best: Track and return best response across iterations
 
     Returns:
-        RefineMiddleware configured with the reward function
+        List containing RefineMiddleware configured with the reward function
 
     Example:
         def evaluate_completeness(response: str, ctx: dict) -> float: