dao-ai 0.0.29__tar.gz → 0.0.31__tar.gz

{dao_ai-0.0.29 → dao_ai-0.0.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.29
+Version: 0.0.31
 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

{dao_ai-0.0.29 → dao_ai-0.0.31}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dao-ai"
-version = "0.0.29"
+version = "0.0.31"
 description = "DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML."
 readme = "README.md"
 license = { text = "MIT" }

{dao_ai-0.0.29 → dao_ai-0.0.31}/src/dao_ai/config.py

@@ -2009,21 +2009,45 @@ class AppConfig(BaseModel):
     def create_agent(
         self,
         w: WorkspaceClient | None = None,
+        vsc: "VectorSearchClient | None" = None,
+        pat: str | None = None,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        workspace_host: str | None = None,
     ) -> None:
         from dao_ai.providers.base import ServiceProvider
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: ServiceProvider = DatabricksProvider(w=w)
+        provider: ServiceProvider = DatabricksProvider(
+            w=w,
+            vsc=vsc,
+            pat=pat,
+            client_id=client_id,
+            client_secret=client_secret,
+            workspace_host=workspace_host,
+        )
         provider.create_agent(self)
 
     def deploy_agent(
         self,
         w: WorkspaceClient | None = None,
+        vsc: "VectorSearchClient | None" = None,
+        pat: str | None = None,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        workspace_host: str | None = None,
     ) -> None:
         from dao_ai.providers.base import ServiceProvider
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: ServiceProvider = DatabricksProvider(w=w)
+        provider: ServiceProvider = DatabricksProvider(
+            w=w,
+            vsc=vsc,
+            pat=pat,
+            client_id=client_id,
+            client_secret=client_secret,
+            workspace_host=workspace_host,
+        )
         provider.deploy_agent(self)
 
     def find_agents(

{dao_ai-0.0.29 → dao_ai-0.0.31}/src/dao_ai/tools/vector_search.py

@@ -1,3 +1,4 @@
+import os
 from typing import Annotated, Any, Callable, List, Optional, Sequence
 
 import mlflow
@@ -100,12 +101,36 @@ def create_vector_search_tool(
     # Initialize the vector store
     # Note: text_column is only required for self-managed embeddings
     # For Databricks-managed embeddings, it's automatically determined from the index
+    
+    # Build client_args for VectorSearchClient from environment variables
+    # This is needed because during MLflow model validation, credentials must be
+    # explicitly passed to VectorSearchClient via client_args.
+    # The workspace_client parameter in DatabricksVectorSearch is only used to detect
+    # model serving mode - it doesn't pass credentials to VectorSearchClient.
+    client_args: dict[str, Any] = {}
+    if os.environ.get("DATABRICKS_HOST"):
+        client_args["workspace_url"] = os.environ.get("DATABRICKS_HOST")
+    if os.environ.get("DATABRICKS_TOKEN"):
+        client_args["personal_access_token"] = os.environ.get("DATABRICKS_TOKEN")
+    if os.environ.get("DATABRICKS_CLIENT_ID"):
+        client_args["service_principal_client_id"] = os.environ.get(
+            "DATABRICKS_CLIENT_ID"
+        )
+    if os.environ.get("DATABRICKS_CLIENT_SECRET"):
+        client_args["service_principal_client_secret"] = os.environ.get(
+            "DATABRICKS_CLIENT_SECRET"
+        )
+
+    logger.debug(f"Creating DatabricksVectorSearch with client_args keys: {list(client_args.keys())}")
+
+    # Pass both workspace_client (for model serving detection) and client_args (for credentials)
     vector_store: DatabricksVectorSearch = DatabricksVectorSearch(
         index_name=index_name,
         text_column=None,  # Let DatabricksVectorSearch determine this from the index
         columns=columns,
         include_score=True,
         workspace_client=vector_store_config.workspace_client,
+        client_args=client_args if client_args else None,
     )
 
     # Register the retriever schema with MLflow for model serving integration