dao-ai 0.0.32__py3-none-any.whl → 0.0.34__py3-none-any.whl

dao_ai/config.py

@@ -12,6 +12,7 @@ from typing import (
     Iterator,
     Literal,
     Optional,
+    Self,
     Sequence,
     TypeAlias,
     Union,
@@ -200,6 +201,15 @@ AnyVariable: TypeAlias = (
 )
 
 
+class ServicePrincipalModel(BaseModel):
+    model_config = ConfigDict(
+        frozen=True,
+        use_enum_values=True,
+    )
+    client_id: AnyVariable
+    client_secret: AnyVariable
+
+
 class Privilege(str, Enum):
     ALL_PRIVILEGES = "ALL_PRIVILEGES"
     USE_CATALOG = "USE_CATALOG"
@@ -226,9 +236,21 @@ class Privilege(str, Enum):
 
 class PermissionModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    principals: list[str] = Field(default_factory=list)
+    principals: list[ServicePrincipalModel | str] = Field(default_factory=list)
     privileges: list[Privilege]
 
+    @model_validator(mode="after")
+    def resolve_principals(self) -> Self:
+        """Resolve ServicePrincipalModel objects to their client_id."""
+        resolved: list[str] = []
+        for principal in self.principals:
+            if isinstance(principal, ServicePrincipalModel):
+                resolved.append(value_of(principal.client_id))
+            else:
+                resolved.append(principal)
+        self.principals = resolved
+        return self
+
 
 class SchemaModel(BaseModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -451,7 +473,7 @@ class GenieRoomModel(BaseModel, IsDatabricksResource):
         ]
 
     @model_validator(mode="after")
-    def update_space_id(self):
+    def update_space_id(self) -> Self:
         self.space_id = value_of(self.space_id)
         return self
 
@@ -530,13 +552,13 @@ class VectorStoreModel(BaseModel, IsDatabricksResource):
     embedding_source_column: str
 
     @model_validator(mode="after")
-    def set_default_embedding_model(self):
+    def set_default_embedding_model(self) -> Self:
         if 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):
+    def set_default_primary_key(self) -> Self:
         if self.primary_key is None:
             from dao_ai.providers.databricks import DatabricksProvider
 
@@ -557,14 +579,14 @@ class VectorStoreModel(BaseModel, IsDatabricksResource):
         return self
 
     @model_validator(mode="after")
-    def set_default_index(self):
+    def set_default_index(self) -> Self:
         if self.index is 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):
+    def set_default_endpoint(self) -> Self:
         if self.endpoint is None:
             from dao_ai.providers.databricks import (
                 DatabricksProvider,
@@ -719,7 +741,7 @@ class WarehouseModel(BaseModel, IsDatabricksResource):
         ]
 
     @model_validator(mode="after")
-    def update_warehouse_id(self):
+    def update_warehouse_id(self) -> Self:
         self.warehouse_id = value_of(self.warehouse_id)
         return self
 
@@ -737,12 +759,28 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
        - Used for: discovering database instance, getting host DNS, checking instance status
        - Controlled by: DATABRICKS_HOST, DATABRICKS_TOKEN env vars, or SDK default config
 
-    2. **Database Connection Authentication** (configured via client_id/client_secret OR user):
+    2. **Database Connection Authentication** (configured via service_principal, client_id/client_secret, OR user):
        - Used for: connecting to the PostgreSQL database as a specific identity
+       - Service Principal: Set service_principal with workspace_host to connect as a service principal
        - OAuth M2M: Set client_id, client_secret, workspace_host to connect as a service principal
        - User Auth: Set user (and optionally password) to connect as a user identity
 
