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

dao_ai/config.py

@@ -725,6 +725,46 @@ class WarehouseModel(BaseModel, IsDatabricksResource):
 
 
 class DatabaseModel(BaseModel, IsDatabricksResource):
+    """
+    Configuration for a Databricks Lakebase (PostgreSQL) database instance.
+
+    Authentication Model:
+    --------------------
+    This model uses TWO separate authentication contexts:
+
+    1. **Workspace API Authentication** (inherited from IsDatabricksResource):
+       - Uses ambient/default authentication (environment variables, notebook context, app service principal)
+       - 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):
+       - Used for: connecting to the PostgreSQL database as a specific identity
+       - 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:
+    ```yaml
+    databases:
+      my_lakebase:
+        name: my-database
+        client_id:
+          env: SERVICE_PRINCIPAL_CLIENT_ID
+        client_secret:
+          scope: my-scope
+          secret: sp-client-secret
+        workspace_host:
+          env: DATABRICKS_HOST
+    ```
+
+    Example User Configuration:
+    ```yaml
+    databases:
+      my_lakebase:
+        name: my-database
+        user: my-user@databricks.com
+    ```
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     instance_name: Optional[str] = None
@@ -883,7 +923,7 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
     def create(self, w: WorkspaceClient | None = None) -> None:
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: DatabricksProvider = DatabricksProvider()
+        provider: DatabricksProvider = DatabricksProvider(w=w)
         provider.create_lakebase(self)
         provider.create_lakebase_instance_role(self)
 
@@ -1613,6 +1653,12 @@ class AppModel(BaseModel):
     chat_history: Optional[ChatHistoryModel] = None
     code_paths: list[str] = Field(default_factory=list)
     pip_requirements: list[str] = Field(default_factory=list)
+    python_version: Optional[str] = Field(
+        default="3.12",
+        description="Python version for Model Serving deployment. Defaults to 3.12 "
+        "which is supported by Databricks Model Serving. This allows deploying from "
+        "environments with different Python versions (e.g., Databricks Apps with 3.11).",
+    )
 
     @model_validator(mode="after")
     def validate_agents_not_empty(self):

dao_ai/prompts.py

@@ -1,10 +1,10 @@
 from typing import Any, Callable, Optional, Sequence
 
-from langchain.prompts import PromptTemplate
 from langchain_core.messages import (
     BaseMessage,
     SystemMessage,
 )
+from langchain_core.prompts import PromptTemplate
 from langchain_core.runnables import RunnableConfig
 from loguru import logger
 

dao_ai/providers/databricks.py

@@ -332,6 +332,23 @@ class DatabricksProvider(ServiceProvider):
 
         logger.debug(f"input_example: {input_example}")
 
+        # Create conda environment with configured Python version
+        # This allows deploying from environments with different Python versions
+        # (e.g., Databricks Apps with Python 3.11 can deploy to Model Serving with 3.12)
+        target_python_version: str = config.app.python_version
+        logger.debug(f"target_python_version: {target_python_version}")
+
+        conda_env: dict[str, Any] = {
+            "name": "mlflow-env",
+            "channels": ["conda-forge"],
+            "dependencies": [
+                f"python={target_python_version}",
+                "pip",
+                {"pip": list(pip_requirements)},
+            ],
+        }
+        logger.debug(f"conda_env: {conda_env}")
+
         with mlflow.start_run(run_name=run_name):
             mlflow.set_tag("type", "agent")
             mlflow.set_tag("dao_ai", dao_ai_version())
@@ -340,7 +357,7 @@ class DatabricksProvider(ServiceProvider):
                 code_paths=code_paths,
                 model_config=config.model_dump(mode="json", by_alias=True),
                 name="agent",
-                pip_requirements=pip_requirements,
+                conda_env=conda_env,
                 input_example=input_example,
                 # resources=all_resources,
                 auth_policy=auth_policy,
@@ -773,6 +790,72 @@ class DatabricksProvider(ServiceProvider):
         logger.debug(f"Vector search index found: {found_endpoint_name}")
         return found_endpoint_name
 
+    def _wait_for_database_available(
+        self,
+        workspace_client: WorkspaceClient,
+        instance_name: str,
+        max_wait_time: int = 600,
+        wait_interval: int = 10,
+    ) -> None:
+        """
+        Wait for a database instance to become AVAILABLE.
+
+        Args:
+            workspace_client: The Databricks workspace client
+            instance_name: Name of the database instance to wait for
+            max_wait_time: Maximum time to wait in seconds (default: 600 = 10 minutes)
+            wait_interval: Time between status checks in seconds (default: 10)
+
+        Raises:
+            TimeoutError: If the database doesn't become AVAILABLE within max_wait_time
+            RuntimeError: If the database enters a failed or deleted state
+        """
+        import time
+        from typing import Any
+
+        logger.info(
+            f"Waiting for database instance {instance_name} to become AVAILABLE..."
+        )
+        elapsed: int = 0
+
+        while elapsed < max_wait_time:
+            try:
+                current_instance: Any = workspace_client.database.get_database_instance(
+                    name=instance_name
+                )
+                current_state: str = current_instance.state
+                logger.debug(
+                    f"Database instance {instance_name} state: {current_state}"
+                )
+
+                if current_state == "AVAILABLE":
+                    logger.info(f"Database instance {instance_name} is now AVAILABLE")
+                    return
+                elif current_state in ["STARTING", "UPDATING", "PROVISIONING"]:
+                    logger.debug(
+                        f"Database instance still in {current_state} state, waiting {wait_interval} seconds..."
+                    )
+                    time.sleep(wait_interval)
+                    elapsed += wait_interval
+                elif current_state in ["STOPPED", "DELETING", "FAILED"]:
+                    raise RuntimeError(
+                        f"Database instance {instance_name} entered unexpected state: {current_state}"
+                    )
+                else:
+                    logger.warning(
+                        f"Unknown database state: {current_state}, continuing to wait..."
+                    )
+                    time.sleep(wait_interval)
+                    elapsed += wait_interval
+            except NotFound:
+                raise RuntimeError(
+                    f"Database instance {instance_name} was deleted while waiting for it to become AVAILABLE"
+                )
+
+        raise TimeoutError(
+            f"Timed out waiting for database instance {instance_name} to become AVAILABLE after {max_wait_time} seconds"
+        )
+
     def create_lakebase(self, database: DatabaseModel) -> None:
         """
         Create a Lakebase database instance using the Databricks workspace client.
