dao-ai 0.0.30__py3-none-any.whl → 0.0.31__py3-none-any.whl

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

{dao_ai-0.0.30.dist-info → dao_ai-0.0.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.30
+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.30.dist-info → dao_ai-0.0.31.dist-info}/RECORD

@@ -33,9 +33,9 @@ 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=h6yCgEtOA3h-anJG0hGB3VvcmC3Os7_qnhz8xxRjv1E,11346
-dao_ai-0.0.30.dist-info/METADATA,sha256=T8RdBoC0m6A9CzhFGYPd81X4USQwabDrrWFTa7aEgCs,42778
-dao_ai-0.0.30.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.0.30.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.30.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.30.dist-info/RECORD,,
+dao_ai/tools/vector_search.py,sha256=6SHw3RRR9Itfx7AMQAr3glXb15o8qIZDNri7uDw9JxI,12614
+dao_ai-0.0.31.dist-info/METADATA,sha256=fq-wYpeRMkTteamhkpnMjp64G0-GnTIxNxousrQLKb4,42778
+dao_ai-0.0.31.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.0.31.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.31.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.31.dist-info/RECORD,,