-    Example OAuth M2M Configuration:
+    Example Service Principal Configuration:
+    ```yaml
+    databases:
+      my_lakebase:
+        name: my-database
+        service_principal:
+          client_id:
+            env: SERVICE_PRINCIPAL_CLIENT_ID
+          client_secret:
+            scope: my-scope
+            secret: sp-client-secret
+        workspace_host:
+          env: DATABRICKS_HOST
+    ```
+
+    Example OAuth M2M Configuration (alternative):
     ```yaml
     databases:
       my_lakebase:
@@ -779,6 +817,7 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
     node_count: Optional[int] = None
     user: Optional[AnyVariable] = None
     password: Optional[AnyVariable] = None
+    service_principal: Optional[ServicePrincipalModel] = None
     client_id: Optional[AnyVariable] = None
     client_secret: Optional[AnyVariable] = None
     workspace_host: Optional[AnyVariable] = None
@@ -796,14 +835,24 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         ]
 
     @model_validator(mode="after")
-    def update_instance_name(self):
+    def update_instance_name(self) -> Self:
         if self.instance_name is None:
             self.instance_name = self.name
 
         return self
 
     @model_validator(mode="after")
-    def update_user(self):
+    def expand_service_principal(self) -> Self:
+        """Expand service_principal into client_id and client_secret if provided."""
+        if self.service_principal is not None:
+            if self.client_id is None:
+                self.client_id = self.service_principal.client_id
+            if self.client_secret is None:
+                self.client_secret = self.service_principal.client_secret
+        return self
+
+    @model_validator(mode="after")
+    def update_user(self) -> Self:
         if self.client_id or self.user:
             return self
 
@@ -816,7 +865,7 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         return self
 
     @model_validator(mode="after")
-    def update_host(self):
+    def update_host(self) -> Self:
         if self.host is not None:
             return self
 
@@ -829,7 +878,7 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         return self
 
     @model_validator(mode="after")
-    def validate_auth_methods(self):
+    def validate_auth_methods(self) -> Self:
         oauth_fields: Sequence[Any] = [
             self.workspace_host,
             self.client_id,
@@ -849,8 +898,8 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         if not has_oauth and not has_user_auth:
             raise ValueError(
                 "At least one authentication method must be provided: "
-                "either OAuth credentials (workspace_host, client_id, client_secret) "
-                "or user credentials (user, password)."
+                "either OAuth credentials (workspace_host, client_id, client_secret), "
+                "service_principal with workspace_host, or user credentials (user, password)."
             )
 
         return self
@@ -997,14 +1046,14 @@ class RetrieverModel(BaseModel):
     )
 
     @model_validator(mode="after")
-    def set_default_columns(self):
+    def set_default_columns(self) -> Self:
         if not self.columns:
             columns: Sequence[str] = self.vector_store.columns
             self.columns = columns
         return self
 
     @model_validator(mode="after")
-    def set_default_reranker(self):
+    def set_default_reranker(self) -> Self:
         """Convert bool to ReRankParametersModel with defaults."""
         if isinstance(self.rerank, bool) and self.rerank:
             self.rerank = RerankParametersModel()
@@ -1091,7 +1140,7 @@ class FactoryFunctionModel(BaseFunctionModel, HasFullName):
         return [create_factory_tool(self, **kwargs)]
 
     @model_validator(mode="after")
-    def update_args(self):
+    def update_args(self) -> Self:
         for key, value in self.args.items():
             self.args[key] = value_of(value)
         return self
@@ -1111,6 +1160,7 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
     headers: dict[str, AnyVariable] = Field(default_factory=dict)
     args: list[str] = Field(default_factory=list)
     pat: Optional[AnyVariable] = None
+    service_principal: Optional[ServicePrincipalModel] = None
     client_id: Optional[AnyVariable] = None
     client_secret: Optional[AnyVariable] = None
     workspace_host: Optional[AnyVariable] = None
@@ -1120,6 +1170,16 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
     sql: Optional[bool] = None
     vector_search: Optional[VectorStoreModel] = None
 
+    @model_validator(mode="after")
+    def expand_service_principal(self) -> Self:
+        """Expand service_principal into client_id and client_secret if provided."""
+        if self.service_principal is not None:
+            if self.client_id is None:
+                self.client_id = self.service_principal.client_id
+            if self.client_secret is None:
+                self.client_secret = self.service_principal.client_secret
+        return self
+
     @property
     def full_name(self) -> str:
         return self.name
@@ -1129,12 +1189,12 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
         Get the workspace host, either from config or from workspace client.
 
         If connection is provided, uses its workspace client.
-        Otherwise, falls back to creating a new workspace client.
+        Otherwise, falls back to the default Databricks host.
 
         Returns:
             str: The workspace host URL without trailing slash
         """
-        from databricks.sdk import WorkspaceClient
+        from dao_ai.utils import get_default_databricks_host
 
         # Try to get workspace_host from config
         workspace_host: str | None = (
@@ -1147,9 +1207,13 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
             if self.connection:
                 workspace_host = self.connection.workspace_client.config.host
             else:
-                # Create a default workspace client
-                w: WorkspaceClient = WorkspaceClient()
-                workspace_host = w.config.host
+                workspace_host = get_default_databricks_host()
+
+        if not workspace_host:
+            raise ValueError(
+                "Could not determine workspace host. "
+                "Please set workspace_host in config or DATABRICKS_HOST environment variable."
+            )
 
         # Remove trailing slash
         return workspace_host.rstrip("/")
@@ -1356,7 +1420,7 @@ class CheckpointerModel(BaseModel):
     database: Optional[DatabaseModel] = None
 
     @model_validator(mode="after")
-    def validate_postgres_requires_database(self):
+    def validate_postgres_requires_database(self) -> Self:
         if self.type == StorageType.POSTGRES and not self.database:
             raise ValueError("Database must be provided when storage type is POSTGRES")
         return self
@@ -1381,7 +1445,7 @@ class StoreModel(BaseModel):
     namespace: Optional[str] = None
 
     @model_validator(mode="after")
-    def validate_postgres_requires_database(self):
+    def validate_postgres_requires_database(self) -> Self:
         if self.type == StorageType.POSTGRES and not self.database:
             raise ValueError("Database must be provided when storage type is POSTGRES")
         return self
@@ -1445,7 +1509,7 @@ class PromptModel(BaseModel, HasFullName):
         return prompt_version
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self):
+    def validate_mutually_exclusive(self) -> Self:
         if self.alias and self.version:
             raise ValueError("Cannot specify both alias and version")
         return self
@@ -1499,7 +1563,7 @@ class OrchestrationModel(BaseModel):
     memory: Optional[MemoryModel] = None
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self):
+    def validate_mutually_exclusive(self) -> Self:
         if self.supervisor is not None and self.swarm is not None:
             raise ValueError("Cannot specify both supervisor and swarm")
         if self.supervisor is None and self.swarm is None:
@@ -1529,9 +1593,21 @@ class Entitlement(str, Enum):
 
 class AppPermissionModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    principals: list[str] = Field(default_factory=list)
+    principals: list[ServicePrincipalModel | str] = Field(default_factory=list)
     entitlements: list[Entitlement]
 
+    @model_validator(mode="after")
+    def resolve_principals(self) -> Self:
+        """Resolve ServicePrincipalModel objects to their client_id."""
+        resolved: list[str] = []
+        for principal in self.principals:
+            if isinstance(principal, ServicePrincipalModel):
+                resolved.append(value_of(principal.client_id))
+            else:
+                resolved.append(principal)
+        self.principals = resolved
+        return self
+
 
 class LogLevel(str, Enum):
     TRACE = "TRACE"
@@ -1626,6 +1702,7 @@ class ChatHistoryModel(BaseModel):
 class AppModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
+    service_principal: Optional[ServicePrincipalModel] = None
     description: Optional[str] = None
     log_level: Optional[LogLevel] = "WARNING"
     registered_model: RegisteredModelModel
@@ -1661,14 +1738,42 @@ class AppModel(BaseModel):
     )
 
     @model_validator(mode="after")
-    def validate_agents_not_empty(self):
+    def set_databricks_env_vars(self) -> Self:
+        """Set Databricks environment variables for Model Serving.
+
+        Sets DATABRICKS_HOST, DATABRICKS_CLIENT_ID, and DATABRICKS_CLIENT_SECRET.
+        Values explicitly provided in environment_vars take precedence.
+        """
+        from dao_ai.utils import get_default_databricks_host
+
+        # Set DATABRICKS_HOST if not already provided
+        if "DATABRICKS_HOST" not in self.environment_vars:
+            host: str | None = get_default_databricks_host()
+            if host:
+                self.environment_vars["DATABRICKS_HOST"] = host
+
+        # Set service principal credentials if provided
+        if self.service_principal is not None:
+            if "DATABRICKS_CLIENT_ID" not in self.environment_vars:
+                self.environment_vars["DATABRICKS_CLIENT_ID"] = (
+                    self.service_principal.client_id
+                )
+            if "DATABRICKS_CLIENT_SECRET" not in self.environment_vars:
+                self.environment_vars["DATABRICKS_CLIENT_SECRET"] = (
+                    self.service_principal.client_secret
+                )
+        return self
+
+    @model_validator(mode="after")
+    def validate_agents_not_empty(self) -> Self:
         if not self.agents:
             raise ValueError("At least one agent must be specified")
         return self
 
     @model_validator(mode="after")
-    def update_environment_vars(self):
+    def resolve_environment_vars(self) -> Self:
         for key, value in self.environment_vars.items():
+            updated_value: str
             if isinstance(value, SecretVariableModel):
                 updated_value = str(value)
             else:
@@ -1678,7 +1783,7 @@ class AppModel(BaseModel):
         return self
 
     @model_validator(mode="after")
-    def set_default_orchestration(self):
+    def set_default_orchestration(self) -> Self:
         if self.orchestration is None:
             if len(self.agents) > 1:
                 default_agent: AgentModel = self.agents[0]
@@ -1698,14 +1803,14 @@ class AppModel(BaseModel):
         return self
 
     @model_validator(mode="after")
-    def set_default_endpoint_name(self):
+    def set_default_endpoint_name(self) -> Self:
         if self.endpoint_name is None:
             self.endpoint_name = self.name
         return self
 
     @model_validator(mode="after")
-    def set_default_agent(self):
-        default_agent_name = self.agents[0].name
+    def set_default_agent(self) -> Self:
+        default_agent_name: str = self.agents[0].name
 
         if self.orchestration.swarm and not self.orchestration.swarm.default_agent:
             self.orchestration.swarm.default_agent = default_agent_name
@@ -1713,7 +1818,7 @@ class AppModel(BaseModel):
         return self
 
     @model_validator(mode="after")
-    def add_code_paths_to_sys_path(self):
+    def add_code_paths_to_sys_path(self) -> Self:
         for code_path in self.code_paths:
             parent_path: str = str(Path(code_path).parent)
             if parent_path not in sys.path:
@@ -1746,7 +1851,7 @@ class EvaluationDatasetExpectationsModel(BaseModel):
     expected_facts: Optional[list[str]] = None
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self):
+    def validate_mutually_exclusive(self) -> Self:
         if self.expected_response is not None and self.expected_facts is not None:
             raise ValueError("Cannot specify both expected_response and expected_facts")
         return self
@@ -1854,7 +1959,7 @@ class PromptOptimizationModel(BaseModel):
         return optimized_prompt
 
     @model_validator(mode="after")
-    def set_defaults(self):
+    def set_defaults(self) -> Self:
         # If no prompt is specified, try to use the agent's prompt
         if self.prompt is None:
             if isinstance(self.agent.prompt, PromptModel):
@@ -1976,6 +2081,7 @@ class ResourcesModel(BaseModel):
 class AppConfig(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     variables: dict[str, AnyVariable] = Field(default_factory=dict)
+    service_principals: dict[str, ServicePrincipalModel] = Field(default_factory=dict)
     schemas: dict[str, SchemaModel] = Field(default_factory=dict)
     resources: Optional[ResourcesModel] = None
     retrievers: dict[str, RetrieverModel] = Field(default_factory=dict)

dao_ai/models.py

@@ -331,13 +331,23 @@ class LanggraphResponsesAgent(ResponsesAgent):
         context: Context = self._convert_request_to_context(request)
         custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
+        # Build the graph input state, including genie_conversation_ids if provided
+        graph_input: dict[str, Any] = {"messages": messages}
+        if request.custom_inputs and "genie_conversation_ids" in request.custom_inputs:
+            graph_input["genie_conversation_ids"] = request.custom_inputs[
+                "genie_conversation_ids"
+            ]
+            logger.debug(
+                f"Including genie_conversation_ids in graph input: {graph_input['genie_conversation_ids']}"
+            )
+
         # Use async ainvoke internally for parallel execution
         import asyncio
 
         async def _async_invoke():
             try:
                 return await self.graph.ainvoke(
-                    {"messages": messages}, context=context, config=custom_inputs
+                    graph_input, context=context, config=custom_inputs
                 )
             except Exception as e:
                 logger.error(f"Error in graph.ainvoke: {e}")
@@ -399,6 +409,16 @@ class LanggraphResponsesAgent(ResponsesAgent):
         context: Context = self._convert_request_to_context(request)
         custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
+        # Build the graph input state, including genie_conversation_ids if provided
+        graph_input: dict[str, Any] = {"messages": messages}
+        if request.custom_inputs and "genie_conversation_ids" in request.custom_inputs:
+            graph_input["genie_conversation_ids"] = request.custom_inputs[
+                "genie_conversation_ids"
+            ]
+            logger.debug(
+                f"Including genie_conversation_ids in graph input: {graph_input['genie_conversation_ids']}"
+            )
+
         # Use async astream internally for parallel execution
         import asyncio
 
@@ -408,7 +428,7 @@ class LanggraphResponsesAgent(ResponsesAgent):
 
             try:
                 async for nodes, stream_mode, messages_batch in self.graph.astream(
-                    {"messages": messages},
+                    graph_input,
                     context=context,
                     config=custom_inputs,
                     stream_mode=["messages", "custom"],

dao_ai/providers/databricks.py

@@ -1151,7 +1151,7 @@ class DatabricksProvider(ServiceProvider):
         If an explicit version or alias is specified in the prompt_model, uses that directly.
         Otherwise, tries to load prompts in this order:
         1. champion alias
-        2. latest version (max version number from search_prompt_versions)
+        2. latest alias
         3. default alias
         4. Register default_template if provided
 
@@ -1166,7 +1166,6 @@ class DatabricksProvider(ServiceProvider):
         """
 
         prompt_name: str = prompt_model.full_name
-        mlflow_client: MlflowClient = MlflowClient()
 
         # If explicit version or alias is specified, use it directly
         if prompt_model.version or prompt_model.alias:
@@ -1197,19 +1196,13 @@ class DatabricksProvider(ServiceProvider):
         except Exception as e:
             logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
 
-        # 2. Try to get latest version by finding the max version number
+        # 2. Try latest alias
         try:
-            versions = mlflow_client.search_prompt_versions(
-                prompt_name, max_results=100
-            )
-            if versions:
-                latest = max(versions, key=lambda v: int(v.version))
-                logger.info(
-                    f"Loaded prompt '{prompt_name}' version {latest.version} (latest by max version)"
-                )
-                return latest
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@latest")
+            logger.info(f"Loaded prompt '{prompt_name}' from latest alias")
+            return prompt_version
         except Exception as e:
-            logger.debug(f"Failed to find latest version for '{prompt_name}': {e}")
+            logger.debug(f"Latest alias not found for '{prompt_name}': {e}")
 
         # 3. Try default alias
         try:
@@ -1225,7 +1218,7 @@ class DatabricksProvider(ServiceProvider):
                 f"No existing prompt found for '{prompt_name}', "
                 "attempting to register default_template"
             )
-            return self._sync_default_template_to_registry(
+            return self._register_default_template(
                 prompt_name, prompt_model.default_template, prompt_model.description
             )
 
@@ -1235,49 +1228,17 @@ class DatabricksProvider(ServiceProvider):
             "and no default_template provided"
         )
 
-    def _sync_default_template_to_registry(
+    def _register_default_template(
         self, prompt_name: str, default_template: str, description: str | None = None
     ) -> PromptVersion:
-        """Get the best available prompt version, or register default_template if possible.
-
-        Tries to load prompts in order: champion → latest (max version) → default.
-        If none found and we have write permissions, registers the default_template.
-        If registration fails (e.g., in Model Serving), logs the error and raises.
-        """
-        mlflow_client: MlflowClient = MlflowClient()
-
-        # Try to find an existing prompt version in priority order
-        # 1. Try champion alias
-        try:
-            champion = mlflow.genai.load_prompt(f"prompts:/{prompt_name}@champion")
-            logger.info(f"Loaded prompt '{prompt_name}' from champion alias")
-            return champion
-        except Exception as e:
-            logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
+        """Register default_template as a new prompt version.
 
-        # 2. Try to get the latest version by finding the max version number
-        try:
-            versions = mlflow_client.search_prompt_versions(
-                prompt_name, max_results=100
-            )
-            if versions:
-                latest = max(versions, key=lambda v: int(v.version))
-                logger.info(
-                    f"Loaded prompt '{prompt_name}' version {latest.version} (latest by max version)"
-                )
-                return latest
-        except Exception as e:
-            logger.debug(f"Failed to search versions for '{prompt_name}': {e}")
-
-        # 3. Try default alias
-        try:
-            default = mlflow.genai.load_prompt(f"prompts:/{prompt_name}@default")
-            logger.info(f"Loaded prompt '{prompt_name}' from default alias")
-            return default
-        except Exception as e:
-            logger.debug(f"Default alias not found for '{prompt_name}': {e}")
+        Called when no existing prompt version is found (champion, latest, default all failed).
+        Registers the template and sets both 'default' and 'champion' aliases.
 
-        # No existing prompt found - try to register if we have a template
+        If registration fails (e.g., in Model Serving with restricted permissions),
+        logs the error and raises.
+        """
         logger.info(
             f"No existing prompt found for '{prompt_name}', attempting to register default_template"
         )

dao_ai/tools/core.py

@@ -35,7 +35,7 @@ def create_tools(tool_models: Sequence[ToolModel]) -> Sequence[RunnableLike]:
         if name in tools:
             logger.warning(f"Tools already registered for: {name}, skipping creation.")
             continue
-        registered_tools: Sequence[RunnableLike] = tool_registry.get(name)
+        registered_tools: Sequence[RunnableLike] | None = tool_registry.get(name)
         if registered_tools is None:
             logger.debug(f"Creating tools for: {name}...")
             function: AnyTool = tool_config.function

dao_ai/tools/genie.py

@@ -5,8 +5,9 @@ from typing import Annotated, Any, Callable
 
 import pandas as pd
 from databricks_ai_bridge.genie import Genie, GenieResponse
+from langchain.tools import tool
 from langchain_core.messages import ToolMessage
-from langchain_core.tools import InjectedToolCallId, tool
+from langchain_core.tools import InjectedToolCallId
 from langgraph.prebuilt import InjectedState
 from langgraph.types import Command
 from loguru import logger
@@ -43,7 +44,7 @@ def create_genie_tool(
     genie_room: GenieRoomModel | dict[str, Any],
     name: str | None = None,
     description: str | None = None,
-    persist_conversation: bool = False,
+    persist_conversation: bool = True,
     truncate_results: bool = False,
 ) -> Callable[..., Command]:
     """
@@ -64,6 +65,16 @@ def create_genie_tool(
     Returns:
         A LangGraph tool that processes natural language queries through Genie
     """
+    logger.debug("create_genie_tool")
+    logger.debug(f"genie_room type: {type(genie_room)}")
+    logger.debug(f"genie_room: {genie_room}")
+    logger.debug(f"persist_conversation: {persist_conversation}")
+    logger.debug(f"truncate_results: {truncate_results}")
+    logger.debug(f"name: {name}")
+    logger.debug(f"description: {description}")
+    logger.debug(f"genie_room: {genie_room}")
+    logger.debug(f"persist_conversation: {persist_conversation}")
+    logger.debug(f"truncate_results: {truncate_results}")
 
     if isinstance(genie_room, dict):
         genie_room = GenieRoomModel(**genie_room)
@@ -106,14 +117,13 @@ GenieResponse: A response object containing the conversation ID and result from
         state: Annotated[dict, InjectedState],
         tool_call_id: Annotated[str, InjectedToolCallId],
     ) -> Command:
-        """Process a natural language question through Databricks Genie."""
-        # Create Genie instance using databricks_langchain implementation
         genie: Genie = Genie(
             space_id=space_id,
             client=genie_room.workspace_client,
             truncate_results=truncate_results,
         )
 
+        """Process a natural language question through Databricks Genie."""
         # Get existing conversation mapping and retrieve conversation ID for this space
         conversation_ids: dict[str, str] = state.get("genie_conversation_ids", {})
         existing_conversation_id: str | None = conversation_ids.get(space_id)
@@ -131,6 +141,7 @@ GenieResponse: A response object containing the conversation ID and result from
         )
 
         # Update the conversation mapping with the new conversation ID for this space
+
         update: dict[str, Any] = {
             "messages": [
                 ToolMessage(_response_to_json(response), tool_call_id=tool_call_id)

dao_ai/tools/unity_catalog.py

@@ -265,23 +265,52 @@ def with_partial_args(
 
     Args:
         tool: ToolModel containing the Unity Catalog function configuration
-        partial_args: Dictionary of arguments to pre-fill in the tool
+        partial_args: Dictionary of arguments to pre-fill in the tool.
+            Supports:
+            - client_id, client_secret: OAuth credentials directly
+            - service_principal: ServicePrincipalModel with client_id and client_secret
+            - host or workspace_host: Databricks workspace host
 
     Returns:
         StructuredTool: A LangChain tool with partial arguments pre-filled
     """
     from unitycatalog.ai.langchain.toolkit import generate_function_input_params_schema
 
+    from dao_ai.config import ServicePrincipalModel
+
     logger.debug(f"with_partial_args: {tool}")
 
     # Convert dict-based variables to CompositeVariableModel and resolve their values
-    resolved_args = {}
+    resolved_args: dict[str, Any] = {}
     for k, v in partial_args.items():
         if isinstance(v, dict):
             resolved_args[k] = value_of(CompositeVariableModel(**v))
         else:
             resolved_args[k] = value_of(v)
 
+    # Handle service_principal - expand into client_id and client_secret
+    if "service_principal" in resolved_args:
+        sp = resolved_args.pop("service_principal")
+        if isinstance(sp, dict):
+            sp = ServicePrincipalModel(**sp)
+        if isinstance(sp, ServicePrincipalModel):
+            if "client_id" not in resolved_args:
+                resolved_args["client_id"] = value_of(sp.client_id)
+            if "client_secret" not in resolved_args:
+                resolved_args["client_secret"] = value_of(sp.client_secret)
+
+    # Normalize host/workspace_host - accept either key
+    if "workspace_host" in resolved_args and "host" not in resolved_args:
+        resolved_args["host"] = resolved_args.pop("workspace_host")
+
+    # Default host from WorkspaceClient if not provided
+    if "host" not in resolved_args:
+        from dao_ai.utils import get_default_databricks_host
+
+        host: str | None = get_default_databricks_host()
+        if host:
+            resolved_args["host"] = host
+
     logger.debug(f"Resolved partial args: {resolved_args.keys()}")
 
     if isinstance(tool, dict):

dao_ai/utils.py

@@ -38,6 +38,32 @@ def normalize_name(name: str) -> str:
     return normalized.strip("_")
 
 
+def get_default_databricks_host() -> str | None:
+    """Get the default Databricks workspace host.
+
+    Attempts to get the host from:
+    1. DATABRICKS_HOST environment variable
+    2. WorkspaceClient ambient authentication (e.g., from ~/.databrickscfg)
+
+    Returns:
+        The Databricks workspace host URL, or None if not available.
+    """
+    # Try environment variable first
+    host: str | None = os.environ.get("DATABRICKS_HOST")
+    if host:
+        return host
+
+    # Fall back to WorkspaceClient
+    try:
+        from databricks.sdk import WorkspaceClient
+
+        w: WorkspaceClient = WorkspaceClient()
+        return w.config.host
+    except Exception:
+        logger.debug("Could not get default Databricks host from WorkspaceClient")
+        return None
+
+
 def dao_ai_version() -> str:
     """
     Get the dao-ai package version, with fallback for source installations.

{dao_ai-0.0.32.dist-info → dao_ai-0.0.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.32
+Version: 0.0.34
 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
@@ -25,7 +25,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
 Requires-Python: >=3.11
-Requires-Dist: databricks-agents>=1.7.0
+Requires-Dist: databricks-agents>=1.8.2
 Requires-Dist: databricks-langchain>=0.11.0
 Requires-Dist: databricks-mcp>=0.3.0
 Requires-Dist: databricks-sdk[openai]>=0.67.0
@@ -33,21 +33,21 @@ Requires-Dist: ddgs>=9.9.3
 Requires-Dist: flashrank>=0.2.8
 Requires-Dist: gepa>=0.0.17
 Requires-Dist: grandalf>=0.8
-Requires-Dist: langchain-mcp-adapters>=0.1.10
+Requires-Dist: langchain-mcp-adapters>=0.2.1
 Requires-Dist: langchain-tavily>=0.2.11
 Requires-Dist: langchain>=1.1.3
-Requires-Dist: langgraph-checkpoint-postgres>=2.0.25
+Requires-Dist: langgraph-checkpoint-postgres>=3.0.2
 Requires-Dist: langgraph-supervisor>=0.0.31
 Requires-Dist: langgraph-swarm>=0.1.0
 Requires-Dist: langgraph>=1.0.4
-Requires-Dist: langmem>=0.0.29
+Requires-Dist: langmem>=0.0.30
 Requires-Dist: loguru>=0.7.3
-Requires-Dist: mcp>=1.17.0
+Requires-Dist: mcp>=1.23.3
 Requires-Dist: mlflow>=3.7.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: psycopg[binary,pool]>=3.3.2
 Requires-Dist: pydantic>=2.12.0
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: pyyaml>=6.0.2

{dao_ai-0.0.32.dist-info → dao_ai-0.0.34.dist-info}/RECORD

@@ -3,16 +3,16 @@ dao_ai/agent_as_code.py,sha256=sviZQV7ZPxE5zkZ9jAbfegI681nra5i8yYxw05e3X7U,552
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/chat_models.py,sha256=uhwwOTeLyHWqoTTgHrs4n5iSyTwe4EQcLKnh3jRxPWI,8626
 dao_ai/cli.py,sha256=gq-nsapWxDA1M6Jua3vajBvIwf0Oa6YLcB58lEtMKUo,22503
-dao_ai/config.py,sha256=sc9iYPui5tHitG5kmOTd9LVjzgLJ2Dn0M6s-Zu3dw04,75022
+dao_ai/config.py,sha256=Jzb0ePrt2TM2WuXI_LtmTafbseKBlJ8J8J2ExyBowbM,79491
 dao_ai/graph.py,sha256=9kjJx0oFZKq5J9-Kpri4-0VCJILHYdYyhqQnj0_noxQ,8913
 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/models.py,sha256=hvEZO2N0EC2sQoMgjJ9mbKmDWcdxnnAb2NqzpXh4Wgk,32691
 dao_ai/nodes.py,sha256=iQ_5vL6mt1UcRnhwgz-l1D8Ww4CMQrSMVnP_Lu7fFjU,8781
 dao_ai/prompts.py,sha256=iA2Iaky7yzjwWT5cxg0cUIgwo1z1UVQua__8WPnvV6g,1633
 dao_ai/state.py,sha256=_lF9krAYYjvFDMUwZzVKOn0ZnXKcOrbjWKdre0C5B54,1137
 dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dao_ai/utils.py,sha256=FLXbiUaCeBva4vJ-czs-sRP7QSxjoKjyDt1Q4yeI7sU,7727
+dao_ai/utils.py,sha256=oIPmz02kZ3LMntbqxUajFXh4nswOhbvEjOTi4e5_cvI,8500
 dao_ai/vector_search.py,sha256=jlaFS_iizJ55wblgzZmswMM3UOL-qOp2BGJc0JqXYSg,2839
 dao_ai/hooks/__init__.py,sha256=LlHGIuiZt6vGW8K5AQo1XJEkBP5vDVtMhq0IdjcLrD4,417
 dao_ai/hooks/core.py,sha256=ZShHctUSoauhBgdf1cecy9-D7J6-sGn-pKjuRMumW5U,6663
@@ -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=rPBMdGcJvdGBRK9FZeBxkLfcTpXyxU1cs14YllyZKbY,67857
+dao_ai/providers/databricks.py,sha256=WEigNPGRTlIPVjwp97My8o1zOHn5ftuMsMrpqrBeaLg,66012
 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=BPM_1Sk5bf7QSCFPPboWWkZKYwBwDwbGhMVp5-QDd10,5956
+dao_ai/tools/core.py,sha256=kN77fWOzVY7qOs4NiW72cUxCsSTC0DnPp73s6VJEZOQ,1991
+dao_ai/tools/genie.py,sha256=hWDLLGUNz1wgwOb69pXnMiLJnMbG_1YmMdfVKt1Qe8o,6426
 dao_ai/tools/human_in_the_loop.py,sha256=yk35MO9eNETnYFH-sqlgR-G24TrEgXpJlnZUustsLkI,3681
 dao_ai/tools/mcp.py,sha256=5aQoRtx2z4xm6zgRslc78rSfEQe-mfhqov2NsiybYfc,8416
 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/unity_catalog.py,sha256=K9t8M4spsbxbecWmV5yEZy16s_AG7AfaoxT-7IDW43I,14438
 dao_ai/tools/vector_search.py,sha256=3cdiUaFpox25GSRNec7FKceY3DuLp7dLVH8FRA0BgeY,12624
-dao_ai-0.0.32.dist-info/METADATA,sha256=1_BlILYdzDHCILhIxFNeWdM6CRg4uKqBNPiP_hjbXtE,42763
-dao_ai-0.0.32.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.0.32.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.32.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.32.dist-info/RECORD,,
+dao_ai-0.0.34.dist-info/METADATA,sha256=vq51NEV-pg7WTOD5z56jyOrC5_6Q-nUIL51RI5lL-Hg,42761
+dao_ai-0.0.34.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.0.34.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.34.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.34.dist-info/RECORD,,