@@ -907,6 +990,12 @@ class DatabricksProvider(ServiceProvider):
                     f"Successfully created database instance: {database.instance_name}"
                 )
 
+                # Wait for the newly created database to become AVAILABLE
+                self._wait_for_database_available(
+                    workspace_client, database.instance_name
+                )
+                return
+
             except Exception as create_error:
                 error_msg: str = str(create_error)
 
@@ -918,6 +1007,10 @@ class DatabricksProvider(ServiceProvider):
                     logger.info(
                         f"Database instance {database.instance_name} was created concurrently by another process"
                     )
+                    # Still need to wait for the database to become AVAILABLE
+                    self._wait_for_database_available(
+                        workspace_client, database.instance_name
+                    )
                     return
                 else:
                     # Re-raise unexpected errors
@@ -1057,9 +1150,10 @@ 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 (if it exists)
-        2. latest alias (if it exists)
-        3. default_template (if provided)
+        1. champion alias
+        2. latest version (max version number from search_prompt_versions)
+        3. default alias
+        4. Register default_template if provided
 
         Args:
             prompt_model: The prompt model configuration
@@ -1070,9 +1164,11 @@ class DatabricksProvider(ServiceProvider):
         Raises:
             ValueError: If no prompt can be loaded from any source
         """
+
         prompt_name: str = prompt_model.full_name
+        mlflow_client: MlflowClient = MlflowClient()
 
-        # If explicit version or alias is specified, use it directly without fallback
+        # If explicit version or alias is specified, use it directly
         if prompt_model.version or prompt_model.alias:
             try:
                 prompt_version: PromptVersion = prompt_model.as_prompt()
@@ -1086,73 +1182,48 @@ class DatabricksProvider(ServiceProvider):
                     f"Failed to load prompt '{prompt_name}' with explicit "
                     f"{'version ' + str(prompt_model.version) if prompt_model.version else 'alias ' + prompt_model.alias}: {e}"
                 )
-                # Fall through to default_template if available
-        else:
-            # No explicit version/alias specified - check if default_template needs syncing first
-            logger.debug(
-                f"No explicit version/alias specified for '{prompt_name}', "
-                "checking if default_template needs syncing"
-            )
-
-            # If we have a default_template, check if it differs from what's in the registry
-            # This ensures we always sync config changes before returning any alias
-            if prompt_model.default_template:
-                try:
-                    default_uri: str = f"prompts:/{prompt_name}@default"
-                    default_version: PromptVersion = load_prompt(default_uri)
-
-                    if (
-                        default_version.to_single_brace_format().strip()
-                        != prompt_model.default_template.strip()
-                    ):
-                        logger.info(
-                            f"Config default_template for '{prompt_name}' differs from registry, syncing..."
-                        )
-                        return self._sync_default_template_to_registry(
-                            prompt_name,
-                            prompt_model.default_template,
-                            prompt_model.description,
-                        )
-                except Exception as e:
-                    logger.debug(f"Could not check default alias for sync: {e}")
+                # Fall through to try other methods
 
-            # Now try aliases in order: champion → latest → default
-            logger.debug(
-                f"Trying fallback order for '{prompt_name}': champion → latest → default"
-            )
+        # Try to load in priority order: champion → latest → default
+        logger.debug(
+            f"Trying fallback order for '{prompt_name}': champion → latest → default"
+        )
 
-            # Try champion alias first
-            try:
-                champion_uri: str = f"prompts:/{prompt_name}@champion"
-                prompt_version: PromptVersion = load_prompt(champion_uri)
-                logger.info(f"Loaded prompt '{prompt_name}' from champion alias")
-                return prompt_version
-            except Exception as e:
-                logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
+        # 1. Try champion alias
+        try:
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@champion")
+            logger.info(f"Loaded prompt '{prompt_name}' from champion alias")
+            return prompt_version
+        except Exception as e:
+            logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
 
-            # Try latest alias next
-            try:
-                latest_uri: str = f"prompts:/{prompt_name}@latest"
-                prompt_version: PromptVersion = load_prompt(latest_uri)
-                logger.info(f"Loaded prompt '{prompt_name}' from latest alias")
-                return prompt_version
-            except Exception as e:
-                logger.debug(f"Latest alias not found for '{prompt_name}': {e}")
+        # 2. Try to get 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 find latest version for '{prompt_name}': {e}")
 
-            # Try default alias last
-            try:
-                default_uri: str = f"prompts:/{prompt_name}@default"
-                prompt_version: PromptVersion = load_prompt(default_uri)
-                logger.info(f"Loaded prompt '{prompt_name}' from default alias")
-                return prompt_version
-            except Exception as e:
-                logger.debug(f"Default alias not found for '{prompt_name}': {e}")
+        # 3. Try default alias
+        try:
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@default")
+            logger.info(f"Loaded prompt '{prompt_name}' from default alias")
+            return prompt_version
+        except Exception as e:
+            logger.debug(f"Default alias not found for '{prompt_name}': {e}")
 
-        # Fall back to registering default_template if provided
+        # 4. Try to register default_template if provided
         if prompt_model.default_template:
             logger.info(
-                f"Registering default_template for '{prompt_name}' "
-                "(no aliases found in registry)"
+                f"No existing prompt found for '{prompt_name}', "
+                "attempting to register default_template"
             )
             return self._sync_default_template_to_registry(
                 prompt_name, prompt_model.default_template, prompt_model.description
@@ -1160,72 +1231,58 @@ class DatabricksProvider(ServiceProvider):
 
         raise ValueError(
             f"Prompt '{prompt_name}' not found in registry "
-            "(tried champion, latest, default aliases) and no default_template provided"
+            "(tried champion alias, latest version, default alias) "
+            "and no default_template provided"
         )
 
     def _sync_default_template_to_registry(
         self, prompt_name: str, default_template: str, description: str | None = None
     ) -> PromptVersion:
-        """Register default_template to prompt registry under 'default' alias if changed."""
-        prompt_version: PromptVersion | None = None
+        """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:
-            # Check if default alias already has the same template
-            try:
-                logger.debug(f"Loading prompt '{prompt_name}' from registry...")
-                existing: PromptVersion = mlflow.genai.load_prompt(
-                    f"prompts:/{prompt_name}@default"
-                )
-                if (
-                    existing.to_single_brace_format().strip()
-                    == default_template.strip()
-                ):
-                    logger.debug(f"Prompt '{prompt_name}' is already up-to-date")
+            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}")
 
