dao-ai 0.0.20__py3-none-any.whl → 0.0.22__py3-none-any.whl

dao_ai/config.py

@@ -23,6 +23,7 @@ from databricks.sdk.credentials_provider import (
     ModelServingUserCredentials,
 )
 from databricks.sdk.service.catalog import FunctionInfo, TableInfo
+from databricks.sdk.service.database import DatabaseInstance
 from databricks.vector_search.client import VectorSearchClient
 from databricks.vector_search.index import VectorSearchIndex
 from databricks_langchain import (
@@ -427,7 +428,8 @@ class GenieRoomModel(BaseModel, IsDatabricksResource):
     def as_resources(self) -> Sequence[DatabricksResource]:
         return [
             DatabricksGenieSpace(
-                genie_space_id=self.space_id, on_behalf_of_user=self.on_behalf_of_user
+                genie_space_id=value_of(self.space_id),
+                on_behalf_of_user=self.on_behalf_of_user,
             )
         ]
 
@@ -437,7 +439,7 @@ class GenieRoomModel(BaseModel, IsDatabricksResource):
         return self
 
 
-class VolumeModel(BaseModel, HasFullName):
+class VolumeModel(BaseModel, HasFullName, IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
     name: str
@@ -455,6 +457,13 @@ class VolumeModel(BaseModel, HasFullName):
         provider: ServiceProvider = DatabricksProvider(w=w)
         provider.create_volume(self)
 
+    @property
+    def api_scopes(self) -> Sequence[str]:
+        return ["files.files", "catalog.volumes"]
+
+    def as_resources(self) -> Sequence[DatabricksResource]:
+        return []
+
 
 class VolumePathModel(BaseModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -683,7 +692,8 @@ class WarehouseModel(BaseModel, IsDatabricksResource):
     def as_resources(self) -> Sequence[DatabricksResource]:
         return [
             DatabricksSQLWarehouse(
-                warehouse_id=self.warehouse_id, on_behalf_of_user=self.on_behalf_of_user
+                warehouse_id=value_of(self.warehouse_id),
+                on_behalf_of_user=self.on_behalf_of_user,
             )
         ]
 
@@ -694,15 +704,18 @@ class WarehouseModel(BaseModel, IsDatabricksResource):
 
 
 class DatabaseModel(BaseModel, IsDatabricksResource):
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
+    instance_name: Optional[str] = None
     description: Optional[str] = None
-    host: Optional[AnyVariable]
+    host: Optional[AnyVariable] = None
     database: Optional[AnyVariable] = "databricks_postgres"
     port: Optional[AnyVariable] = 5432
     connection_kwargs: Optional[dict[str, Any]] = Field(default_factory=dict)
     max_pool_size: Optional[int] = 10
-    timeout_seconds: Optional[int] = 5
+    timeout_seconds: Optional[int] = 10
+    capacity: Optional[Literal["CU_1", "CU_2"]] = "CU_2"
+    node_count: Optional[int] = None
     user: Optional[AnyVariable] = None
     password: Optional[AnyVariable] = None
     client_id: Optional[AnyVariable] = None
@@ -716,11 +729,44 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
     def as_resources(self) -> Sequence[DatabricksResource]:
         return [
             DatabricksLakebase(
-                database_instance_name=self.name,
+                database_instance_name=self.instance_name,
                 on_behalf_of_user=self.on_behalf_of_user,
             )
         ]
 
+    @model_validator(mode="after")
+    def update_instance_name(self):
+        if self.instance_name is None:
+            self.instance_name = self.name
+
+        return self
+
+    @model_validator(mode="after")
+    def update_user(self):
+        if self.client_id or self.user:
+            return self
+
+        self.user = self.workspace_client.current_user.me().user_name
+        if not self.user:
+            raise ValueError(
+                "Unable to determine current user. Please provide a user name or OAuth credentials."
+            )
+
+        return self
+
+    @model_validator(mode="after")
+    def update_host(self):
+        if self.host is not None:
+            return self
+
+        existing_instance: DatabaseInstance = (
+            self.workspace_client.database.get_database_instance(
+                name=self.instance_name
+            )
+        )
+        self.host = existing_instance.read_write_dns
+        return self
+
     @model_validator(mode="after")
     def validate_auth_methods(self):
         oauth_fields: Sequence[Any] = [
@@ -730,7 +776,7 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         ]
         has_oauth: bool = all(field is not None for field in oauth_fields)
 
-        pat_fields: Sequence[Any] = [self.user, self.password]
+        pat_fields: Sequence[Any] = [self.user]
         has_user_auth: bool = all(field is not None for field in pat_fields)
 
         if has_oauth and has_user_auth:
@@ -749,7 +795,14 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         return self
 
     @property
-    def connection_url(self) -> str:
+    def connection_params(self) -> dict[str, Any]:
+        """
+        Get database connection parameters as a dictionary.
+
+        Returns a dict with connection parameters suitable for psycopg ConnectionPool.
+        If username is configured, it will be included; otherwise it will be omitted
+        to allow Lakebase to authenticate using the token's identity.
+        """
         from dao_ai.providers.base import ServiceProvider
         from dao_ai.providers.databricks import DatabricksProvider
 
@@ -757,7 +810,7 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
 
         if self.client_id and self.client_secret and self.workspace_host:
             username = value_of(self.client_id)
-        else:
+        elif self.user:
             username = value_of(self.user)
 
         host: str = value_of(self.host)
@@ -770,11 +823,48 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
             workspace_host=value_of(self.workspace_host),
             pat=value_of(self.password),
         )
-        token: str = provider.create_token()
 
-        return (
-            f"postgresql://{username}:{token}@{host}:{port}/{database}?sslmode=require"
-        )
+        token: str = provider.lakebase_password_provider(self.instance_name)
+
+        # Build connection parameters dictionary
+        params: dict[str, Any] = {
+            "dbname": database,
+            "host": host,
+            "port": port,
+            "password": token,
+            "sslmode": "require",
+        }
+
+        # Only include user if explicitly configured
+        if username:
+            params["user"] = username
+            logger.debug(
+                f"Connection params: dbname={database} user={username} host={host} port={port} password=******** sslmode=require"
+            )
+        else:
+            logger.debug(
+                f"Connection params: dbname={database} host={host} port={port} password=******** sslmode=require (using token identity)"
+            )
+
+        return params
+
+    @property
+    def connection_url(self) -> str:
+        """
+        Get database connection URL as a string (for backwards compatibility).
+
+        Note: It's recommended to use connection_params instead for better flexibility.
+        """
+        params = self.connection_params
+        parts = [f"{k}={v}" for k, v in params.items()]
+        return " ".join(parts)
+
+    def create(self, w: WorkspaceClient | None = None) -> None:
+        from dao_ai.providers.databricks import DatabricksProvider
+
+        provider: DatabricksProvider = DatabricksProvider()
+        provider.create_lakebase(self)
+        provider.create_lakebase_instance_role(self)
 
 
 class SearchParametersModel(BaseModel):
@@ -867,7 +957,7 @@ class PythonFunctionModel(BaseFunctionModel, HasFullName):
 
 class FactoryFunctionModel(BaseFunctionModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    args: Optional[dict[str, AnyVariable]] = Field(default_factory=dict)
+    args: Optional[dict[str, Any]] = Field(default_factory=dict)
     type: Literal[FunctionType.FACTORY] = FunctionType.FACTORY
 
     @property
@@ -879,6 +969,12 @@ class FactoryFunctionModel(BaseFunctionModel, HasFullName):
 
         return [create_factory_tool(self, **kwargs)]
 
+    @model_validator(mode="after")
+    def update_args(self):
+        for key, value in self.args.items():
+            self.args[key] = value_of(value)
+        return self
+
 
 class TransportType(str, Enum):
     STREAMABLE_HTTP = "streamable_http"
@@ -1078,6 +1174,7 @@ class AgentModel(BaseModel):
 class SupervisorModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     model: LLMModel
+    tools: list[ToolModel] = Field(default_factory=list)
     prompt: Optional[str] = None
 
 
@@ -1262,12 +1359,12 @@ class AppModel(BaseModel):
             if len(self.agents) > 1:
                 default_agent: AgentModel = self.agents[0]
                 self.orchestration = OrchestrationModel(
-                    swarm=SupervisorModel(model=default_agent.model)
+                    supervisor=SupervisorModel(model=default_agent.model)
                 )
             elif len(self.agents) == 1:
                 default_agent: AgentModel = self.agents[0]
                 self.orchestration = OrchestrationModel(
-                    supervisor=SwarmModel(
+                    swarm=SwarmModel(
                         model=default_agent.model, default_agent=default_agent
                     )
                 )

dao_ai/graph.py

@@ -27,6 +27,7 @@ from dao_ai.nodes import (
 )
 from dao_ai.prompts import make_prompt
 from dao_ai.state import Context, IncomingState, OutgoingState, SharedState
+from dao_ai.tools import create_tools
 
 
 def route_message(state: SharedState) -> str:
@@ -91,6 +92,8 @@ def _create_supervisor_graph(config: AppConfig) -> CompiledStateGraph:
     orchestration: OrchestrationModel = config.app.orchestration
     supervisor: SupervisorModel = orchestration.supervisor
 
+    tools += create_tools(orchestration.supervisor.tools)
+
     store: BaseStore = None
     if orchestration.memory and orchestration.memory.store:
         store = orchestration.memory.store.as_store()

dao_ai/memory/core.py

@@ -70,10 +70,14 @@ class StoreManager:
             case StorageType.POSTGRES:
                 from dao_ai.memory.postgres import PostgresStoreManager
 
-                store_manager = cls.store_managers.get(store_model.database.name)
+                store_manager = cls.store_managers.get(
+                    store_model.database.instance_name
+                )
                 if store_manager is None:
                     store_manager = PostgresStoreManager(store_model)
-                    cls.store_managers[store_model.database.name] = store_manager
+                    cls.store_managers[store_model.database.instance_name] = (
+                        store_manager
+                    )
             case _:
                 raise ValueError(f"Unknown store type: {store_model.type}")
 
@@ -102,15 +106,15 @@ class CheckpointManager:
                 from dao_ai.memory.postgres import AsyncPostgresCheckpointerManager
 
                 checkpointer_manager = cls.checkpoint_managers.get(
-                    checkpointer_model.database.name
+                    checkpointer_model.database.instance_name
                 )
                 if checkpointer_manager is None:
                     checkpointer_manager = AsyncPostgresCheckpointerManager(
                         checkpointer_model
                     )
-                    cls.checkpoint_managers[checkpointer_model.database.name] = (
-                        checkpointer_manager
-                    )
+                    cls.checkpoint_managers[
+                        checkpointer_model.database.instance_name
+                    ] = checkpointer_manager
             case _:
                 raise ValueError(f"Unknown store type: {checkpointer_model.type}")
 

dao_ai/memory/postgres.py

@@ -20,6 +20,59 @@ from dao_ai.memory.base import (
 )
 
 
+def _create_pool(
+    connection_params: dict[str, Any],
+    database_name: str,
+    max_pool_size: int,
+    timeout_seconds: int,
+    kwargs: dict,
+) -> ConnectionPool:
+    """Create a connection pool using the provided connection parameters."""
+    logger.debug(
+        f"Connection params for {database_name}: {', '.join(k + '=' + (str(v) if k != 'password' else '***') for k, v in connection_params.items())}"
+    )
+
+    # Merge connection_params into kwargs for psycopg
+    connection_kwargs = kwargs | connection_params
+    pool = ConnectionPool(
+        conninfo="",  # Empty conninfo, params come from kwargs
+        min_size=1,
+        max_size=max_pool_size,
+        open=False,
+        timeout=timeout_seconds,
+        kwargs=connection_kwargs,
+    )
+    pool.open(wait=True, timeout=timeout_seconds)
+    logger.info(f"Successfully connected to {database_name}")
+    return pool
+
+
+async def _create_async_pool(
+    connection_params: dict[str, Any],
+    database_name: str,
+    max_pool_size: int,
+    timeout_seconds: int,
+    kwargs: dict,
+) -> AsyncConnectionPool:
+    """Create an async connection pool using the provided connection parameters."""
+    logger.debug(
+        f"Connection params for {database_name}: {', '.join(k + '=' + (str(v) if k != 'password' else '***') for k, v in connection_params.items())}"
+    )
+
+    # Merge connection_params into kwargs for psycopg
+    connection_kwargs = kwargs | connection_params
+    pool = AsyncConnectionPool(
+        conninfo="",  # Empty conninfo, params come from kwargs
+        max_size=max_pool_size,
+        open=False,
+        timeout=timeout_seconds,
+        kwargs=connection_kwargs,
+    )
+    await pool.open(wait=True, timeout=timeout_seconds)
+    logger.info(f"Successfully connected to {database_name}")
+    return pool
+
+
 class AsyncPostgresPoolManager:
     _pools: dict[str, AsyncConnectionPool] = {}
     _lock: asyncio.Lock = asyncio.Lock()
@@ -27,7 +80,7 @@ class AsyncPostgresPoolManager:
     @classmethod
     async def get_pool(cls, database: DatabaseModel) -> AsyncConnectionPool:
         connection_key: str = database.name
-        connection_url: str = database.connection_url
+        connection_params: dict[str, Any] = database.connection_params
 
         async with cls._lock:
             if connection_key in cls._pools:
@@ -41,23 +94,17 @@ class AsyncPostgresPoolManager:
                 "autocommit": True,
             } | database.connection_kwargs or {}
 
-            pool: AsyncConnectionPool = AsyncConnectionPool(
-                conninfo=connection_url,
-                max_size=database.max_pool_size,
-                open=False,
-                timeout=database.timeout_seconds,
+            # Create connection pool
+            pool: AsyncConnectionPool = await _create_async_pool(
+                connection_params=connection_params,
+                database_name=database.name,
+                max_pool_size=database.max_pool_size,
+                timeout_seconds=database.timeout_seconds,
                 kwargs=kwargs,
             )
 
-            try:
-                await pool.open(wait=True, timeout=database.timeout_seconds)
-                cls._pools[connection_key] = pool
-                return pool
-            except Exception as e:
-                logger.error(
-                    f"Failed to create PostgreSQL pool for {database.name}: {e}"
-                )
-                raise e
+            cls._pools[connection_key] = pool
+            return pool
 
     @classmethod
     async def close_pool(cls, database: DatabaseModel):
@@ -74,8 +121,17 @@ class AsyncPostgresPoolManager:
         async with cls._lock:
             for connection_key, pool in cls._pools.items():
                 try:
-                    await pool.close()
+                    # Use a short timeout to avoid blocking on pool closure
+                    await asyncio.wait_for(pool.close(), timeout=2.0)
                     logger.debug(f"Closed PostgreSQL pool: {connection_key}")
+                except asyncio.TimeoutError:
+                    logger.warning(
+                        f"Timeout closing pool {connection_key}, forcing closure"
+                    )
+                except asyncio.CancelledError:
+                    logger.warning(
+                        f"Pool closure cancelled for {connection_key} (shutdown in progress)"
+                    )
                 except Exception as e:
                     logger.error(f"Error closing pool {connection_key}: {e}")
             cls._pools.clear()
@@ -209,7 +265,7 @@ class PostgresPoolManager:
     @classmethod
     def get_pool(cls, database: DatabaseModel) -> ConnectionPool:
         connection_key: str = str(database.name)
-        connection_url: str = database.connection_url
+        connection_params: dict[str, Any] = database.connection_params
 
         with cls._lock:
             if connection_key in cls._pools:
@@ -223,23 +279,17 @@ class PostgresPoolManager:
                 "autocommit": True,
             } | database.connection_kwargs or {}
 
-            pool: ConnectionPool = ConnectionPool(
-                conninfo=connection_url,
-                max_size=database.max_pool_size,
-                open=False,
-                timeout=database.timeout_seconds,
+            # Create connection pool
+            pool: ConnectionPool = _create_pool(
+                connection_params=connection_params,
+                database_name=database.name,
+                max_pool_size=database.max_pool_size,
+                timeout_seconds=database.timeout_seconds,
                 kwargs=kwargs,
             )
 
-            try:
-                pool.open(wait=True, timeout=database.timeout_seconds)
-                cls._pools[connection_key] = pool
-                return pool
-            except Exception as e:
-                logger.error(
-                    f"Failed to create PostgreSQL pool for {database.name}: {e}"
-                )
-                raise e
+            cls._pools[connection_key] = pool
+            return pool
 
     @classmethod
     def close_pool(cls, database: DatabaseModel):
@@ -369,8 +419,27 @@ def _shutdown_pools():
 
 def _shutdown_async_pools():
     try:
-        asyncio.run(AsyncPostgresPoolManager.close_all_pools())
-        logger.debug("Successfully closed all asynchronous PostgreSQL pools")
+        # Try to get the current event loop first
+        try:
+            loop = asyncio.get_running_loop()
+            # If we're already in an event loop, create a task
+            loop.create_task(AsyncPostgresPoolManager.close_all_pools())
+            logger.debug("Scheduled async pool closure in running event loop")
+        except RuntimeError:
+            # No running loop, try to get or create one
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_closed():
+                    # Loop is closed, create a new one
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                loop.run_until_complete(AsyncPostgresPoolManager.close_all_pools())
+                logger.debug("Successfully closed all asynchronous PostgreSQL pools")
+            except Exception as inner_e:
+                # If all else fails, just log the error
+                logger.warning(
+                    f"Could not close async pools cleanly during shutdown: {inner_e}"
+                )
     except Exception as e:
         logger.error(
             f"Error closing asynchronous PostgreSQL pools during shutdown: {e}"

dao_ai/models.py

@@ -5,6 +5,7 @@ from typing import Any, Generator, Optional, Sequence, Union
 
 from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
 from langgraph.graph.state import CompiledStateGraph
+from langgraph.types import StateSnapshot
 from loguru import logger
 from mlflow import MlflowClient
 from mlflow.pyfunc import ChatAgent, ChatModel, ResponsesAgent
@@ -59,6 +60,113 @@ def get_latest_model_version(model_name: str) -> int:
     return latest_version
 
 
+async def get_state_snapshot_async(
+    graph: CompiledStateGraph, thread_id: str
+) -> Optional[StateSnapshot]:
+    """
+    Retrieve the state snapshot from the graph for a given thread_id asynchronously.
+
+    This utility function accesses the graph's checkpointer to retrieve the current
+    state snapshot, which contains the full state values and metadata.
+
+    Args:
+        graph: The compiled LangGraph state machine
+        thread_id: The thread/conversation ID to retrieve state for
+
+    Returns:
+        StateSnapshot if found, None otherwise
+    """
+    logger.debug(f"Retrieving state snapshot for thread_id: {thread_id}")
+    try:
+        # Check if graph has a checkpointer
+        if graph.checkpointer is None:
+            logger.debug("No checkpointer available in graph")
+            return None
+
+        # Get the current state from the checkpointer (use async version)
+        config: dict[str, Any] = {"configurable": {"thread_id": thread_id}}
+        state_snapshot: Optional[StateSnapshot] = await graph.aget_state(config)
+
+        if state_snapshot is None:
+            logger.debug(f"No state found for thread_id: {thread_id}")
+            return None
+
+        return state_snapshot
+
+    except Exception as e:
+        logger.warning(f"Error retrieving state snapshot for thread {thread_id}: {e}")
+        return None
+
+
+def get_state_snapshot(
+    graph: CompiledStateGraph, thread_id: str
+) -> Optional[StateSnapshot]:
+    """
+    Retrieve the state snapshot from the graph for a given thread_id.
+
+    This is a synchronous wrapper around get_state_snapshot_async.
+    Use this for backward compatibility in synchronous contexts.
+
+    Args:
+        graph: The compiled LangGraph state machine
+        thread_id: The thread/conversation ID to retrieve state for
+
+    Returns:
+        StateSnapshot if found, None otherwise
+    """
+    import asyncio
+
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+    try:
+        return loop.run_until_complete(get_state_snapshot_async(graph, thread_id))
+    except Exception as e:
+        logger.warning(f"Error in synchronous state snapshot retrieval: {e}")
+        return None
+
+
+def get_genie_conversation_ids_from_state(
+    state_snapshot: Optional[StateSnapshot],
+) -> dict[str, str]:
+    """
+    Extract genie_conversation_ids from a state snapshot.
+
+    This function extracts the genie_conversation_ids dictionary from the state
+    snapshot values if present.
+
+    Args:
+        state_snapshot: The state snapshot to extract conversation IDs from
+
+    Returns:
+        A dictionary mapping genie space_id to conversation_id, or empty dict if not found
+    """
+    if state_snapshot is None:
+        return {}
+
+    try:
+        # Extract state values - these contain the actual state data
+        state_values: dict[str, Any] = state_snapshot.values
+
+        # Extract genie_conversation_ids from state values
+        genie_conversation_ids: dict[str, str] = state_values.get(
+            "genie_conversation_ids", {}
+        )
+
+        if genie_conversation_ids:
+            logger.debug(f"Retrieved genie_conversation_ids: {genie_conversation_ids}")
+            return genie_conversation_ids
+
+        return {}
+
+    except Exception as e:
+        logger.warning(f"Error extracting genie_conversation_ids from state: {e}")
+        return {}
+
+
 class LanggraphChatModel(ChatModel):
     """
     ChatModel that delegates requests to a LangGraph CompiledStateGraph.
@@ -257,7 +365,19 @@ class LanggraphResponsesAgent(ResponsesAgent):
             text=last_message.content, id=f"msg_{uuid.uuid4().hex[:8]}"
         )
 
-        custom_outputs = custom_inputs
+        # Retrieve genie_conversation_ids from state if available
+        custom_outputs: dict[str, Any] = custom_inputs.copy()
+        thread_id: Optional[str] = context.thread_id
+        if thread_id:
+            state_snapshot: Optional[StateSnapshot] = loop.run_until_complete(
+                get_state_snapshot_async(self.graph, thread_id)
+            )
+            genie_conversation_ids: dict[str, str] = (
+                get_genie_conversation_ids_from_state(state_snapshot)
+            )
+            if genie_conversation_ids:
+                custom_outputs["genie_conversation_ids"] = genie_conversation_ids
+
         return ResponsesAgentResponse(
             output=[output_item], custom_outputs=custom_outputs
         )
@@ -318,7 +438,22 @@ class LanggraphResponsesAgent(ResponsesAgent):
                                 **self.create_text_delta(delta=content, item_id=item_id)
                             )
 
-                custom_outputs = custom_inputs
+                # Retrieve genie_conversation_ids from state if available
+                custom_outputs: dict[str, Any] = custom_inputs.copy()
+                thread_id: Optional[str] = context.thread_id
+
+                if thread_id:
+                    state_snapshot: Optional[
+                        StateSnapshot
+                    ] = await get_state_snapshot_async(self.graph, thread_id)
+                    genie_conversation_ids: dict[str, str] = (
+                        get_genie_conversation_ids_from_state(state_snapshot)
+                    )
+                    if genie_conversation_ids:
+                        custom_outputs["genie_conversation_ids"] = (
+                            genie_conversation_ids
+                        )
+
                 # Yield final output item
                 yield ResponsesAgentStreamEvent(
                     type="response.output_item.done",