dao-ai 0.0.25__py3-none-any.whl → 0.1.2__py3-none-any.whl

dao_ai/config.py

@@ -12,6 +12,7 @@ from typing import (
     Iterator,
     Literal,
     Optional,
+    Self,
     Sequence,
     TypeAlias,
     Union,
@@ -22,22 +23,33 @@ from databricks.sdk.credentials_provider import (
     CredentialsStrategy,
     ModelServingUserCredentials,
 )
+from databricks.sdk.errors.platform import NotFound
 from databricks.sdk.service.catalog import FunctionInfo, TableInfo
+from databricks.sdk.service.dashboards import GenieSpace
 from databricks.sdk.service.database import DatabaseInstance
+from databricks.sdk.service.sql import GetWarehouseResponse
 from databricks.vector_search.client import VectorSearchClient
 from databricks.vector_search.index import VectorSearchIndex
 from databricks_langchain import (
+    ChatDatabricks,
+    DatabricksEmbeddings,
     DatabricksFunctionClient,
 )
+from langchain.agents.structured_output import ProviderStrategy, ToolStrategy
+from langchain_core.embeddings import Embeddings
 from langchain_core.language_models import LanguageModelLike
+from langchain_core.messages import BaseMessage, messages_from_dict
 from langchain_core.runnables.base import RunnableLike
 from langchain_openai import ChatOpenAI
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.store.base import BaseStore
 from loguru import logger
+from mlflow.genai.datasets import EvaluationDataset, create_dataset, get_dataset
+from mlflow.genai.prompts import PromptVersion, load_prompt
 from mlflow.models import ModelConfig
 from mlflow.models.resources import (
+    DatabricksApp,
     DatabricksFunction,
     DatabricksGenieSpace,
     DatabricksLakebase,
@@ -49,14 +61,20 @@ from mlflow.models.resources import (
     DatabricksVectorSearchIndex,
 )
 from mlflow.pyfunc import ChatModel, ResponsesAgent
+from mlflow.types.responses import (
+    ResponsesAgentRequest,
+)
 from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
+    PrivateAttr,
     field_serializer,
     model_validator,
 )
 
+from dao_ai.utils import normalize_name
+
 
 class HasValue(ABC):
     @abstractmethod
@@ -75,27 +93,6 @@ class HasFullName(ABC):
     def full_name(self) -> str: ...
 
 
-class IsDatabricksResource(ABC):
-    on_behalf_of_user: Optional[bool] = False
-
-    @abstractmethod
-    def as_resources(self) -> Sequence[DatabricksResource]: ...
-
-    @property
-    @abstractmethod
-    def api_scopes(self) -> Sequence[str]: ...
-
-    @property
-    def workspace_client(self) -> WorkspaceClient:
-        credentials_strategy: CredentialsStrategy = None
-        if self.on_behalf_of_user:
-            credentials_strategy = ModelServingUserCredentials()
-        logger.debug(
-            f"Creating WorkspaceClient with credentials strategy: {credentials_strategy}"
-        )
-        return WorkspaceClient(credentials_strategy=credentials_strategy)
-
-
 class EnvironmentVariableModel(BaseModel, HasValue):
     model_config = ConfigDict(
         frozen=True,
@@ -194,6 +191,162 @@ AnyVariable: TypeAlias = (
 )
 
 
+class ServicePrincipalModel(BaseModel):
+    model_config = ConfigDict(
+        frozen=True,
+        use_enum_values=True,
+    )
+    client_id: AnyVariable
+    client_secret: AnyVariable
+
+
+class IsDatabricksResource(ABC, BaseModel):
+    """
+    Base class for Databricks resources with authentication support.
+
+    Authentication Options:
+    ----------------------
+    1. **On-Behalf-Of User (OBO)**: Set on_behalf_of_user=True to use the
+       calling user's identity via ModelServingUserCredentials.
+
+    2. **Service Principal (OAuth M2M)**: Provide service_principal or
+       (client_id + client_secret + workspace_host) for service principal auth.
+
+    3. **Personal Access Token (PAT)**: Provide pat (and optionally workspace_host)
+       to authenticate with a personal access token.
+
+    4. **Ambient Authentication**: If no credentials provided, uses SDK defaults
+       (environment variables, notebook context, etc.)
+
+    Authentication Priority:
+    1. OBO (on_behalf_of_user=True)
+    2. Service Principal (client_id + client_secret + workspace_host)
+    3. PAT (pat + workspace_host)
+    4. Ambient/default authentication
+    """
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    on_behalf_of_user: Optional[bool] = False
+    service_principal: Optional[ServicePrincipalModel] = None
+    client_id: Optional[AnyVariable] = None
+    client_secret: Optional[AnyVariable] = None
+    workspace_host: Optional[AnyVariable] = None
+    pat: Optional[AnyVariable] = None
+
+    # Private attribute to cache the workspace client (lazy instantiation)
+    _workspace_client: Optional[WorkspaceClient] = PrivateAttr(default=None)
+
+    @abstractmethod
+    def as_resources(self) -> Sequence[DatabricksResource]: ...
+
+    @property
+    @abstractmethod
+    def api_scopes(self) -> Sequence[str]: ...
+
+    @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
+
+    @model_validator(mode="after")
+    def _validate_auth_not_mixed(self) -> Self:
+        """Validate that OAuth and PAT authentication are not both provided."""
+        has_oauth: bool = self.client_id is not None and self.client_secret is not None
+        has_pat: bool = self.pat is not None
+
+        if has_oauth and has_pat:
+            raise ValueError(
+                "Cannot use both OAuth and user authentication methods. "
+                "Please provide either OAuth credentials or user credentials."
+            )
+        return self
+
+    @property
+    def workspace_client(self) -> WorkspaceClient:
+        """
+        Get a WorkspaceClient configured with the appropriate authentication.
+
+        The client is lazily instantiated on first access and cached for subsequent calls.
+
+        Authentication priority:
+        1. If on_behalf_of_user is True, uses ModelServingUserCredentials (OBO)
+        2. If service principal credentials are configured (client_id, client_secret,
+           workspace_host), uses OAuth M2M
+        3. If PAT is configured, uses token authentication
+        4. Otherwise, uses default/ambient authentication
+        """
+        # Return cached client if already instantiated
+        if self._workspace_client is not None:
+            return self._workspace_client
+
+        from dao_ai.utils import normalize_host
+
+        # Check for OBO first (highest priority)
+        if self.on_behalf_of_user:
+            credentials_strategy: CredentialsStrategy = ModelServingUserCredentials()
+            logger.debug(
+                f"Creating WorkspaceClient for {self.__class__.__name__} "
+                f"with OBO credentials strategy"
+            )
+            self._workspace_client = WorkspaceClient(
+                credentials_strategy=credentials_strategy
+            )
+            return self._workspace_client
+
+        # Check for service principal credentials
+        client_id_value: str | None = (
+            value_of(self.client_id) if self.client_id else None
+        )
+        client_secret_value: str | None = (
+            value_of(self.client_secret) if self.client_secret else None
+        )
+        workspace_host_value: str | None = (
+            normalize_host(value_of(self.workspace_host))
+            if self.workspace_host
+            else None
+        )
+
+        if client_id_value and client_secret_value and workspace_host_value:
+            logger.debug(
+                f"Creating WorkspaceClient for {self.__class__.__name__} with service principal: "
+                f"client_id={client_id_value}, host={workspace_host_value}"
+            )
+            self._workspace_client = WorkspaceClient(
+                host=workspace_host_value,
+                client_id=client_id_value,
+                client_secret=client_secret_value,
+                auth_type="oauth-m2m",
+            )
+            return self._workspace_client
+
+        # Check for PAT authentication
+        pat_value: str | None = value_of(self.pat) if self.pat else None
+        if pat_value:
+            logger.debug(
+                f"Creating WorkspaceClient for {self.__class__.__name__} with PAT"
+            )
+            self._workspace_client = WorkspaceClient(
+                host=workspace_host_value,
+                token=pat_value,
+                auth_type="pat",
+            )
+            return self._workspace_client
+
+        # Default: use ambient authentication
+        logger.debug(
+            f"Creating WorkspaceClient for {self.__class__.__name__} "
+            "with default/ambient authentication"
+        )
+        self._workspace_client = WorkspaceClient()
+        return self._workspace_client
+
+
 class Privilege(str, Enum):
     ALL_PRIVILEGES = "ALL_PRIVILEGES"
     USE_CATALOG = "USE_CATALOG"
@@ -220,9 +373,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")
@@ -242,7 +407,26 @@ class SchemaModel(BaseModel, HasFullName):
         provider.create_schema(self)
 
 
-class TableModel(BaseModel, HasFullName, IsDatabricksResource):
+class DatabricksAppModel(IsDatabricksResource, HasFullName):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: str
+    url: str
+
+    @property
+    def full_name(self) -> str:
+        return self.name
+
+    @property
+    def api_scopes(self) -> Sequence[str]:
+        return ["apps.apps"]
+
+    def as_resources(self) -> Sequence[DatabricksResource]:
+        return [
+            DatabricksApp(app_name=self.name, on_behalf_of_user=self.on_behalf_of_user)
+        ]
+
+
+class TableModel(IsDatabricksResource, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
     name: Optional[str] = None
@@ -268,6 +452,22 @@ class TableModel(BaseModel, HasFullName, IsDatabricksResource):
     def api_scopes(self) -> Sequence[str]:
         return []
 
+    def exists(self) -> bool:
+        """Check if the table exists in Unity Catalog.
+
+        Returns:
+            True if the table exists, False otherwise.
+        """
+        try:
+            self.workspace_client.tables.get(full_name=self.full_name)
+            return True
+        except NotFound:
+            logger.debug(f"Table not found: {self.full_name}")
+            return False
+        except Exception as e:
+            logger.warning(f"Error checking table existence for {self.full_name}: {e}")
+            return False
+
     def as_resources(self) -> Sequence[DatabricksResource]:
         resources: list[DatabricksResource] = []
 
@@ -311,12 +511,17 @@ class TableModel(BaseModel, HasFullName, IsDatabricksResource):
         return resources
 
 
-class LLMModel(BaseModel, IsDatabricksResource):
+class LLMModel(IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
+    description: Optional[str] = None
     temperature: Optional[float] = 0.1
     max_tokens: Optional[int] = 8192
     fallbacks: Optional[list[Union[str, "LLMModel"]]] = Field(default_factory=list)
+    use_responses_api: Optional[bool] = Field(
+        default=False,
+        description="Use Responses API for ResponsesAgent endpoints",
+    )
 
     @property
     def api_scopes(self) -> Sequence[str]:
@@ -324,6 +529,10 @@ class LLMModel(BaseModel, IsDatabricksResource):
             "serving.serving-endpoints",
         ]
 
+    @property
+    def uri(self) -> str:
+        return f"databricks:/{self.name}"
+
     def as_resources(self) -> Sequence[DatabricksResource]:
         return [
             DatabricksServingEndpoint(
@@ -332,19 +541,12 @@ class LLMModel(BaseModel, IsDatabricksResource):
         ]
 
     def as_chat_model(self) -> LanguageModelLike:
-        # Retrieve langchain chat client from workspace client to enable OBO
-        # ChatOpenAI does not allow additional inputs at the moment, so we cannot use it directly
-        # chat_client: LanguageModelLike = self.as_open_ai_client()
-
-        # Create ChatDatabricksWrapper instance directly
-        from dao_ai.chat_models import ChatDatabricksFiltered
-
-        chat_client: LanguageModelLike = ChatDatabricksFiltered(
-            model=self.name, temperature=self.temperature, max_tokens=self.max_tokens
+        chat_client: LanguageModelLike = ChatDatabricks(
+            model=self.name,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+            use_responses_api=self.use_responses_api,
         )
-        # chat_client: LanguageModelLike = ChatDatabricks(
-        #     model=self.name, temperature=self.temperature, max_tokens=self.max_tokens
-        # )
 
         fallbacks: Sequence[LanguageModelLike] = []
         for fallback in self.fallbacks:
@@ -376,6 +578,9 @@ class LLMModel(BaseModel, IsDatabricksResource):
 
         return chat_client
 
+    def as_embeddings_model(self) -> Embeddings:
+        return DatabricksEmbeddings(endpoint=self.name)
+
 
 class VectorSearchEndpointType(str, Enum):
     STANDARD = "STANDARD"
@@ -387,8 +592,15 @@ class VectorSearchEndpoint(BaseModel):
     name: str
     type: VectorSearchEndpointType = VectorSearchEndpointType.STANDARD
 
+    @field_serializer("type")
+    def serialize_type(self, value: VectorSearchEndpointType) -> str:
+        """Ensure enum is serialized to string value."""
+        if isinstance(value, VectorSearchEndpointType):
+            return value.value
+        return str(value)
+
 
-class IndexModel(BaseModel, HasFullName, IsDatabricksResource):
+class IndexModel(IsDatabricksResource, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
     name: str
@@ -413,12 +625,297 @@ class IndexModel(BaseModel, HasFullName, IsDatabricksResource):
         ]
 
 
-class GenieRoomModel(BaseModel, IsDatabricksResource):
+class FunctionModel(IsDatabricksResource, HasFullName):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
+    name: Optional[str] = None
+
+    @model_validator(mode="after")
+    def validate_name_or_schema_required(self) -> Self:
+        if not self.name and not self.schema_model:
+            raise ValueError(
+                "Either 'name' or 'schema_model' must be provided for FunctionModel"
+            )
+        return self
+
+    @property
+    def full_name(self) -> str:
+        if self.schema_model:
+            name: str = ""
+            if self.name:
+                name = f".{self.name}"
+            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}{name}"
+        return self.name
+
+    def exists(self) -> bool:
+        """Check if the function exists in Unity Catalog.
+
+        Returns:
+            True if the function exists, False otherwise.
+        """
+        try:
+            self.workspace_client.functions.get(name=self.full_name)
+            return True
+        except NotFound:
+            logger.debug(f"Function not found: {self.full_name}")
+            return False
+        except Exception as e:
+            logger.warning(
+                f"Error checking function existence for {self.full_name}: {e}"
+            )
+            return False
+
+    def as_resources(self) -> Sequence[DatabricksResource]:
+        resources: list[DatabricksResource] = []
+        if self.name:
+            resources.append(
+                DatabricksFunction(
+                    function_name=self.full_name,
+                    on_behalf_of_user=self.on_behalf_of_user,
+                )
+            )
+        else:
+            w: WorkspaceClient = self.workspace_client
+            schema_full_name: str = self.schema_model.full_name
+            functions: Iterator[FunctionInfo] = w.functions.list(
+                catalog_name=self.schema_model.catalog_name,
+                schema_name=self.schema_model.schema_name,
+            )
+            resources.extend(
+                [
+                    DatabricksFunction(
+                        function_name=f"{schema_full_name}.{function.name}",
+                        on_behalf_of_user=self.on_behalf_of_user,
+                    )
+                    for function in functions
+                ]
+            )
+
+        return resources
+
+    @property
+    def api_scopes(self) -> Sequence[str]:
+        return ["sql.statement-execution"]
+
+
+class WarehouseModel(IsDatabricksResource):
+    model_config = ConfigDict()
+    name: str
+    description: Optional[str] = None
+    warehouse_id: AnyVariable
+
+    @property
+    def api_scopes(self) -> Sequence[str]:
+        return [
+            "sql.warehouses",
+            "sql.statement-execution",
+        ]
+
+    def as_resources(self) -> Sequence[DatabricksResource]:
+        return [
+            DatabricksSQLWarehouse(
+                warehouse_id=value_of(self.warehouse_id),
+                on_behalf_of_user=self.on_behalf_of_user,
+            )
+        ]
+
+    @model_validator(mode="after")
+    def update_warehouse_id(self) -> Self:
+        self.warehouse_id = value_of(self.warehouse_id)
+        return self
+
+
+class GenieRoomModel(IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     description: Optional[str] = None
     space_id: AnyVariable
 
+    _space_details: Optional[GenieSpace] = PrivateAttr(default=None)
+
+    def _get_space_details(self) -> GenieSpace:
+        if self._space_details is None:
+            self._space_details = self.workspace_client.genie.get_space(
+                space_id=self.space_id, include_serialized_space=True
+            )
+        return self._space_details
+
+    def _parse_serialized_space(self) -> dict[str, Any]:
+        """Parse the serialized_space JSON string and return the parsed data."""
+        import json
+
+        space_details = self._get_space_details()
+        if not space_details.serialized_space:
+            return {}
+
+        try:
+            return json.loads(space_details.serialized_space)
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse serialized_space: {e}")
+            return {}
+
+    @property
+    def warehouse(self) -> Optional[WarehouseModel]:
+        """Extract warehouse information from the Genie space.
+
+        Returns:
+            WarehouseModel instance if warehouse_id is available, None otherwise.
+        """
+        space_details: GenieSpace = self._get_space_details()
+
+        if not space_details.warehouse_id:
+            return None
+
+        try:
+            response: GetWarehouseResponse = self.workspace_client.warehouses.get(
+                space_details.warehouse_id
+            )
+            warehouse_name: str = response.name or space_details.warehouse_id
+
+            warehouse_model = WarehouseModel(
+                name=warehouse_name,
+                warehouse_id=space_details.warehouse_id,
+                on_behalf_of_user=self.on_behalf_of_user,
+                service_principal=self.service_principal,
+                client_id=self.client_id,
+                client_secret=self.client_secret,
+                workspace_host=self.workspace_host,
+                pat=self.pat,
+            )
+
+            # Share the cached workspace client if available
+            if self._workspace_client is not None:
+                warehouse_model._workspace_client = self._workspace_client
+
+            return warehouse_model
+        except Exception as e:
+            logger.warning(
+                f"Failed to fetch warehouse details for {space_details.warehouse_id}: {e}"
+            )
+            return None
+
+    @property
+    def tables(self) -> list[TableModel]:
+        """Extract tables from the serialized Genie space.
+
+        Databricks Genie stores tables in: data_sources.tables[].identifier
+        Only includes tables that actually exist in Unity Catalog.
+        """
+        parsed_space = self._parse_serialized_space()
+        tables_list: list[TableModel] = []
+
+        # Primary structure: data_sources.tables with 'identifier' field
+        if "data_sources" in parsed_space:
+            data_sources = parsed_space["data_sources"]
+            if isinstance(data_sources, dict) and "tables" in data_sources:
+                tables_data = data_sources["tables"]
+                if isinstance(tables_data, list):
+                    for table_item in tables_data:
+                        table_name: str | None = None
+                        if isinstance(table_item, dict):
+                            # Standard Databricks structure uses 'identifier'
+                            table_name = table_item.get("identifier") or table_item.get(
+                                "name"
+                            )
+                        elif isinstance(table_item, str):
+                            table_name = table_item
+
+                        if table_name:
+                            table_model = TableModel(
+                                name=table_name,
+                                on_behalf_of_user=self.on_behalf_of_user,
+                                service_principal=self.service_principal,
+                                client_id=self.client_id,
+                                client_secret=self.client_secret,
+                                workspace_host=self.workspace_host,
+                                pat=self.pat,
+                            )
+                            # Share the cached workspace client if available
+                            if self._workspace_client is not None:
+                                table_model._workspace_client = self._workspace_client
+
+                            # Verify the table exists before adding
+                            if not table_model.exists():
+                                continue
+
+                            tables_list.append(table_model)
+
+        return tables_list
+
+    @property
+    def functions(self) -> list[FunctionModel]:
+        """Extract functions from the serialized Genie space.
+
+        Databricks Genie stores functions in multiple locations:
+        - instructions.sql_functions[].identifier (SQL functions)
+        - data_sources.functions[].identifier (other functions)
+        Only includes functions that actually exist in Unity Catalog.
+        """
+        parsed_space = self._parse_serialized_space()
+        functions_list: list[FunctionModel] = []
+        seen_functions: set[str] = set()
+
+        def add_function_if_exists(function_name: str) -> None:
+            """Helper to add a function if it exists and hasn't been added."""
+            if function_name in seen_functions:
+                return
+
+            seen_functions.add(function_name)
+            function_model = FunctionModel(
+                name=function_name,
+                on_behalf_of_user=self.on_behalf_of_user,
+                service_principal=self.service_principal,
+                client_id=self.client_id,
+                client_secret=self.client_secret,
+                workspace_host=self.workspace_host,
+                pat=self.pat,
+            )
+            # Share the cached workspace client if available
+            if self._workspace_client is not None:
+                function_model._workspace_client = self._workspace_client
+
+            # Verify the function exists before adding
+            if not function_model.exists():
+                return
+
+            functions_list.append(function_model)
+
+        # Primary structure: instructions.sql_functions with 'identifier' field
+        if "instructions" in parsed_space:
+            instructions = parsed_space["instructions"]
+            if isinstance(instructions, dict) and "sql_functions" in instructions:
+                sql_functions_data = instructions["sql_functions"]
+                if isinstance(sql_functions_data, list):
+                    for function_item in sql_functions_data:
+                        if isinstance(function_item, dict):
+                            # SQL functions use 'identifier' field
+                            function_name = function_item.get(
+                                "identifier"
+                            ) or function_item.get("name")
+                            if function_name:
+                                add_function_if_exists(function_name)
+
+        # Secondary structure: data_sources.functions with 'identifier' field
+        if "data_sources" in parsed_space:
+            data_sources = parsed_space["data_sources"]
+            if isinstance(data_sources, dict) and "functions" in data_sources:
+                functions_data = data_sources["functions"]
+                if isinstance(functions_data, list):
+                    for function_item in functions_data:
+                        function_name: str | None = None
+                        if isinstance(function_item, dict):
+                            # Standard Databricks structure uses 'identifier'
+                            function_name = function_item.get(
+                                "identifier"
+                            ) or function_item.get("name")
+                        elif isinstance(function_item, str):
+                            function_name = function_item
+
+                        if function_name:
+                            add_function_if_exists(function_name)
+
+        return functions_list
+
     @property
     def api_scopes(self) -> Sequence[str]:
         return [
@@ -434,12 +931,24 @@ 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
 
+    @model_validator(mode="after")
+    def update_description_from_space(self) -> Self:
+        """Populate description from GenieSpace if not provided."""
+        if not self.description:
+            try:
+                space_details = self._get_space_details()
+                if space_details.description:
+                    self.description = space_details.description
+            except Exception as e:
+                logger.debug(f"Could not fetch description from Genie space: {e}")
+        return self
 
-class VolumeModel(BaseModel, HasFullName, IsDatabricksResource):
+
+class VolumeModel(IsDatabricksResource, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
     name: str
@@ -499,7 +1008,7 @@ class VolumePathModel(BaseModel, HasFullName):
         provider.create_path(self)
 
 
-class VectorStoreModel(BaseModel, IsDatabricksResource):
+class VectorStoreModel(IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     embedding_model: Optional[LLMModel] = None
     index: Optional[IndexModel] = None
@@ -513,13 +1022,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
 
@@ -540,14 +1049,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,
@@ -598,64 +1107,9 @@ class VectorStoreModel(BaseModel, IsDatabricksResource):
         provider.create_vector_store(self)
 
 
-class FunctionModel(BaseModel, HasFullName, IsDatabricksResource):
+class ConnectionModel(IsDatabricksResource, HasFullName):
     model_config = ConfigDict()
-    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
-    name: Optional[str] = None
-
-    @model_validator(mode="after")
-    def validate_name_or_schema_required(self) -> "FunctionModel":
-        if not self.name and not self.schema_model:
-            raise ValueError(
-                "Either 'name' or 'schema_model' must be provided for FunctionModel"
-            )
-        return self
-
-    @property
-    def full_name(self) -> str:
-        if self.schema_model:
-            name: str = ""
-            if self.name:
-                name = f".{self.name}"
-            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}{name}"
-        return self.name
-
-    def as_resources(self) -> Sequence[DatabricksResource]:
-        resources: list[DatabricksResource] = []
-        if self.name:
-            resources.append(
-                DatabricksFunction(
-                    function_name=self.full_name,
-                    on_behalf_of_user=self.on_behalf_of_user,
-                )
-            )
-        else:
-            w: WorkspaceClient = self.workspace_client
-            schema_full_name: str = self.schema_model.full_name
-            functions: Iterator[FunctionInfo] = w.functions.list(
-                catalog_name=self.schema_model.catalog_name,
-                schema_name=self.schema_model.schema_name,
-            )
-            resources.extend(
-                [
-                    DatabricksFunction(
-                        function_name=f"{schema_full_name}.{function.name}",
-                        on_behalf_of_user=self.on_behalf_of_user,
-                    )
-                    for function in functions
-                ]
-            )
-
-        return resources
-
-    @property
-    def api_scopes(self) -> Sequence[str]:
-        return ["sql.statement-execution"]
-
-
-class ConnectionModel(BaseModel, HasFullName, IsDatabricksResource):
-    model_config = ConfigDict()
-    name: str
+    name: str
 
     @property
     def full_name(self) -> str:
@@ -680,34 +1134,58 @@ class ConnectionModel(BaseModel, HasFullName, IsDatabricksResource):
         ]
 
 
-class WarehouseModel(BaseModel, IsDatabricksResource):
-    model_config = ConfigDict()
-    name: str
-    description: Optional[str] = None
-    warehouse_id: AnyVariable
-
-    @property
-    def api_scopes(self) -> Sequence[str]:
-        return [
-            "sql.warehouses",
-            "sql.statement-execution",
-        ]
-
-    def as_resources(self) -> Sequence[DatabricksResource]:
-        return [
-            DatabricksSQLWarehouse(
-                warehouse_id=value_of(self.warehouse_id),
-                on_behalf_of_user=self.on_behalf_of_user,
-            )
-        ]
-
-    @model_validator(mode="after")
-    def update_warehouse_id(self):
-        self.warehouse_id = value_of(self.warehouse_id)
-        return self
-
+class DatabaseModel(IsDatabricksResource):
+    """
+    Configuration for database connections supporting both Databricks Lakebase and standard PostgreSQL.
+
+    Authentication is inherited from IsDatabricksResource. Additionally supports:
+    - user/password: For user-based database authentication
+
+    Connection Types (determined by fields provided):
+    - Databricks Lakebase: Provide `instance_name` (authentication optional, supports ambient auth)
+    - Standard PostgreSQL: Provide `host` (authentication required via user/password)
+
+    Note: `instance_name` and `host` are mutually exclusive. Provide one or the other.
+
+    Example Databricks Lakebase with Service Principal:
+    ```yaml
+    databases:
+      my_lakebase:
+        name: my-database
+        instance_name: my-lakebase-instance
+        service_principal:
+          client_id:
+            env: SERVICE_PRINCIPAL_CLIENT_ID
+          client_secret:
+            scope: my-scope
+            secret: sp-client-secret
+        workspace_host:
+          env: DATABRICKS_HOST
+    ```
+
+    Example Databricks Lakebase with Ambient Authentication:
+    ```yaml
+    databases:
+      my_lakebase:
+        name: my-database
+        instance_name: my-lakebase-instance
+        on_behalf_of_user: true
+    ```
+
+    Example Standard PostgreSQL:
+    ```yaml
+    databases:
+      my_postgres:
+        name: my-database
+        host: my-postgres-host.example.com
+        port: 5432
+        database: my_db
+        user: my_user
+        password:
+          env: PGPASSWORD
+    ```
+    """
 
-class DatabaseModel(BaseModel, IsDatabricksResource):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     instance_name: Optional[str] = None
@@ -720,80 +1198,137 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
     timeout_seconds: Optional[int] = 10
     capacity: Optional[Literal["CU_1", "CU_2"]] = "CU_2"
     node_count: Optional[int] = None
+    # Database-specific auth (user identity for DB connection)
     user: Optional[AnyVariable] = None
     password: Optional[AnyVariable] = None
-    client_id: Optional[AnyVariable] = None
-    client_secret: Optional[AnyVariable] = None
-    workspace_host: Optional[AnyVariable] = None
 
     @property
     def api_scopes(self) -> Sequence[str]:
-        return []
+        return ["database.database-instances"]
+
+    @property
+    def is_lakebase(self) -> bool:
+        """Returns True if this is a Databricks Lakebase connection (instance_name provided)."""
+        return self.instance_name is not None
 
     def as_resources(self) -> Sequence[DatabricksResource]:
-        return [
-            DatabricksLakebase(
-                database_instance_name=self.instance_name,
-                on_behalf_of_user=self.on_behalf_of_user,
-            )
-        ]
+        if self.is_lakebase:
+            return [
+                DatabricksLakebase(
+                    database_instance_name=self.instance_name,
+                    on_behalf_of_user=self.on_behalf_of_user,
+                )
+            ]
+        return []
 
     @model_validator(mode="after")
-    def update_instance_name(self):
-        if self.instance_name is None:
-            self.instance_name = self.name
+    def validate_connection_type(self) -> Self:
+        """Validate connection configuration based on type.
 
+        - If instance_name is provided: Databricks Lakebase connection
+          (host is optional - will be fetched from API if not provided)
+        - If only host is provided: Standard PostgreSQL connection
+          (must not have instance_name)
+        """
+        if not self.instance_name and not self.host:
+            raise ValueError(
+                "Either instance_name (Databricks Lakebase) or host (PostgreSQL) must be provided."
+            )
         return self
 
     @model_validator(mode="after")
-    def update_user(self):
-        if self.client_id or self.user:
+    def update_user(self) -> Self:
+        # Skip if using OBO (passive auth), explicit credentials, or explicit user
+        if self.on_behalf_of_user or self.client_id or self.user or self.pat:
             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."
-            )
+        # For standard PostgreSQL, we need explicit user credentials
+        # For Lakebase with no auth, ambient auth is allowed
+        if not self.is_lakebase:
+            # Standard PostgreSQL - try to determine current user for local development
+            try:
+                self.user = self.workspace_client.current_user.me().user_name
+            except Exception as e:
+                logger.warning(
+                    f"Could not determine current user for PostgreSQL database: {e}. "
+                    f"Please provide explicit user credentials."
+                )
+        else:
+            # For Lakebase, try to determine current user but don't fail if we can't
+            try:
+                self.user = self.workspace_client.current_user.me().user_name
+            except Exception:
+                # If we can't determine user and no explicit auth, that's okay
+                # for Lakebase with ambient auth - credentials will be injected at runtime
+                pass
 
         return self
 
     @model_validator(mode="after")
-    def update_host(self):
+    def update_host(self) -> 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
+        # If instance_name is provided (Lakebase), try to fetch host from existing instance
+        # This may fail for OBO/ambient auth during model logging (before deployment)
+        if self.is_lakebase:
+            try:
+                existing_instance: DatabaseInstance = (
+                    self.workspace_client.database.get_database_instance(
+                        name=self.instance_name
+                    )
+                )
+                self.host = existing_instance.read_write_dns
+            except Exception as e:
+                # For Lakebase with OBO/ambient auth, we can't fetch at config time
+                # The host will need to be provided explicitly or fetched at runtime
+                if self.on_behalf_of_user:
+                    logger.debug(
+                        f"Could not fetch host for database {self.instance_name} "
+                        f"(Lakebase with OBO mode - will be resolved at runtime): {e}"
+                    )
+                else:
+                    raise ValueError(
+                        f"Could not fetch host for database {self.instance_name}. "
+                        f"Please provide the 'host' explicitly or ensure the instance exists: {e}"
+                    )
         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,
             self.client_secret,
         ]
         has_oauth: bool = all(field is not None for field in oauth_fields)
+        has_user_auth: bool = self.user is not None
+        has_obo: bool = self.on_behalf_of_user is True
+        has_pat: bool = self.pat is not None
 
-        pat_fields: Sequence[Any] = [self.user]
-        has_user_auth: bool = all(field is not None for field in pat_fields)
+        # Count how many auth methods are configured
+        auth_methods_count: int = sum([has_oauth, has_user_auth, has_obo, has_pat])
 
-        if has_oauth and has_user_auth:
+        if auth_methods_count > 1:
             raise ValueError(
-                "Cannot use both OAuth and user authentication methods. "
-                "Please provide either OAuth credentials or user credentials."
+                "Cannot mix authentication methods. "
+                "Please provide exactly one of: "
+                "on_behalf_of_user=true (for passive auth in model serving), "
+                "OAuth credentials (service_principal or client_id + client_secret + workspace_host), "
+                "PAT (personal access token), "
+                "or user credentials (user)."
             )
 
-        if not has_oauth and not has_user_auth:
+        # For standard PostgreSQL (host-based), at least one auth method must be configured
+        # For Lakebase (instance_name-based), auth is optional (supports ambient authentication)
+        if not self.is_lakebase and auth_methods_count == 0:
             raise ValueError(
-                "At least one authentication method must be provided: "
-                "either OAuth credentials (workspace_host, client_id, client_secret) "
-                "or user credentials (user, password)."
+                "PostgreSQL databases require explicit authentication. "
+                "Please provide one of: "
+                "OAuth credentials (workspace_host, client_id, client_secret), "
+                "service_principal with workspace_host, "
+                "PAT (personal access token), "
+                "or user credentials (user)."
             )
 
         return self
@@ -804,38 +1339,76 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
         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.
+
+        For Lakebase: Uses Databricks-generated credentials (token-based auth).
+        For standard PostgreSQL: Uses provided user/password credentials.
         """
-        from dao_ai.providers.base import ServiceProvider
-        from dao_ai.providers.databricks import DatabricksProvider
+        import uuid as _uuid
 
+        from databricks.sdk.service.database import DatabaseCredential
+
+        host: str
+        port: int
+        database: str
         username: str | None = None
+        password_value: str | None = None
+
+        # Resolve host - may need to fetch at runtime for OBO mode
+        host_value: Any = self.host
+        if host_value is None and self.is_lakebase and self.on_behalf_of_user:
+            # Fetch host at runtime for OBO mode
+            existing_instance: DatabaseInstance = (
+                self.workspace_client.database.get_database_instance(
+                    name=self.instance_name
+                )
+            )
+            host_value = existing_instance.read_write_dns
 
-        if self.client_id and self.client_secret and self.workspace_host:
-            username = value_of(self.client_id)
-        elif self.user:
-            username = value_of(self.user)
+        if host_value is None:
+            instance_or_name = self.instance_name if self.is_lakebase else self.name
+            raise ValueError(
+                f"Database host not configured for {instance_or_name}. "
+                "Please provide 'host' explicitly."
+            )
 
-        host: str = value_of(self.host)
-        port: int = value_of(self.port)
-        database: str = value_of(self.database)
+        host = value_of(host_value)
+        port = value_of(self.port)
+        database = value_of(self.database)
 
-        provider: ServiceProvider = DatabricksProvider(
-            client_id=value_of(self.client_id),
-            client_secret=value_of(self.client_secret),
-            workspace_host=value_of(self.workspace_host),
-            pat=value_of(self.password),
-        )
+        if self.is_lakebase:
+            # Lakebase: Use Databricks-generated credentials
+            if self.client_id and self.client_secret and self.workspace_host:
+                username = value_of(self.client_id)
+            elif self.user:
+                username = value_of(self.user)
+            # For OBO mode, no username is needed - the token identity is used
 
-        token: str = provider.lakebase_password_provider(self.instance_name)
+            # Generate Databricks database credential (token)
+            w: WorkspaceClient = self.workspace_client
+            cred: DatabaseCredential = w.database.generate_database_credential(
+                request_id=str(_uuid.uuid4()),
+                instance_names=[self.instance_name],
+            )
+            password_value = cred.token
+        else:
+            # Standard PostgreSQL: Use provided credentials
+            if self.user:
+                username = value_of(self.user)
+            if self.password:
+                password_value = value_of(self.password)
+
+            if not username or not password_value:
+                raise ValueError(
+                    f"Standard PostgreSQL databases require both 'user' and 'password'. "
+                    f"Database: {self.name}"
+                )
 
         # Build connection parameters dictionary
         params: dict[str, Any] = {
             "dbname": database,
             "host": host,
             "port": port,
-            "password": token,
+            "password": password_value,
             "sslmode": "require",
         }
 
@@ -866,11 +1439,86 @@ class DatabaseModel(BaseModel, IsDatabricksResource):
     def create(self, w: WorkspaceClient | None = None) -> None:
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: DatabricksProvider = DatabricksProvider()
+        # Use provided workspace client or fall back to resource's own workspace_client
+        if w is None:
+            w = self.workspace_client
+        provider: DatabricksProvider = DatabricksProvider(w=w)
         provider.create_lakebase(self)
         provider.create_lakebase_instance_role(self)
 
 
+class GenieLRUCacheParametersModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    capacity: int = 1000
+    time_to_live_seconds: int | None = (
+        60 * 60 * 24
+    )  # 1 day default, None or negative = never expires
+    warehouse: WarehouseModel
+
+
+class GenieSemanticCacheParametersModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    time_to_live_seconds: int | None = (
+        60 * 60 * 24
+    )  # 1 day default, None or negative = never expires
+    similarity_threshold: float = 0.85  # Minimum similarity for question matching (L2 distance converted to 0-1 scale)
+    context_similarity_threshold: float = 0.80  # Minimum similarity for context matching (L2 distance converted to 0-1 scale)
+    question_weight: Optional[float] = (
+        0.6  # Weight for question similarity in combined score (0-1). If not provided, computed as 1 - context_weight
+    )
+    context_weight: Optional[float] = (
+        None  # Weight for context similarity in combined score (0-1). If not provided, computed as 1 - question_weight
+    )
+    embedding_model: str | LLMModel = "databricks-gte-large-en"
+    embedding_dims: int | None = None  # Auto-detected if None
+    database: DatabaseModel
+    warehouse: WarehouseModel
+    table_name: str = "genie_semantic_cache"
+    context_window_size: int = 3  # Number of previous turns to include for context
+    max_context_tokens: int = (
+        2000  # Maximum context length to prevent extremely long embeddings
+    )
+
+    @model_validator(mode="after")
+    def compute_and_validate_weights(self) -> Self:
+        """
+        Compute missing weight and validate that question_weight + context_weight = 1.0.
+
+        Either question_weight or context_weight (or both) can be provided.
+        The missing one will be computed as 1.0 - provided_weight.
+        If both are provided, they must sum to 1.0.
+        """
+        if self.question_weight is None and self.context_weight is None:
+            # Both missing - use defaults
+            self.question_weight = 0.6
+            self.context_weight = 0.4
+        elif self.question_weight is None:
+            # Compute question_weight from context_weight
+            if not (0.0 <= self.context_weight <= 1.0):
+                raise ValueError(
+                    f"context_weight must be between 0.0 and 1.0, got {self.context_weight}"
+                )
+            self.question_weight = 1.0 - self.context_weight
+        elif self.context_weight is None:
+            # Compute context_weight from question_weight
+            if not (0.0 <= self.question_weight <= 1.0):
+                raise ValueError(
+                    f"question_weight must be between 0.0 and 1.0, got {self.question_weight}"
+                )
+            self.context_weight = 1.0 - self.question_weight
+        else:
+            # Both provided - validate they sum to 1.0
+            total_weight = self.question_weight + self.context_weight
+            if not abs(total_weight - 1.0) < 0.0001:  # Allow small floating point error
+                raise ValueError(
+                    f"question_weight ({self.question_weight}) + context_weight ({self.context_weight}) "
+                    f"must equal 1.0 (got {total_weight}). These weights determine the relative importance "
+                    f"of question vs context similarity in the combined score."
+                )
+
+        return self
+
+
 class SearchParametersModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     num_results: Optional[int] = 10
@@ -878,6 +1526,55 @@ class SearchParametersModel(BaseModel):
     query_type: Optional[str] = "ANN"
 
 
+class RerankParametersModel(BaseModel):
+    """
+    Configuration for reranking retrieved documents using FlashRank.
+
+    FlashRank provides fast, local reranking without API calls using lightweight
+    cross-encoder models. Reranking improves retrieval quality by reordering results
+    based on semantic relevance to the query.
+
+    Typical workflow:
+    1. Retrieve more documents than needed (e.g., 50 via num_results)
+    2. Rerank all retrieved documents
+    3. Return top_n best matches (e.g., 5)
+
+    Example:
+        ```yaml
+        retriever:
+          search_parameters:
+            num_results: 50  # Retrieve more candidates
+          rerank:
+            model: ms-marco-MiniLM-L-12-v2
+            top_n: 5  # Return top 5 after reranking
+        ```
+
+    Available models (from fastest to most accurate):
+    - "ms-marco-TinyBERT-L-2-v2" (fastest, smallest)
+    - "ms-marco-MiniLM-L-6-v2"
+    - "ms-marco-MiniLM-L-12-v2" (default, good balance)
+    - "rank-T5-flan" (most accurate, slower)
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+
+    model: str = Field(
+        default="ms-marco-MiniLM-L-12-v2",
+        description="FlashRank model name. Default provides good balance of speed and accuracy.",
+    )
+    top_n: Optional[int] = Field(
+        default=None,
+        description="Number of documents to return after reranking. If None, uses search_parameters.num_results.",
+    )
+    cache_dir: Optional[str] = Field(
+        default="~/.dao_ai/cache/flashrank",
+        description="Directory to cache downloaded model weights. Supports tilde expansion (e.g., ~/.dao_ai).",
+    )
+    columns: Optional[list[str]] = Field(
+        default_factory=list, description="Columns to rerank using DatabricksReranker"
+    )
+
+
 class RetrieverModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     vector_store: VectorStoreModel
@@ -885,14 +1582,25 @@ class RetrieverModel(BaseModel):
     search_parameters: SearchParametersModel = Field(
         default_factory=SearchParametersModel
     )
+    rerank: Optional[RerankParametersModel | bool] = Field(
+        default=None,
+        description="Optional reranking configuration. Set to true for defaults, or provide ReRankParametersModel for custom settings.",
+    )
 
     @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) -> Self:
+        """Convert bool to ReRankParametersModel with defaults."""
+        if isinstance(self.rerank, bool) and self.rerank:
+            self.rerank = RerankParametersModel()
+        return self
+
 
 class FunctionType(str, Enum):
     PYTHON = "python"
@@ -901,28 +1609,47 @@ class FunctionType(str, Enum):
     MCP = "mcp"
 
 
-class HumanInTheLoopActionType(str, Enum):
-    """Supported action types for human-in-the-loop interactions."""
+class HumanInTheLoopModel(BaseModel):
+    """
+    Configuration for Human-in-the-Loop tool approval.
 
-    ACCEPT = "accept"
-    EDIT = "edit"
-    RESPONSE = "response"
-    DECLINE = "decline"
+    This model configures when and how tools require human approval before execution.
+    It maps to LangChain's HumanInTheLoopMiddleware.
 
+    LangChain supports three decision types:
+    - "approve": Execute tool with original arguments
+    - "edit": Modify arguments before execution
+    - "reject": Skip execution with optional feedback message
+    """
 
-class HumanInTheLoopModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    review_prompt: str = "Please review the tool call"
-    interrupt_config: dict[str, Any] = Field(
-        default_factory=lambda: {
-            "allow_accept": True,
-            "allow_edit": True,
-            "allow_respond": True,
-            "allow_decline": True,
-        }
+
+    review_prompt: Optional[str] = Field(
+        default=None,
+        description="Message shown to the reviewer when approval is requested",
+    )
+
+    allowed_decisions: list[Literal["approve", "edit", "reject"]] = Field(
+        default_factory=lambda: ["approve", "edit", "reject"],
+        description="List of allowed decision types for this tool",
     )
-    decline_message: str = "Tool call declined by user"
-    custom_actions: Optional[dict[str, str]] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def validate_and_normalize_decisions(self) -> Self:
+        """Validate and normalize allowed decisions."""
+        if not self.allowed_decisions:
+            raise ValueError("At least one decision type must be allowed")
+
+        # Remove duplicates while preserving order
+        seen = set()
+        unique_decisions = []
+        for decision in self.allowed_decisions:
+            if decision not in seen:
+                seen.add(decision)
+                unique_decisions.append(decision)
+        self.allowed_decisions = unique_decisions
+
+        return self
 
 
 class BaseFunctionModel(ABC, BaseModel):
@@ -931,7 +1658,6 @@ class BaseFunctionModel(ABC, BaseModel):
         discriminator="type",
     )
     type: FunctionType
-    name: str
     human_in_the_loop: Optional[HumanInTheLoopModel] = None
 
     @abstractmethod
@@ -948,6 +1674,7 @@ class BaseFunctionModel(ABC, BaseModel):
 class PythonFunctionModel(BaseFunctionModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     type: Literal[FunctionType.PYTHON] = FunctionType.PYTHON
+    name: str
 
     @property
     def full_name(self) -> str:
@@ -961,8 +1688,9 @@ class PythonFunctionModel(BaseFunctionModel, HasFullName):
 
 class FactoryFunctionModel(BaseFunctionModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    args: Optional[dict[str, Any]] = Field(default_factory=dict)
     type: Literal[FunctionType.FACTORY] = FunctionType.FACTORY
+    name: str
+    args: Optional[dict[str, Any]] = Field(default_factory=dict)
 
     @property
     def full_name(self) -> str:
@@ -974,7 +1702,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
@@ -985,24 +1713,148 @@ class TransportType(str, Enum):
     STDIO = "stdio"
 
 
-class McpFunctionModel(BaseFunctionModel, HasFullName):
+class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
+    """
+    MCP Function Model with authentication inherited from IsDatabricksResource.
+
+    Authentication for MCP connections uses the same options as other resources:
+    - Service Principal (client_id + client_secret + workspace_host)
+    - PAT (pat + workspace_host)
+    - OBO (on_behalf_of_user)
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     type: Literal[FunctionType.MCP] = FunctionType.MCP
-
     transport: TransportType = TransportType.STREAMABLE_HTTP
     command: Optional[str] = "python"
     url: Optional[AnyVariable] = None
-    connection: Optional[ConnectionModel] = None
     headers: dict[str, AnyVariable] = Field(default_factory=dict)
     args: list[str] = Field(default_factory=list)
-    pat: Optional[AnyVariable] = None
-    client_id: Optional[AnyVariable] = None
-    client_secret: Optional[AnyVariable] = None
-    workspace_host: Optional[AnyVariable] = None
+    # MCP-specific fields
+    connection: Optional[ConnectionModel] = None
+    functions: Optional[SchemaModel] = None
+    genie_room: Optional[GenieRoomModel] = None
+    sql: Optional[bool] = None
+    vector_search: Optional[VectorStoreModel] = None
 
     @property
-    def full_name(self) -> str:
-        return self.name
+    def api_scopes(self) -> Sequence[str]:
+        """API scopes for MCP connections."""
+        return [
+            "serving.serving-endpoints",
+            "mcp.genie",
+            "mcp.functions",
+            "mcp.vectorsearch",
+            "mcp.external",
+        ]
+
+    def as_resources(self) -> Sequence[DatabricksResource]:
+        """MCP functions don't declare static resources."""
+        return []
+
+    def _get_workspace_host(self) -> str:
+        """
+        Get the workspace host, either from config or from workspace client.
+
+        If connection is provided, uses its workspace client.
+        Otherwise, falls back to the default Databricks host.
+
+        Returns:
+            str: The workspace host URL with https:// scheme and without trailing slash
+        """
+        from dao_ai.utils import get_default_databricks_host, normalize_host
+
+        # Try to get workspace_host from config
+        workspace_host: str | None = (
+            normalize_host(value_of(self.workspace_host))
+            if self.workspace_host
+            else None
+        )
+
+        # If no workspace_host in config, get it from workspace client
+        if not workspace_host:
+            # Use connection's workspace client if available
+            if self.connection:
+                workspace_host = normalize_host(
+                    self.connection.workspace_client.config.host
+                )
+            else:
+                # get_default_databricks_host already normalizes the 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("/")
+
+    @property
+    def mcp_url(self) -> str:
+        """
+        Get the MCP URL for this function.
+
+        Returns the URL based on the configured source:
+        - If url is set, returns it directly
+        - If connection is set, constructs URL from connection
+        - If genie_room is set, constructs Genie MCP URL
+        - If sql is set, constructs DBSQL MCP URL (serverless)
+        - If vector_search is set, constructs Vector Search MCP URL
+        - If functions is set, constructs UC Functions MCP URL
+
+        URL patterns (per https://docs.databricks.com/aws/en/generative-ai/mcp/managed-mcp):
+        - Genie: https://{host}/api/2.0/mcp/genie/{space_id}
+        - DBSQL: https://{host}/api/2.0/mcp/sql (serverless, workspace-level)
+        - Vector Search: https://{host}/api/2.0/mcp/vector-search/{catalog}/{schema}
+        - UC Functions: https://{host}/api/2.0/mcp/functions/{catalog}/{schema}
+        - Connection: https://{host}/api/2.0/mcp/external/{connection_name}
+        """
+        # Direct URL provided
+        if self.url:
+            return self.url
+
+        # Get workspace host (from config, connection, or default workspace client)
+        workspace_host: str = self._get_workspace_host()
+
+        # UC Connection
+        if self.connection:
+            connection_name: str = self.connection.name
+            return f"{workspace_host}/api/2.0/mcp/external/{connection_name}"
+
+        # Genie Room
+        if self.genie_room:
+            space_id: str = value_of(self.genie_room.space_id)
+            return f"{workspace_host}/api/2.0/mcp/genie/{space_id}"
+
+        # DBSQL MCP server (serverless, workspace-level)
+        if self.sql:
+            return f"{workspace_host}/api/2.0/mcp/sql"
+
+        # Vector Search
+        if self.vector_search:
+            if (
+                not self.vector_search.index
+                or not self.vector_search.index.schema_model
+            ):
+                raise ValueError(
+                    "vector_search must have an index with a schema (catalog/schema) configured"
+                )
+            catalog: str = self.vector_search.index.schema_model.catalog_name
+            schema: str = self.vector_search.index.schema_model.schema_name
+            return f"{workspace_host}/api/2.0/mcp/vector-search/{catalog}/{schema}"
+
+        # UC Functions MCP server
+        if self.functions:
+            catalog: str = self.functions.catalog_name
+            schema: str = self.functions.schema_name
+            return f"{workspace_host}/api/2.0/mcp/functions/{catalog}/{schema}"
+
+        raise ValueError(
+            "No URL source configured. Provide one of: url, connection, genie_room, "
+            "sql, vector_search, or functions"
+        )
 
     @field_serializer("transport")
     def serialize_transport(self, value) -> str:
@@ -1011,71 +1863,65 @@ class McpFunctionModel(BaseFunctionModel, HasFullName):
         return str(value)
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self):
-        if self.transport == TransportType.STREAMABLE_HTTP and not (
-            self.url or self.connection
-        ):
-            raise ValueError(
-                "url or connection must be provided for STREAMABLE_HTTP transport"
-            )
-        if self.transport == TransportType.STDIO and not self.command:
-            raise ValueError("command must not be provided for STDIO transport")
-        if self.transport == TransportType.STDIO and not self.args:
-            raise ValueError("args must not be provided for STDIO transport")
+    def validate_mutually_exclusive(self) -> "McpFunctionModel":
+        """Validate that exactly one URL source is provided."""
+        # Count how many URL sources are provided
+        url_sources: list[tuple[str, Any]] = [
+            ("url", self.url),
+            ("connection", self.connection),
+            ("genie_room", self.genie_room),
+            ("sql", self.sql),
+            ("vector_search", self.vector_search),
+            ("functions", self.functions),
+        ]
+
+        provided_sources: list[str] = [
+            name for name, value in url_sources if value is not None
+        ]
+
+        if self.transport == TransportType.STREAMABLE_HTTP:
+            if len(provided_sources) == 0:
+                raise ValueError(
+                    "For STREAMABLE_HTTP transport, exactly one of the following must be provided: "
+                    "url, connection, genie_room, sql, vector_search, or functions"
+                )
+            if len(provided_sources) > 1:
+                raise ValueError(
+                    f"For STREAMABLE_HTTP transport, only one URL source can be provided. "
+                    f"Found: {', '.join(provided_sources)}. "
+                    f"Please provide only one of: url, connection, genie_room, sql, vector_search, or functions"
+                )
+
+        if self.transport == TransportType.STDIO:
+            if not self.command:
+                raise ValueError("command must be provided for STDIO transport")
+            if not self.args:
+                raise ValueError("args must be provided for STDIO transport")
+
         return self
 
     @model_validator(mode="after")
-    def update_url(self):
+    def update_url(self) -> "McpFunctionModel":
         self.url = value_of(self.url)
         return self
 
     @model_validator(mode="after")
-    def update_headers(self):
+    def update_headers(self) -> "McpFunctionModel":
         for key, value in self.headers.items():
             self.headers[key] = value_of(value)
         return self
 
-    @model_validator(mode="after")
-    def validate_auth_methods(self):
-        oauth_fields: Sequence[Any] = [
-            self.client_id,
-            self.client_secret,
-        ]
-        has_oauth: bool = all(field is not None for field in oauth_fields)
-
-        pat_fields: Sequence[Any] = [self.pat]
-        has_user_auth: bool = all(field is not None for field in pat_fields)
-
-        if has_oauth and has_user_auth:
-            raise ValueError(
-                "Cannot use both OAuth and user authentication methods. "
-                "Please provide either OAuth credentials or user credentials."
-            )
-
-        if (has_oauth or has_user_auth) and not self.workspace_host:
-            raise ValueError(
-                "Workspace host must be provided when using OAuth or user credentials."
-            )
-
-        return self
-
     def as_tools(self, **kwargs: Any) -> Sequence[RunnableLike]:
         from dao_ai.tools import create_mcp_tools
 
         return create_mcp_tools(self)
 
 
-class UnityCatalogFunctionModel(BaseFunctionModel, HasFullName):
+class UnityCatalogFunctionModel(BaseFunctionModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
-    partial_args: Optional[dict[str, AnyVariable]] = Field(default_factory=dict)
     type: Literal[FunctionType.UNITY_CATALOG] = FunctionType.UNITY_CATALOG
-
-    @property
-    def full_name(self) -> str:
-        if self.schema_model:
-            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}.{self.name}"
-        return self.name
+    resource: FunctionModel
+    partial_args: Optional[dict[str, AnyVariable]] = Field(default_factory=dict)
 
     def as_tools(self, **kwargs: Any) -> Sequence[RunnableLike]:
         from dao_ai.tools import create_uc_tools
@@ -1100,12 +1946,97 @@ class ToolModel(BaseModel):
     function: AnyTool
 
 
-class GuardrailModel(BaseModel):
-    model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    name: str
-    model: LLMModel
-    prompt: str
-    num_retries: Optional[int] = 3
+class PromptModel(BaseModel, HasFullName):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
+    name: str
+    description: Optional[str] = None
+    default_template: Optional[str] = None
+    alias: Optional[str] = None
+    version: Optional[int] = None
+    tags: Optional[dict[str, Any]] = Field(default_factory=dict)
+    auto_register: bool = Field(
+        default=False,
+        description="Whether to automatically register the default_template to the prompt registry. "
+        "If False, the prompt will only be loaded from the registry (never created/updated). "
+        "Defaults to True for backward compatibility.",
+    )
+
+    @property
+    def template(self) -> str:
+        from dao_ai.providers.databricks import DatabricksProvider
+
+        provider: DatabricksProvider = DatabricksProvider()
+        prompt_version = provider.get_prompt(self)
+        return prompt_version.to_single_brace_format()
+
+    @property
+    def full_name(self) -> str:
+        prompt_name: str = self.name
+        if self.schema_model:
+            prompt_name = f"{self.schema_model.full_name}.{prompt_name}"
+        return prompt_name
+
+    @property
+    def uri(self) -> str:
+        prompt_uri: str = f"prompts:/{self.full_name}"
+
+        if self.alias:
+            prompt_uri = f"prompts:/{self.full_name}@{self.alias}"
+        elif self.version:
+            prompt_uri = f"prompts:/{self.full_name}/{self.version}"
+        else:
+            prompt_uri = f"prompts:/{self.full_name}@latest"
+
+        return prompt_uri
+
+    def as_prompt(self) -> PromptVersion:
+        prompt_version: PromptVersion = load_prompt(self.uri)
+        return prompt_version
+
+    @model_validator(mode="after")
+    def validate_mutually_exclusive(self) -> Self:
+        if self.alias and self.version:
+            raise ValueError("Cannot specify both alias and version")
+        return self
+
+
+class GuardrailModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: str
+    model: str | LLMModel
+    prompt: str | PromptModel
+    num_retries: Optional[int] = 3
+
+    @model_validator(mode="after")
+    def validate_llm_model(self) -> Self:
+        if isinstance(self.model, str):
+            self.model = LLMModel(name=self.model)
+        return self
+
+
+class MiddlewareModel(BaseModel):
+    """Configuration for middleware that can be applied to agents.
+
+    Middleware is defined at the AppConfig level and can be referenced by name
+    in agent configurations using YAML anchors for reusability.
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: str = Field(
+        description="Fully qualified name of the middleware factory function"
+    )
+    args: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Arguments to pass to the middleware factory function",
+    )
+
+    @model_validator(mode="after")
+    def resolve_args(self) -> Self:
+        """Resolve any variable references in args."""
+        for key, value in self.args.items():
+            self.args[key] = value_of(value)
+        return self
 
 
 class StorageType(str, Enum):
@@ -1116,14 +2047,12 @@ class StorageType(str, Enum):
 class CheckpointerModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
-    type: Optional[StorageType] = StorageType.MEMORY
     database: Optional[DatabaseModel] = None
 
-    @model_validator(mode="after")
-    def validate_postgres_requires_database(self):
-        if self.type == StorageType.POSTGRES and not self.database:
-            raise ValueError("Database must be provided when storage type is POSTGRES")
-        return self
+    @property
+    def storage_type(self) -> StorageType:
+        """Infer storage type from database presence."""
+        return StorageType.POSTGRES if self.database else StorageType.MEMORY
 
     def as_checkpointer(self) -> BaseCheckpointSaver:
         from dao_ai.memory import CheckpointManager
@@ -1139,16 +2068,14 @@ class StoreModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     embedding_model: Optional[LLMModel] = None
-    type: Optional[StorageType] = StorageType.MEMORY
     dims: Optional[int] = 1536
     database: Optional[DatabaseModel] = None
     namespace: Optional[str] = None
 
-    @model_validator(mode="after")
-    def validate_postgres_requires_database(self):
-        if self.type == StorageType.POSTGRES and not self.database:
-            raise ValueError("Database must be provided when storage type is POSTGRES")
-        return self
+    @property
+    def storage_type(self) -> StorageType:
+        """Infer storage type from database presence."""
+        return StorageType.POSTGRES if self.database else StorageType.MEMORY
 
     def as_store(self) -> BaseStore:
         from dao_ai.memory import StoreManager
@@ -1166,41 +2093,158 @@ class MemoryModel(BaseModel):
 FunctionHook: TypeAlias = PythonFunctionModel | FactoryFunctionModel | str
 
 
-class PromptModel(BaseModel, HasFullName):
+class ResponseFormatModel(BaseModel):
+    """
+    Configuration for structured response formats.
+
+    The response_schema field accepts either a type or a string:
+    - Type (Pydantic model, dataclass, etc.): Used directly for structured output
+    - String: First attempts to load as a fully qualified type name, falls back to JSON schema string
+
+    This unified approach simplifies the API while maintaining flexibility.
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
-    name: str
-    description: Optional[str] = None
-    default_template: Optional[str] = None
-    alias: Optional[str] = None
-    version: Optional[int] = None
-    tags: Optional[dict[str, Any]] = Field(default_factory=dict)
+    use_tool: Optional[bool] = Field(
+        default=None,
+        description=(
+            "Strategy for structured output: "
+            "None (default) = auto-detect from model capabilities, "
+            "False = force ProviderStrategy (native), "
+            "True = force ToolStrategy (function calling)"
+        ),
+    )
+    response_schema: Optional[str | type] = Field(
+        default=None,
+        description="Type or string for response format. String attempts FQN import, falls back to JSON schema.",
+    )
 
-    @property
-    def template(self) -> str:
-        from dao_ai.providers.databricks import DatabricksProvider
+    def as_strategy(self) -> ProviderStrategy | ToolStrategy:
+        """
+        Convert response_schema to appropriate LangChain strategy.
 
-        provider: DatabricksProvider = DatabricksProvider()
-        prompt: str = provider.get_prompt(self)
-        return prompt
+        Returns:
+            - None if no response_schema configured
+            - Raw schema/type for auto-detection (when use_tool=None)
+            - ToolStrategy wrapping the schema (when use_tool=True)
+            - ProviderStrategy wrapping the schema (when use_tool=False)
 
-    @property
-    def full_name(self) -> str:
-        if self.schema_model:
-            name: str = ""
-            if self.name:
-                name = f".{self.name}"
-            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}{name}"
-        return self.name
+        Raises:
+            ValueError: If response_schema is a JSON schema string that cannot be parsed
+        """
+
+        if self.response_schema is None:
+            return None
+
+        schema = self.response_schema
+
+        # Handle type schemas (Pydantic, dataclass, etc.)
+        if self.is_type_schema:
+            if self.use_tool is None:
+                # Auto-detect: Pass schema directly, let LangChain decide
+                return schema
+            elif self.use_tool is True:
+                # Force ToolStrategy (function calling)
+                return ToolStrategy(schema)
+            else:  # use_tool is False
+                # Force ProviderStrategy (native structured output)
+                return ProviderStrategy(schema)
+
+        # Handle JSON schema strings
+        elif self.is_json_schema:
+            import json
+
+            try:
+                schema_dict = json.loads(schema)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON schema string: {e}") from e
+
+            # Apply same use_tool logic as type schemas
+            if self.use_tool is None:
+                # Auto-detect
+                return schema_dict
+            elif self.use_tool is True:
+                # Force ToolStrategy
+                return ToolStrategy(schema_dict)
+            else:  # use_tool is False
+                # Force ProviderStrategy
+                return ProviderStrategy(schema_dict)
+
+        return None
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self):
-        if self.alias and self.version:
-            raise ValueError("Cannot specify both alias and version")
-        return self
+    def validate_response_schema(self) -> Self:
+        """
+        Validate and convert response_schema.
+
+        Processing logic:
+        1. If None: no response format specified
+        2. If type: use directly as structured output type
+        3. If str: try to load as FQN using type_from_fqn
+           - Success: response_schema becomes the loaded type
+           - Failure: keep as string (treated as JSON schema)
+
+        After validation, response_schema is one of:
+        - None (no schema)
+        - type (use for structured output)
+        - str (JSON schema)
+
+        Returns:
+            Self with validated response_schema
+        """
+        if self.response_schema is None:
+            return self
+
+        # If already a type, return
+        if isinstance(self.response_schema, type):
+            return self
+
+        # If it's a string, try to load as type, fallback to json_schema
+        if isinstance(self.response_schema, str):
+            from dao_ai.utils import type_from_fqn
+
+            try:
+                resolved_type = type_from_fqn(self.response_schema)
+                self.response_schema = resolved_type
+                logger.debug(
+                    f"Resolved response_schema string to type: {resolved_type}"
+                )
+                return self
+            except (ValueError, ImportError, AttributeError, TypeError) as e:
+                # Keep as string - it's a JSON schema
+                logger.debug(
+                    f"Could not resolve '{self.response_schema}' as type: {e}. "
+                    f"Treating as JSON schema string."
+                )
+                return self
+
+        # Invalid type
+        raise ValueError(
+            f"response_schema must be None, type, or str, got {type(self.response_schema)}"
+        )
+
+    @property
+    def is_type_schema(self) -> bool:
+        """Returns True if response_schema is a type (not JSON schema string)."""
+        return isinstance(self.response_schema, type)
+
+    @property
+    def is_json_schema(self) -> bool:
+        """Returns True if response_schema is a JSON schema string (not a type)."""
+        return isinstance(self.response_schema, str)
 
 
 class AgentModel(BaseModel):
+    """
+    Configuration model for an agent in the DAO AI framework.
+
+    Agents combine an LLM with tools and middleware to create systems that can
+    reason about tasks, decide which tools to use, and iteratively work towards solutions.
+
+    Middleware replaces the previous pre_agent_hook and post_agent_hook patterns,
+    providing a more flexible and composable way to customize agent behavior.
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     name: str
     description: Optional[str] = None
@@ -1209,9 +2253,54 @@ class AgentModel(BaseModel):
     guardrails: list[GuardrailModel] = Field(default_factory=list)
     prompt: Optional[str | PromptModel] = None
     handoff_prompt: Optional[str] = None
-    create_agent_hook: Optional[FunctionHook] = None
-    pre_agent_hook: Optional[FunctionHook] = None
-    post_agent_hook: Optional[FunctionHook] = None
+    middleware: list[MiddlewareModel] = Field(
+        default_factory=list,
+        description="List of middleware to apply to this agent",
+    )
+    response_format: Optional[ResponseFormatModel | type | str] = None
+
+    @model_validator(mode="after")
+    def validate_response_format(self) -> Self:
+        """
+        Validate and normalize response_format.
+
+        Accepts:
+        - None (no response format)
+        - ResponseFormatModel (already validated)
+        - type (Pydantic model, dataclass, etc.) - converts to ResponseFormatModel
+        - str (FQN or json_schema) - converts to ResponseFormatModel (smart fallback)
+
+        ResponseFormatModel handles the logic of trying FQN import and falling back to JSON schema.
+        """
+        if self.response_format is None or isinstance(
+            self.response_format, ResponseFormatModel
+        ):
+            return self
+
+        # Convert type or str to ResponseFormatModel
+        # ResponseFormatModel's validator will handle the smart type loading and fallback
+        if isinstance(self.response_format, (type, str)):
+            self.response_format = ResponseFormatModel(
+                response_schema=self.response_format
+            )
+            return self
+
+        # Invalid type
+        raise ValueError(
+            f"response_format must be None, ResponseFormatModel, type, or str, "
+            f"got {type(self.response_format)}"
+        )
+
+    def as_runnable(self) -> RunnableLike:
+        from dao_ai.nodes import create_agent_node
+
+        return create_agent_node(self)
+
+    def as_responses_agent(self) -> ResponsesAgent:
+        from dao_ai.models import create_responses_agent
+
+        graph: CompiledStateGraph = self.as_runnable()
+        return create_responses_agent(graph)
 
 
 class SupervisorModel(BaseModel):
@@ -1219,12 +2308,20 @@ class SupervisorModel(BaseModel):
     model: LLMModel
     tools: list[ToolModel] = Field(default_factory=list)
     prompt: Optional[str] = None
+    middleware: list[MiddlewareModel] = Field(
+        default_factory=list,
+        description="List of middleware to apply to the supervisor",
+    )
 
 
 class SwarmModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     model: LLMModel
     default_agent: Optional[AgentModel | str] = None
+    middleware: list[MiddlewareModel] = Field(
+        default_factory=list,
+        description="List of middleware to apply to all agents in the swarm",
+    )
     handoffs: Optional[dict[str, Optional[list[AgentModel | str]]]] = Field(
         default_factory=dict
     )
@@ -1237,7 +2334,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:
@@ -1267,9 +2364,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"
@@ -1330,27 +2439,81 @@ class ChatPayload(BaseModel):
 
         return self
 
+    @model_validator(mode="after")
+    def ensure_thread_id(self) -> "ChatPayload":
+        """Ensure thread_id or conversation_id is present in configurable, generating UUID if needed."""
+        import uuid
+
+        if self.custom_inputs is None:
+            self.custom_inputs = {}
+
+        # Get or create configurable section
+        configurable: dict[str, Any] = self.custom_inputs.get("configurable", {})
+
+        # Check if thread_id or conversation_id exists
+        has_thread_id = configurable.get("thread_id") is not None
+        has_conversation_id = configurable.get("conversation_id") is not None
+
+        # If neither is provided, generate a UUID for conversation_id
+        if not has_thread_id and not has_conversation_id:
+            configurable["conversation_id"] = str(uuid.uuid4())
+            self.custom_inputs["configurable"] = configurable
+
+        return self
+
+    def as_messages(self) -> Sequence[BaseMessage]:
+        return messages_from_dict(
+            [{"type": m.role, "content": m.content} for m in self.messages]
+        )
+
+    def as_agent_request(self) -> ResponsesAgentRequest:
+        from mlflow.types.responses_helpers import Message as _Message
+
+        return ResponsesAgentRequest(
+            input=[_Message(role=m.role, content=m.content) for m in self.messages],
+            custom_inputs=self.custom_inputs,
+        )
+
 
 class ChatHistoryModel(BaseModel):
+    """
+    Configuration for chat history summarization.
+
+    Attributes:
+        model: The LLM to use for generating summaries.
+        max_tokens: Maximum tokens to keep after summarization (the "keep" threshold).
+            After summarization, recent messages totaling up to this many tokens are preserved.
+        max_tokens_before_summary: Token threshold that triggers summarization.
+            When conversation exceeds this, summarization runs. Mutually exclusive with
+            max_messages_before_summary. If neither is set, defaults to max_tokens * 10.
+        max_messages_before_summary: Message count threshold that triggers summarization.
+            When conversation exceeds this many messages, summarization runs.
+            Mutually exclusive with max_tokens_before_summary.
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     model: LLMModel
-    max_tokens: int = 256
-    max_tokens_before_summary: Optional[int] = None
-    max_messages_before_summary: Optional[int] = None
-    max_summary_tokens: int = 255
-
-    @model_validator(mode="after")
-    def validate_max_summary_tokens(self) -> "ChatHistoryModel":
-        if self.max_summary_tokens >= self.max_tokens:
-            raise ValueError(
-                f"max_summary_tokens ({self.max_summary_tokens}) must be less than max_tokens ({self.max_tokens})"
-            )
-        return self
+    max_tokens: int = Field(
+        default=2048,
+        gt=0,
+        description="Maximum tokens to keep after summarization",
+    )
+    max_tokens_before_summary: Optional[int] = Field(
+        default=None,
+        gt=0,
+        description="Token threshold that triggers summarization",
+    )
+    max_messages_before_summary: Optional[int] = Field(
+        default=None,
+        gt=0,
+        description="Message count threshold that triggers summarization",
+    )
 
 
 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
@@ -1371,23 +2534,54 @@ class AppModel(BaseModel):
     shutdown_hooks: Optional[FunctionHook | list[FunctionHook]] = Field(
         default_factory=list
     )
-    message_hooks: Optional[FunctionHook | list[FunctionHook]] = Field(
-        default_factory=list
-    )
     input_example: Optional[ChatPayload] = None
     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):
+    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:
@@ -1397,7 +2591,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]
@@ -1417,14 +2611,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
@@ -1432,7 +2626,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:
@@ -1459,6 +2653,202 @@ class EvaluationModel(BaseModel):
     guidelines: list[GuidelineModel] = Field(default_factory=list)
 
 
+class EvaluationDatasetExpectationsModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    expected_response: Optional[str] = None
+    expected_facts: Optional[list[str]] = None
+
+    @model_validator(mode="after")
+    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
+
+
+class EvaluationDatasetEntryModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    inputs: ChatPayload
+    expectations: EvaluationDatasetExpectationsModel
+
+    def to_mlflow_format(self) -> dict[str, Any]:
+        """
+        Convert to MLflow evaluation dataset format.
+
+        Flattens the expectations fields to the top level alongside inputs,
+        which is the format expected by MLflow's Correctness scorer.
+
+        Returns:
+            dict: Flattened dictionary with inputs and expectation fields at top level
+        """
+        result: dict[str, Any] = {"inputs": self.inputs.model_dump()}
+
+        # Flatten expectations to top level for MLflow compatibility
+        if self.expectations.expected_response is not None:
+            result["expected_response"] = self.expectations.expected_response
+        if self.expectations.expected_facts is not None:
+            result["expected_facts"] = self.expectations.expected_facts
+
+        return result
+
+
+class EvaluationDatasetModel(BaseModel, HasFullName):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    schema_model: Optional[SchemaModel] = Field(default=None, alias="schema")
+    name: str
+    data: Optional[list[EvaluationDatasetEntryModel]] = Field(default_factory=list)
+    overwrite: Optional[bool] = False
+
+    def as_dataset(self, w: WorkspaceClient | None = None) -> EvaluationDataset:
+        evaluation_dataset: EvaluationDataset
+        needs_creation: bool = False
+
+        try:
+            evaluation_dataset = get_dataset(name=self.full_name)
+            if self.overwrite:
+                logger.warning(f"Overwriting dataset {self.full_name}")
+                workspace_client: WorkspaceClient = w if w else WorkspaceClient()
+                logger.debug(f"Dropping table: {self.full_name}")
+                workspace_client.tables.delete(full_name=self.full_name)
+                needs_creation = True
+        except Exception:
+            logger.warning(
+                f"Dataset {self.full_name} not found, will create new dataset"
+            )
+            needs_creation = True
+
+        # Create dataset if needed (either new or after overwrite)
+        if needs_creation:
+            evaluation_dataset = create_dataset(name=self.full_name)
+            if self.data:
+                logger.debug(
+                    f"Merging {len(self.data)} entries into dataset {self.full_name}"
+                )
+                # Use to_mlflow_format() to flatten expectations for MLflow compatibility
+                evaluation_dataset.merge_records(
+                    [e.to_mlflow_format() for e in self.data]
+                )
+
+        return evaluation_dataset
+
+    @property
+    def full_name(self) -> str:
+        if self.schema_model:
+            return f"{self.schema_model.catalog_name}.{self.schema_model.schema_name}.{self.name}"
+        return self.name
+
+
+class PromptOptimizationModel(BaseModel):
+    """Configuration for prompt optimization using GEPA.
+
+    GEPA (Generative Evolution of Prompts and Agents) is an evolutionary
+    optimizer that uses reflective mutation to improve prompts based on
+    evaluation feedback.
+
+    Example:
+        prompt_optimization:
+          name: optimize_my_prompt
+          prompt: *my_prompt
+          agent: *my_agent
+          dataset: *my_training_dataset
+          reflection_model: databricks-meta-llama-3-3-70b-instruct
+          num_candidates: 50
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    name: str
+    prompt: Optional[PromptModel] = None
+    agent: AgentModel
+    dataset: EvaluationDatasetModel  # Training dataset with examples
+    reflection_model: Optional[LLMModel | str] = None
+    num_candidates: Optional[int] = 50
+
+    def optimize(self, w: WorkspaceClient | None = None) -> PromptModel:
+        """
+        Optimize the prompt using GEPA.
+
+        Args:
+            w: Optional WorkspaceClient (not used, kept for API compatibility)
+
+        Returns:
+            PromptModel: The optimized prompt model
+        """
+        from dao_ai.optimization import OptimizationResult, optimize_prompt
+
+        # Get reflection model name
+        reflection_model_name: str | None = None
+        if self.reflection_model:
+            if isinstance(self.reflection_model, str):
+                reflection_model_name = self.reflection_model
+            else:
+                reflection_model_name = self.reflection_model.uri
+
+        # Ensure prompt is set
+        prompt = self.prompt
+        if prompt is None:
+            raise ValueError(
+                f"Prompt optimization '{self.name}' requires a prompt to be set"
+            )
+
+        result: OptimizationResult = optimize_prompt(
+            prompt=prompt,
+            agent=self.agent,
+            dataset=self.dataset,
+            reflection_model=reflection_model_name,
+            num_candidates=self.num_candidates or 50,
+            register_if_improved=True,
+        )
+
+        return result.optimized_prompt
+
+    @model_validator(mode="after")
+    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):
+                self.prompt = self.agent.prompt
+            else:
+                raise ValueError(
+                    f"Prompt optimization '{self.name}' requires either an explicit prompt "
+                    f"or an agent with a prompt configured"
+                )
+
+        return self
+
+
+class OptimizationsModel(BaseModel):
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+    training_datasets: dict[str, EvaluationDatasetModel] = Field(default_factory=dict)
+    prompt_optimizations: dict[str, PromptOptimizationModel] = Field(
+        default_factory=dict
+    )
+
+    def optimize(self, w: WorkspaceClient | None = None) -> dict[str, PromptModel]:
+        """
+        Optimize all prompts in this configuration.
+
+        This method:
+        1. Ensures all training datasets are created/registered in MLflow
+        2. Runs each prompt optimization
+
+        Args:
+            w: Optional WorkspaceClient for Databricks operations
+
+        Returns:
+            dict[str, PromptModel]: Dictionary mapping optimization names to optimized prompts
+        """
+        # First, ensure all training datasets are created/registered in MLflow
+        logger.info(f"Ensuring {len(self.training_datasets)} training datasets exist")
+        for dataset_name, dataset_model in self.training_datasets.items():
+            logger.debug(f"Creating/updating dataset: {dataset_name}")
+            dataset_model.as_dataset()
+
+        # Run optimizations
+        results: dict[str, PromptModel] = {}
+        for name, optimization in self.prompt_optimizations.items():
+            results[name] = optimization.optimize(w)
+        return results
+
+
 class DatasetFormat(str, Enum):
     CSV = "csv"
     DELTA = "delta"
@@ -1494,7 +2884,7 @@ class UnityCatalogFunctionSqlTestModel(BaseModel):
 
 class UnityCatalogFunctionSqlModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    function: UnityCatalogFunctionModel
+    function: FunctionModel
     ddl: str
     parameters: Optional[dict[str, Any]] = Field(default_factory=dict)
     test: Optional[UnityCatalogFunctionSqlTestModel] = None
@@ -1522,21 +2912,132 @@ class ResourcesModel(BaseModel):
     warehouses: dict[str, WarehouseModel] = Field(default_factory=dict)
     databases: dict[str, DatabaseModel] = Field(default_factory=dict)
     connections: dict[str, ConnectionModel] = Field(default_factory=dict)
+    apps: dict[str, DatabricksAppModel] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def update_genie_warehouses(self) -> Self:
+        """
+        Automatically populate warehouses from genie_rooms.
+
+        Warehouses are extracted from each Genie room and added to the
+        resources if they don't already exist (based on warehouse_id).
+        """
+        if not self.genie_rooms:
+            return self
+
+        # Process warehouses from all genie rooms
+        for genie_room in self.genie_rooms.values():
+            genie_room: GenieRoomModel
+            warehouse: Optional[WarehouseModel] = genie_room.warehouse
+
+            if warehouse is None:
+                continue
+
+            # Check if warehouse already exists based on warehouse_id
+            warehouse_exists: bool = any(
+                existing_warehouse.warehouse_id == warehouse.warehouse_id
+                for existing_warehouse in self.warehouses.values()
+            )
+
+            if not warehouse_exists:
+                warehouse_key: str = normalize_name(
+                    "_".join([genie_room.name, warehouse.warehouse_id])
+                )
+                self.warehouses[warehouse_key] = warehouse
+                logger.trace(
+                    "Added warehouse from Genie room",
+                    room=genie_room.name,
+                    warehouse=warehouse.warehouse_id,
+                    key=warehouse_key,
+                )
+
+        return self
+
+    @model_validator(mode="after")
+    def update_genie_tables(self) -> Self:
+        """
+        Automatically populate tables from genie_rooms.
+
+        Tables are extracted from each Genie room and added to the
+        resources if they don't already exist (based on full_name).
+        """
+        if not self.genie_rooms:
+            return self
+
+        # Process tables from all genie rooms
+        for genie_room in self.genie_rooms.values():
+            genie_room: GenieRoomModel
+            for table in genie_room.tables:
+                table: TableModel
+                table_exists: bool = any(
+                    existing_table.full_name == table.full_name
+                    for existing_table in self.tables.values()
+                )
+                if not table_exists:
+                    table_key: str = normalize_name(
+                        "_".join([genie_room.name, table.full_name])
+                    )
+                    self.tables[table_key] = table
+                    logger.trace(
+                        "Added table from Genie room",
+                        room=genie_room.name,
+                        table=table.name,
+                        key=table_key,
+                    )
+
+        return self
+
+    @model_validator(mode="after")
+    def update_genie_functions(self) -> Self:
+        """
+        Automatically populate functions from genie_rooms.
+
+        Functions are extracted from each Genie room and added to the
+        resources if they don't already exist (based on full_name).
+        """
+        if not self.genie_rooms:
+            return self
+
+        # Process functions from all genie rooms
+        for genie_room in self.genie_rooms.values():
+            genie_room: GenieRoomModel
+            for function in genie_room.functions:
+                function: FunctionModel
+                function_exists: bool = any(
+                    existing_function.full_name == function.full_name
+                    for existing_function in self.functions.values()
+                )
+                if not function_exists:
+                    function_key: str = normalize_name(
+                        "_".join([genie_room.name, function.full_name])
+                    )
+                    self.functions[function_key] = function
+                    logger.trace(
+                        "Added function from Genie room",
+                        room=genie_room.name,
+                        function=function.name,
+                        key=function_key,
+                    )
+
+        return self
 
 
 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)
     tools: dict[str, ToolModel] = Field(default_factory=dict)
     guardrails: dict[str, GuardrailModel] = Field(default_factory=dict)
+    middleware: dict[str, MiddlewareModel] = Field(default_factory=dict)
     memory: Optional[MemoryModel] = None
     prompts: dict[str, PromptModel] = Field(default_factory=dict)
     agents: dict[str, AgentModel] = Field(default_factory=dict)
     app: Optional[AppModel] = None
     evaluation: Optional[EvaluationModel] = None
+    optimizations: Optional[OptimizationsModel] = None
     datasets: Optional[list[DatasetModel]] = Field(default_factory=list)
     unity_catalog_functions: Optional[list[UnityCatalogFunctionSqlModel]] = Field(
         default_factory=list
@@ -1558,10 +3059,10 @@ class AppConfig(BaseModel):
 
     def initialize(self) -> None:
         from dao_ai.hooks.core import create_hooks
+        from dao_ai.logging import configure_logging
 
         if self.app and self.app.log_level:
-            logger.remove()
-            logger.add(sys.stderr, level=self.app.log_level)
+            configure_logging(level=self.app.log_level)
 
         logger.debug("Calling initialization hooks...")
         initialization_functions: Sequence[Callable[..., Any]] = create_hooks(
@@ -1605,21 +3106,45 @@ class AppConfig(BaseModel):
     def create_agent(
         self,
         w: WorkspaceClient | None = None,
+        vsc: "VectorSearchClient | None" = None,
+        pat: str | None = None,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        workspace_host: str | None = None,
     ) -> None:
         from dao_ai.providers.base import ServiceProvider
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: ServiceProvider = DatabricksProvider(w=w)
+        provider: ServiceProvider = DatabricksProvider(
+            w=w,
+            vsc=vsc,
+            pat=pat,
+            client_id=client_id,
+            client_secret=client_secret,
+            workspace_host=workspace_host,
+        )
         provider.create_agent(self)
 
     def deploy_agent(
         self,
         w: WorkspaceClient | None = None,
+        vsc: "VectorSearchClient | None" = None,
+        pat: str | None = None,
+        client_id: str | None = None,
+        client_secret: str | None = None,
+        workspace_host: str | None = None,
     ) -> None:
         from dao_ai.providers.base import ServiceProvider
         from dao_ai.providers.databricks import DatabricksProvider
 
-        provider: ServiceProvider = DatabricksProvider(w=w)
+        provider: ServiceProvider = DatabricksProvider(
+            w=w,
+            vsc=vsc,
+            pat=pat,
+            client_id=client_id,
+            client_secret=client_secret,
+            workspace_host=workspace_host,
+        )
         provider.deploy_agent(self)
 
     def find_agents(