-                    # Ensure the "latest" and "champion" aliases also exist and point to the same version
-                    # This handles prompts created before the fix that added these aliases
-                    try:
-                        latest_version: PromptVersion = mlflow.genai.load_prompt(
-                            f"prompts:/{prompt_name}@latest"
-                        )
-                        logger.debug(
-                            f"Latest alias already exists for '{prompt_name}' pointing to version {latest_version.version}"
-                        )
-                    except Exception:
-                        logger.info(
-                            f"Setting 'latest' alias for existing prompt '{prompt_name}' v{existing.version}"
-                        )
-                        mlflow.genai.set_prompt_alias(
-                            name=prompt_name,
-                            alias="latest",
-                            version=existing.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}")
 
-                    # Ensure champion alias exists for first-time deployments
-                    try:
-                        champion_version: PromptVersion = mlflow.genai.load_prompt(
-                            f"prompts:/{prompt_name}@champion"
-                        )
-                        logger.debug(
-                            f"Champion alias already exists for '{prompt_name}' pointing to version {champion_version.version}"
-                        )
-                    except Exception:
-                        logger.info(
-                            f"Setting 'champion' alias for existing prompt '{prompt_name}' v{existing.version}"
-                        )
-                        mlflow.genai.set_prompt_alias(
-                            name=prompt_name,
-                            alias="champion",
-                            version=existing.version,
-                        )
+        # 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}")
 
-                    return existing  # Already up-to-date, return existing version
-            except Exception:
-                logger.debug(
-                    f"Default alias for prompt '{prompt_name}' doesn't exist yet"
-                )
+        # No existing prompt found - try to register if we have a template
+        logger.info(
+            f"No existing prompt found for '{prompt_name}', attempting to register default_template"
+        )
 
-            # Register new version and set as default alias
+        try:
             commit_message = description or "Auto-synced from default_template"
             prompt_version = mlflow.genai.register_prompt(
                 name=prompt_name,
@@ -1234,35 +1291,36 @@ class DatabricksProvider(ServiceProvider):
                 tags={"dao_ai": dao_ai_version()},
             )
 
-            logger.debug(
-                f"Setting default, latest, and champion aliases for prompt '{prompt_name}'"
-            )
-            mlflow.genai.set_prompt_alias(
-                name=prompt_name,
-                alias="default",
-                version=prompt_version.version,
-            )
-            mlflow.genai.set_prompt_alias(
-                name=prompt_name,
-                alias="latest",
-                version=prompt_version.version,
-            )
-            mlflow.genai.set_prompt_alias(
-                name=prompt_name,
-                alias="champion",
-                version=prompt_version.version,
-            )
+            # Try to set aliases (may fail in restricted environments)
+            try:
+                mlflow.genai.set_prompt_alias(
+                    name=prompt_name, alias="default", version=prompt_version.version
+                )
+                mlflow.genai.set_prompt_alias(
+                    name=prompt_name, alias="champion", version=prompt_version.version
+                )
+                logger.info(
+                    f"Registered prompt '{prompt_name}' v{prompt_version.version} with aliases"
+                )
+            except Exception as alias_error:
+                logger.warning(
+                    f"Registered prompt '{prompt_name}' v{prompt_version.version} "
+                    f"but failed to set aliases: {alias_error}"
+                )
 
-            logger.info(
-                f"Synced prompt '{prompt_name}' v{prompt_version.version} to registry with 'default', 'latest', and 'champion' aliases"
-            )
             return prompt_version
 
-        except Exception as e:
-            logger.error(f"Failed to sync '{prompt_name}' to registry: {e}")
-            raise ValueError(
-                f"Failed to sync prompt '{prompt_name}' to registry and unable to retrieve existing version"
-            ) from e
+        except Exception as reg_error:
+            logger.error(
+                f"Failed to register prompt '{prompt_name}': {reg_error}. "
+                f"Please register the prompt from a notebook with write permissions before deployment."
+            )
+            return PromptVersion(
+                name=prompt_name,
+                version=1,
+                template=default_template,
+                tags={"dao_ai": dao_ai_version()},
+            )
 
     def optimize_prompt(self, optimization: PromptOptimizationModel) -> PromptModel:
         """

dao_ai/tools/genie.py

@@ -1,15 +1,10 @@
-import bisect
 import json
 import os
-import time
-from dataclasses import asdict, dataclass
-from datetime import datetime
 from textwrap import dedent
-from typing import Annotated, Any, Callable, Optional, Union
+from typing import Annotated, Any, Callable
 
-import mlflow
 import pandas as pd
-from databricks.sdk import WorkspaceClient
+from databricks_ai_bridge.genie import Genie, GenieResponse
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import InjectedToolCallId, tool
 from langgraph.prebuilt import InjectedState
@@ -19,28 +14,6 @@ from pydantic import BaseModel, Field
 
 from dao_ai.config import AnyVariable, CompositeVariableModel, GenieRoomModel, value_of
 
-MAX_TOKENS_OF_DATA: int = 20000
-MAX_ITERATIONS: int = 50
-DEFAULT_POLLING_INTERVAL_SECS: int = 2
-
-
-def _count_tokens(text):
-    import tiktoken
-
-    encoding = tiktoken.encoding_for_model("gpt-4o")
-    return len(encoding.encode(text))
-
-
-@dataclass
-class GenieResponse:
-    conversation_id: str
-    result: Union[str, pd.DataFrame]
-    query: Optional[str] = ""
-    description: Optional[str] = ""
-
-    def to_json(self):
-        return json.dumps(asdict(self))
-
 
 class GenieToolInput(BaseModel):
     """Input schema for the Genie tool."""
@@ -50,235 +23,29 @@ class GenieToolInput(BaseModel):
     )
 
 
-def _truncate_result(dataframe: pd.DataFrame) -> str:
-    query_result = dataframe.to_markdown()
-    tokens_used = _count_tokens(query_result)
-
-    # If the full result fits, return it
-    if tokens_used <= MAX_TOKENS_OF_DATA:
-        return query_result.strip()
-
-    def is_too_big(n):
-        return _count_tokens(dataframe.iloc[:n].to_markdown()) > MAX_TOKENS_OF_DATA
-
-    # Use bisect_left to find the cutoff point of rows within the max token data limit in a O(log n) complexity
-    # Passing True, as this is the target value we are looking for when _is_too_big returns
-    cutoff = bisect.bisect_left(range(len(dataframe) + 1), True, key=is_too_big)
-
-    # Slice to the found limit
-    truncated_df = dataframe.iloc[:cutoff]
-
-    # Edge case: Cannot return any rows because of tokens so return an empty string
-    if len(truncated_df) == 0:
-        return ""
-
-    truncated_result = truncated_df.to_markdown()
-
-    # Double-check edge case if we overshot by one
-    if _count_tokens(truncated_result) > MAX_TOKENS_OF_DATA:
-        truncated_result = truncated_df.iloc[:-1].to_markdown()
-    return truncated_result
-
-
-@mlflow.trace(span_type="PARSER")
-def _parse_query_result(resp, truncate_results) -> Union[str, pd.DataFrame]:
-    output = resp["result"]
-    if not output:
-        return "EMPTY"
-
-    columns = resp["manifest"]["schema"]["columns"]
-    header = [str(col["name"]) for col in columns]
-    rows = []
-
-    for item in output["data_array"]:
-        row = []
-        for column, value in zip(columns, item):
-            type_name = column["type_name"]
-            if value is None:
-                row.append(None)
-                continue
-
-            if type_name in ["INT", "LONG", "SHORT", "BYTE"]:
-                row.append(int(value))
-            elif type_name in ["FLOAT", "DOUBLE", "DECIMAL"]:
-                row.append(float(value))
-            elif type_name == "BOOLEAN":
-                row.append(value.lower() == "true")
-            elif type_name == "DATE" or type_name == "TIMESTAMP":
-                row.append(datetime.strptime(value[:10], "%Y-%m-%d").date())
-            elif type_name == "BINARY":
-                row.append(bytes(value, "utf-8"))
-            else:
-                row.append(value)
-
-        rows.append(row)
+def _response_to_json(response: GenieResponse) -> str:
+    """Convert GenieResponse to JSON string, handling DataFrame results."""
+    # Convert result to string if it's a DataFrame
+    result: str | pd.DataFrame = response.result
+    if isinstance(result, pd.DataFrame):
+        result = result.to_markdown()
 
-    dataframe = pd.DataFrame(rows, columns=header)
-
-    if truncate_results:
-        query_result = _truncate_result(dataframe)
-    else:
-        query_result = dataframe.to_markdown()
-
-    return query_result.strip()
-
-
-class Genie:
-    def __init__(
-        self,
-        space_id,
-        client: WorkspaceClient | None = None,
-        truncate_results: bool = False,
-        polling_interval: int = DEFAULT_POLLING_INTERVAL_SECS,
-    ):
-        self.space_id = space_id
-        workspace_client = client or WorkspaceClient()
-        self.genie = workspace_client.genie
-        self.description = self.genie.get_space(space_id).description
-        self.headers = {
-            "Accept": "application/json",
-            "Content-Type": "application/json",
-        }
-        self.truncate_results = truncate_results
-        if polling_interval < 1 or polling_interval > 30:
-            raise ValueError("poll_interval must be between 1 and 30 seconds")
-        self.poll_interval = polling_interval
-
-    @mlflow.trace()
-    def start_conversation(self, content):
-        resp = self.genie._api.do(
-            "POST",
-            f"/api/2.0/genie/spaces/{self.space_id}/start-conversation",
-            body={"content": content},
-            headers=self.headers,
-        )
-        return resp
-
-    @mlflow.trace()
-    def create_message(self, conversation_id, content):
-        resp = self.genie._api.do(
-            "POST",
-            f"/api/2.0/genie/spaces/{self.space_id}/conversations/{conversation_id}/messages",
-            body={"content": content},
-            headers=self.headers,
-        )
-        return resp
-
-    @mlflow.trace()
-    def poll_for_result(self, conversation_id, message_id):
-        @mlflow.trace()
-        def poll_query_results(attachment_id, query_str, description):
-            iteration_count = 0
-            while iteration_count < MAX_ITERATIONS:
-                iteration_count += 1
-                resp = self.genie._api.do(
-                    "GET",
-                    f"/api/2.0/genie/spaces/{self.space_id}/conversations/{conversation_id}/messages/{message_id}/attachments/{attachment_id}/query-result",
-                    headers=self.headers,
-                )["statement_response"]
-                state = resp["status"]["state"]
-                if state == "SUCCEEDED":
-                    result = _parse_query_result(resp, self.truncate_results)
-                    return GenieResponse(
-                        conversation_id, result, query_str, description
-                    )
-                elif state in ["RUNNING", "PENDING"]:
-                    logger.debug("Waiting for query result...")
-                    time.sleep(self.poll_interval)
-                else:
-                    return GenieResponse(
-                        conversation_id,
-                        f"No query result: {resp['state']}",
-                        query_str,
-                        description,
-                    )
-            return GenieResponse(
-                conversation_id,
-                f"Genie query for result timed out after {MAX_ITERATIONS} iterations of {self.poll_interval} seconds",
-                query_str,
-                description,
-            )
-
-        @mlflow.trace()
-        def poll_result():
-            iteration_count = 0
-            while iteration_count < MAX_ITERATIONS:
-                iteration_count += 1
-                resp = self.genie._api.do(
-                    "GET",
-                    f"/api/2.0/genie/spaces/{self.space_id}/conversations/{conversation_id}/messages/{message_id}",
-                    headers=self.headers,
-                )
-                if resp["status"] == "COMPLETED":
-                    # Check if attachments key exists in response
-                    attachments = resp.get("attachments", [])
-                    if not attachments:
-                        # Handle case where response has no attachments
-                        return GenieResponse(
-                            conversation_id,
-                            result=f"Genie query completed but no attachments found. Response: {resp}",
-                        )
-
-                    attachment = next((r for r in attachments if "query" in r), None)
-                    if attachment:
-                        query_obj = attachment["query"]
-                        description = query_obj.get("description", "")
-                        query_str = query_obj.get("query", "")
-                        attachment_id = attachment["attachment_id"]
-                        return poll_query_results(attachment_id, query_str, description)
-                    if resp["status"] == "COMPLETED":
-                        text_content = next(
-                            (r for r in attachments if "text" in r), None
-                        )
-                        if text_content:
-                            return GenieResponse(
-                                conversation_id, result=text_content["text"]["content"]
-                            )
-                        return GenieResponse(
-                            conversation_id,
-                            result="Genie query completed but no text content found in attachments.",
-                        )
-                elif resp["status"] in {"CANCELLED", "QUERY_RESULT_EXPIRED"}:
-                    return GenieResponse(
-                        conversation_id, result=f"Genie query {resp['status'].lower()}."
-                    )
-                elif resp["status"] == "FAILED":
-                    return GenieResponse(
-                        conversation_id,
-                        result=f"Genie query failed with error: {resp.get('error', 'Unknown error')}",
-                    )
-                # includes EXECUTING_QUERY, Genie can retry after this status
-                else:
-                    logger.debug(f"Waiting...: {resp['status']}")
-                    time.sleep(self.poll_interval)
-            return GenieResponse(
-                conversation_id,
-                f"Genie query timed out after {MAX_ITERATIONS} iterations of {self.poll_interval} seconds",
-            )
-
-        return poll_result()
-
-    @mlflow.trace()
-    def ask_question(self, question: str, conversation_id: str | None = None):
-        logger.debug(
-            f"ask_question called with question: {question}, conversation_id: {conversation_id}"
-        )
-        if conversation_id:
-            resp = self.create_message(conversation_id, question)
-        else:
-            resp = self.start_conversation(question)
-        logger.debug(f"ask_question response: {resp}")
-        return self.poll_for_result(resp["conversation_id"], resp["message_id"])
+    data: dict[str, Any] = {
+        "result": result,
+        "query": response.query,
+        "description": response.description,
+        "conversation_id": response.conversation_id,
+    }
+    return json.dumps(data)
 
 
 def create_genie_tool(
     genie_room: GenieRoomModel | dict[str, Any],
-    name: Optional[str] = None,
-    description: Optional[str] = None,
+    name: str | None = None,
+    description: str | None = None,
     persist_conversation: bool = False,
     truncate_results: bool = False,
-    poll_interval: int = DEFAULT_POLLING_INTERVAL_SECS,
-) -> Callable[[str], GenieResponse]:
+) -> Callable[..., Command]:
     """
     Create a tool for interacting with Databricks Genie for natural language queries to databases.
 
@@ -290,6 +57,9 @@ def create_genie_tool(
         genie_room: GenieRoomModel or dict containing Genie configuration
         name: Optional custom name for the tool. If None, uses default "genie_tool"
         description: Optional custom description for the tool. If None, uses default description
+        persist_conversation: Whether to persist conversation IDs across tool calls for
+            multi-turn conversations within the same Genie space
+        truncate_results: Whether to truncate large query results to fit token limits
 
     Returns:
         A LangGraph tool that processes natural language queries through Genie
@@ -305,13 +75,6 @@ def create_genie_tool(
         space_id = CompositeVariableModel(**space_id)
     space_id = value_of(space_id)
 
-    # genie: Genie = Genie(
-    #     space_id=space_id,
-    #     client=genie_room.workspace_client,
-    #     truncate_results=truncate_results,
-    #     polling_interval=poll_interval,
-    # )
-
     default_description: str = dedent("""
     This tool lets you have a conversation and chat with tabular data about <topic>. You should ask
     questions about the data and the tool will try to answer them.
@@ -343,14 +106,14 @@ 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,
-            polling_interval=poll_interval,
         )
 
-        """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)
@@ -368,9 +131,10 @@ 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(), tool_call_id=tool_call_id)],
+            "messages": [
+                ToolMessage(_response_to_json(response), tool_call_id=tool_call_id)
+            ],
         }
 
         if persist_conversation:

dao_ai/tools/vector_search.py

@@ -101,7 +101,7 @@ 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.
@@ -121,7 +121,9 @@ def create_vector_search_tool(
             "DATABRICKS_CLIENT_SECRET"
         )
 
-    logger.debug(f"Creating DatabricksVectorSearch with client_args keys: {list(client_args.keys())}")
+    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(

dao_ai/utils.py

@@ -99,7 +99,7 @@ def get_installed_packages() -> dict[str, str]:
         f"databricks-langchain=={version('databricks-langchain')}",
         f"databricks-mcp=={version('databricks-mcp')}",
         f"databricks-sdk[openai]=={version('databricks-sdk')}",
-        f"duckduckgo-search=={version('duckduckgo-search')}",
+        f"ddgs=={version('ddgs')}",
         f"flashrank=={version('flashrank')}",
         f"langchain=={version('langchain')}",
         f"langchain-mcp-adapters=={version('langchain-mcp-adapters')}",
@@ -141,12 +141,12 @@ def load_function(function_name: str) -> Callable[..., Any]:
              "module.submodule.function_name"
 
     Returns:
-        The imported callable function
+        The imported callable function or langchain tool
 
     Raises:
         ImportError: If the module cannot be imported
         AttributeError: If the function doesn't exist in the module
-        TypeError: If the resolved object is not callable
+        TypeError: If the resolved object is not callable or invocable
 
     Example:
         >>> func = callable_from_fqn("dao_ai.models.get_latest_model_version")
@@ -164,9 +164,14 @@ def load_function(function_name: str) -> Callable[..., Any]:
         # Get the function from the module
         func = getattr(module, func_name)
 
-        # Verify that the resolved object is callable
-        if not callable(func):
-            raise TypeError(f"Function {func_name} is not callable.")
+        # Verify that the resolved object is callable or is a langchain tool
+        # In langchain 1.x, StructuredTool objects are not directly callable
+        # but have an invoke() method
+        is_callable = callable(func)
+        is_langchain_tool = hasattr(func, "invoke") and hasattr(func, "name")
+
+        if not is_callable and not is_langchain_tool:
+            raise TypeError(f"Function {func_name} is not callable or invocable.")
 
         return func
     except (ImportError, AttributeError, TypeError) as e:
@@ -175,4 +180,26 @@ def load_function(function_name: str) -> Callable[..., Any]:
 
 
 def is_in_model_serving() -> bool:
-    return os.environ.get("IS_IN_DB_MODEL_SERVING_ENV", "false").lower() == "true"
+    """Check if running in Databricks Model Serving environment.
+
+    Detects Model Serving by checking for environment variables that are
+    typically set in that environment.
+    """
+    # Primary check - explicit Databricks Model Serving env var
+    if os.environ.get("IS_IN_DB_MODEL_SERVING_ENV", "false").lower() == "true":
+        return True
+
+    # Secondary check - Model Serving sets these environment variables
+    if os.environ.get("DATABRICKS_MODEL_SERVING_ENV"):
+        return True
+
+    # Check for cluster type indicator
+    cluster_type = os.environ.get("DATABRICKS_CLUSTER_TYPE", "")
+    if "model-serving" in cluster_type.lower():
+        return True
+
+    # Check for model serving specific paths
+    if os.path.exists("/opt/conda/envs/mlflow-env"):
+        return True
+
+    return False

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.31
+Version: 0.0.32
 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
@@ -26,24 +26,24 @@ 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-langchain>=0.8.1
+Requires-Dist: databricks-langchain>=0.11.0
 Requires-Dist: databricks-mcp>=0.3.0
 Requires-Dist: databricks-sdk[openai]>=0.67.0
-Requires-Dist: duckduckgo-search>=8.0.2
+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-tavily>=0.2.11
-Requires-Dist: langchain>=0.3.27
+Requires-Dist: langchain>=1.1.3
 Requires-Dist: langgraph-checkpoint-postgres>=2.0.25
-Requires-Dist: langgraph-supervisor>=0.0.29
-Requires-Dist: langgraph-swarm>=0.0.14
-Requires-Dist: langgraph>=0.6.10
+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: loguru>=0.7.3
 Requires-Dist: mcp>=1.17.0
-Requires-Dist: mlflow>=3.5.1
+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
@@ -55,7 +55,7 @@ Requires-Dist: rich>=14.0.0
 Requires-Dist: scipy<=1.15
 Requires-Dist: sqlparse>=0.5.3
 Requires-Dist: tomli>=2.3.0
-Requires-Dist: unitycatalog-ai[databricks]>=0.3.0
+Requires-Dist: unitycatalog-ai[databricks]>=0.3.2
 Provides-Extra: databricks
 Requires-Dist: databricks-connect>=15.0.0; extra == 'databricks'
 Requires-Dist: databricks-vectorsearch>=0.63; extra == 'databricks'

{dao_ai-0.0.31.dist-info → dao_ai-0.0.32.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=r4mhailZn2edyJpElUGPvJS5rYyt6kclNWIK6RG328g,73279
+dao_ai/config.py,sha256=sc9iYPui5tHitG5kmOTd9LVjzgLJ2Dn0M6s-Zu3dw04,75022
 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/nodes.py,sha256=iQ_5vL6mt1UcRnhwgz-l1D8Ww4CMQrSMVnP_Lu7fFjU,8781
-dao_ai/prompts.py,sha256=7Hcstmv514P0s9s-TVoIlbkDV2XXOphGCW6gcPeyUYE,1628
+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=4FV9y0EVn0tmxfkn4EdUAkOewoAF_T0pHDAHe6hJx-M,6708
+dao_ai/utils.py,sha256=FLXbiUaCeBva4vJ-czs-sRP7QSxjoKjyDt1Q4yeI7sU,7727
 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=QwQS0Mad6nZNOB2k_h3lTrwLRzI5ApLQjxOz4CpvhlI,66130
+dao_ai/providers/databricks.py,sha256=rPBMdGcJvdGBRK9FZeBxkLfcTpXyxU1cs14YllyZKbY,67857
 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=wt6pVykukNStOUlrTyjWUj-S2Wb47IuMr5HztEcN3Wg,14925
+dao_ai/tools/genie.py,sha256=BPM_1Sk5bf7QSCFPPboWWkZKYwBwDwbGhMVp5-QDd10,5956
 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/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,,
+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,,