dao-ai 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

dao_ai/config.py

@@ -373,15 +373,19 @@ class IsDatabricksResource(ABC, BaseModel):
         """
         from dao_ai.utils import normalize_host
 
-        logger.trace(f"workspace_client_from called", context=context, on_behalf_of_user=self.on_behalf_of_user)
+        logger.trace(
+            "workspace_client_from called",
+            context=context,
+            on_behalf_of_user=self.on_behalf_of_user,
+        )
 
         # Check if we have headers in context for OBO
         if context and context.headers and self.on_behalf_of_user:
             headers = context.headers
             # Try both lowercase and title-case header names (HTTP headers are case-insensitive)
-            forwarded_token: str = headers.get("x-forwarded-access-token") or headers.get(
-                "X-Forwarded-Access-Token"
-            )
+            forwarded_token: str = headers.get(
+                "x-forwarded-access-token"
+            ) or headers.get("X-Forwarded-Access-Token")
 
             if forwarded_token:
                 forwarded_user = headers.get("x-forwarded-user") or headers.get(
@@ -1398,13 +1402,20 @@ class DatabaseModel(IsDatabricksResource):
     - 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.
+    Note: For Lakebase connections, `name` is optional and defaults to `instance_name`.
+    For PostgreSQL connections, `name` is required.
+
+    Example Databricks Lakebase (minimal):
+    ```yaml
+    databases:
+      my_lakebase:
+        instance_name: my-lakebase-instance  # name defaults to instance_name
+    ```
 
     Example Databricks Lakebase with Service Principal:
     ```yaml
     databases:
       my_lakebase:
-        name: my-database
         instance_name: my-lakebase-instance
         service_principal:
           client_id:
@@ -1420,7 +1431,6 @@ class DatabaseModel(IsDatabricksResource):
     ```yaml
     databases:
       my_lakebase:
-        name: my-database
         instance_name: my-lakebase-instance
         on_behalf_of_user: true
     ```
@@ -1440,7 +1450,7 @@ class DatabaseModel(IsDatabricksResource):
     """
 
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    name: str
+    name: Optional[str] = None
     instance_name: Optional[str] = None
     description: Optional[str] = None
     host: Optional[AnyVariable] = None
@@ -1489,6 +1499,17 @@ class DatabaseModel(IsDatabricksResource):
             )
         return self
 
+    @model_validator(mode="after")
+    def populate_name_from_instance_name(self) -> Self:
+        """Populate name from instance_name if not provided for Lakebase connections."""
+        if self.name is None and self.instance_name:
+            self.name = self.instance_name
+        elif self.name is None:
+            raise ValueError(
+                "Either 'name' or 'instance_name' must be provided for DatabaseModel."
+            )
+        return self
+
     @model_validator(mode="after")
     def update_user(self) -> Self:
         # Skip if using OBO (passive auth), explicit credentials, or explicit user
@@ -1586,10 +1607,10 @@ class DatabaseModel(IsDatabricksResource):
         username: str | None = None
         password_value: str | None = None
 
-        # Resolve host - may need to fetch at runtime for OBO mode
+        # Resolve host - fetch from API at runtime for Lakebase if not provided
         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
+        if host_value is None and self.is_lakebase:
+            # Fetch host from Lakebase instance API
             existing_instance: DatabaseInstance = (
                 self.workspace_client.database.get_database_instance(
                     name=self.instance_name
@@ -1759,43 +1780,83 @@ class SearchParametersModel(BaseModel):
     query_type: Optional[str] = "ANN"
 
 
+class InstructionAwareRerankModel(BaseModel):
+    """
+    LLM-based reranking considering user instructions and constraints.
+
+    Use fast models (GPT-3.5, Haiku, Llama 3 8B) to minimize latency (~100ms).
+    Runs AFTER FlashRank as an additional constraint-aware reranking stage.
+    Skipped for 'standard' mode when auto_bypass=true in router config.
+
+    Example:
+        ```yaml
+        rerank:
+          model: ms-marco-MiniLM-L-12-v2
+          top_n: 20
+          instruction_aware:
+            model: *fast_llm
+            instructions: |
+              Prioritize results matching price and brand constraints.
+            top_n: 10
+        ```
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+
+    model: Optional["LLMModel"] = Field(
+        default=None,
+        description="LLM for instruction reranking (fast model recommended)",
+    )
+    instructions: Optional[str] = Field(
+        default=None,
+        description="Custom reranking instructions for constraint prioritization",
+    )
+    top_n: Optional[int] = Field(
+        default=None,
+        description="Number of documents to return after instruction reranking",
+    )
+
+
 class RerankParametersModel(BaseModel):
     """
-    Configuration for reranking retrieved documents using FlashRank.
+    Configuration for reranking retrieved documents.
 
-    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.
+    Supports three reranking options that can be combined:
+    1. FlashRank (local cross-encoder) - set `model`
+    2. Databricks server-side reranking - set `columns`
+    3. LLM instruction-aware reranking - set `instruction_aware`
 
-    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 with Databricks columns + instruction-aware (no FlashRank):
+        ```yaml
+        rerank:
+          columns:                    # Databricks server-side reranking
+            - product_name
+            - brand_name
+          instruction_aware:          # LLM-based constraint reranking
+            model: *fast_llm
+            instructions: "Prioritize by brand preferences"
+            top_n: 10
+        ```
 
-    Example:
+    Example with FlashRank:
         ```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
+        rerank:
+          model: ms-marco-MiniLM-L-12-v2  # FlashRank model
+          top_n: 10
         ```
 
-    Available models (see https://github.com/PrithivirajDamodaran/FlashRank):
+    Available FlashRank models (see https://github.com/PrithivirajDamodaran/FlashRank):
     - "ms-marco-TinyBERT-L-2-v2" (~4MB, fastest)
-    - "ms-marco-MiniLM-L-12-v2" (~34MB, best cross-encoder, default)
+    - "ms-marco-MiniLM-L-12-v2" (~34MB, best cross-encoder)
     - "rank-T5-flan" (~110MB, best non cross-encoder)
     - "ms-marco-MultiBERT-L-12" (~150MB, multilingual 100+ languages)
-    - "ce-esci-MiniLM-L12-v2" (e-commerce optimized, Amazon ESCI)
-    - "miniReranker_arabic_v1" (Arabic language)
     """
 
     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.",
+    model: Optional[str] = Field(
+        default=None,
+        description="FlashRank model name. If None, FlashRank is not used (use columns for Databricks reranking).",
     )
     top_n: Optional[int] = Field(
         default=None,
@@ -1808,6 +1869,289 @@ class RerankParametersModel(BaseModel):
     columns: Optional[list[str]] = Field(
         default_factory=list, description="Columns to rerank using DatabricksReranker"
     )
+    instruction_aware: Optional[InstructionAwareRerankModel] = Field(
+        default=None,
+        description="Optional LLM-based reranking stage after FlashRank",
+    )
+
+
+class FilterItem(BaseModel):
+    """A metadata filter for vector search.
+
+    Filters constrain search results by matching column values.
+    Use column names from the provided schema description.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+    key: str = Field(
+        description=(
+            "Column name with optional operator suffix. "
+            "Operators: (none) for equality, NOT for exclusion, "
+            "< <= > >= for numeric comparison, "
+            "LIKE for token match, NOT LIKE to exclude tokens."
+        )
+    )
+    value: Union[str, int, float, bool, list[Union[str, int, float, bool]]] = Field(
+        description=(
+            "The filter value matching the column type. "
+            "Use an array for IN-style matching multiple values."
+        )
+    )
+
+
+class SearchQuery(BaseModel):
+    """A single search query with optional metadata filters.
+
+    Represents one focused search intent extracted from the user's request.
+    The text should be a natural language query optimized for semantic search.
+    Filters constrain results to match specific metadata values.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+    text: str = Field(
+        description=(
+            "Natural language search query text optimized for semantic similarity. "
+            "Should be focused on a single search intent. "
+            "Do NOT include filter criteria in the text; use the filters field instead."
+        )
+    )
+    filters: Optional[list[FilterItem]] = Field(
+        default=None,
+        description=(
+            "Metadata filters to constrain search results. "
+            "Set to null if no filters apply. "
+            "Extract filter values from explicit constraints in the user query."
+        ),
+    )
+
+
+class DecomposedQueries(BaseModel):
+    """Decomposed search queries extracted from a user request.
+
+    Break down complex user queries into multiple focused search queries.
+    Each query targets a distinct search intent with appropriate filters.
+    Generate 1-3 queries depending on the complexity of the user request.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+    queries: list[SearchQuery] = Field(
+        description=(
+            "List of search queries extracted from the user request. "
+            "Each query should target a distinct search intent. "
+            "Order queries by importance, with the most relevant first."
+        )
+    )
+
+
+class ColumnInfo(BaseModel):
+    """Column metadata for dynamic schema generation in structured output.
+
+    When provided, column information is embedded directly into the JSON schema
+    that with_structured_output sends to the LLM, improving filter accuracy.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: str = Field(description="Column name as it appears in the database")
+    type: Literal["string", "number", "boolean", "datetime"] = Field(
+        default="string",
+        description="Column data type for value validation",
+    )
+    operators: list[str] = Field(
+        default=["", "NOT", "<", "<=", ">", ">=", "LIKE", "NOT LIKE"],
+        description="Valid filter operators for this column",
+    )
+
+
+class InstructedRetrieverModel(BaseModel):
+    """
+    Configuration for instructed retrieval with query decomposition and RRF merging.
+
+    Instructed retrieval decomposes user queries into multiple subqueries with
+    metadata filters, executes them in parallel, and merges results using
+    Reciprocal Rank Fusion (RRF) before reranking.
+
+    Example:
+        ```yaml
+        retriever:
+          vector_store: *products_vector_store
+          instructed:
+            decomposition_model: *fast_llm
+            schema_description: |
+              Products table: product_id, brand_name, category, price, updated_at
+              Filter operators: {"col": val}, {"col >": val}, {"col NOT": val}
+            columns:
+              - name: brand_name
+                type: string
+              - name: price
+                type: number
+                operators: ["", "<", "<=", ">", ">="]
+            constraints:
+              - "Prefer recent products"
+            max_subqueries: 3
+            examples:
+              - query: "cheap drills"
+                filters: {"price <": 100}
+        ```
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+
+    decomposition_model: Optional["LLMModel"] = Field(
+        default=None,
+        description="LLM for query decomposition (smaller/faster model recommended)",
+    )
+    schema_description: str = Field(
+        description="Column names, types, and valid filter syntax for the LLM"
+    )
+    columns: Optional[list[ColumnInfo]] = Field(
+        default=None,
+        description=(
+            "Structured column info for dynamic schema generation. "
+            "When provided, column names are embedded in the JSON schema for better LLM accuracy."
+        ),
+    )
+    constraints: Optional[list[str]] = Field(
+        default=None, description="Default constraints to always apply"
+    )
+    max_subqueries: int = Field(
+        default=3, description="Maximum number of parallel subqueries"
+    )
+    rrf_k: int = Field(
+        default=60,
+        description="RRF constant (lower values weight top ranks more heavily)",
+    )
+    examples: Optional[list[dict[str, Any]]] = Field(
+        default=None,
+        description="Few-shot examples for domain-specific filter translation",
+    )
+    normalize_filter_case: Optional[Literal["uppercase", "lowercase"]] = Field(
+        default=None,
+        description="Auto-normalize filter string values to uppercase or lowercase",
+    )
+
+
+class RouterModel(BaseModel):
+    """
+    Select internal execution mode based on query characteristics.
+
+    Use fast models (GPT-3.5, Haiku, Llama 3 8B) to minimize latency (~50-100ms).
+    Routes to internal modes within the same retriever, not external retrievers.
+    Cross-index routing belongs at the agent/tool-selection level.
+
+    Execution Modes:
+    - "standard": Single similarity_search() for simple keyword/product searches
+    - "instructed": Decompose -> Parallel Search -> RRF for constrained queries
+
+    Example:
+        ```yaml
+        retriever:
+          router:
+            model: *fast_llm
+            default_mode: standard
+            auto_bypass: true
+        ```
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+
+    model: Optional["LLMModel"] = Field(
+        default=None,
+        description="LLM for routing decision (fast model recommended)",
+    )
+    default_mode: Literal["standard", "instructed"] = Field(
+        default="standard",
+        description="Fallback mode if routing fails",
+    )
+    auto_bypass: bool = Field(
+        default=True,
+        description="Skip Instruction Reranker and Verifier for standard mode",
+    )
+
+
+class VerificationResult(BaseModel):
+    """Verification of whether search results satisfy the user's constraints.
+
+    Analyze the retrieved results against the original query and any explicit
+    constraints to determine if a retry with modified filters is needed.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    passed: bool = Field(
+        description="True if results satisfy the user's query intent and constraints."
+    )
+    confidence: float = Field(
+        ge=0.0,
+        le=1.0,
+        description="Confidence in the verification decision, from 0.0 (uncertain) to 1.0 (certain).",
+    )
+    feedback: Optional[str] = Field(
+        default=None,
+        description="Explanation of why verification passed or failed. Include specific issues found.",
+    )
+    suggested_filter_relaxation: Optional[dict[str, Any]] = Field(
+        default=None,
+        description=(
+            "Suggested filter modifications for retry. "
+            "Keys are column names, values indicate changes (e.g., 'REMOVE', 'WIDEN', or new values)."
+        ),
+    )
+    unmet_constraints: Optional[list[str]] = Field(
+        default=None,
+        description="List of user constraints that the results failed to satisfy.",
+    )
+
+
+class VerifierModel(BaseModel):
+    """
+    Validate results against user constraints with structured feedback.
+
+    Use fast models (GPT-3.5, Haiku, Llama 3 8B) to minimize latency (~50-100ms).
+    Skipped for 'standard' mode when auto_bypass=true in router config.
+    Returns structured feedback for intelligent retry, not blind retry.
+
+    Example:
+        ```yaml
+        retriever:
+          verifier:
+            model: *fast_llm
+            on_failure: warn_and_retry
+            max_retries: 1
+        ```
+    """
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+
+    model: Optional["LLMModel"] = Field(
+        default=None,
+        description="LLM for verification (fast model recommended)",
+    )
+    on_failure: Literal["warn", "retry", "warn_and_retry"] = Field(
+        default="warn",
+        description="Behavior when verification fails",
+    )
+    max_retries: int = Field(
+        default=1,
+        description="Maximum retry attempts before returning with warning",
+    )
+
+
+class RankedDocument(BaseModel):
+    """Single ranked document."""
+
+    index: int = Field(description="Document index from input list")
+    score: float = Field(description="0.0-1.0 relevance score")
+    reason: str = Field(default="", description="Why this score")
+
+
+class RankingResult(BaseModel):
+    """Reranking output."""
+
+    rankings: list[RankedDocument] = Field(
+        default_factory=list,
+        description="Ranked documents, highest score first",
+    )
 
 
 class RetrieverModel(BaseModel):
@@ -1817,10 +2161,22 @@ class RetrieverModel(BaseModel):
     search_parameters: SearchParametersModel = Field(
         default_factory=SearchParametersModel
     )
+    router: Optional[RouterModel] = Field(
+        default=None,
+        description="Optional query router for selecting execution mode (standard vs instructed).",
+    )
     rerank: Optional[RerankParametersModel | bool] = Field(
         default=None,
         description="Optional reranking configuration. Set to true for defaults, or provide ReRankParametersModel for custom settings.",
     )
+    instructed: Optional[InstructedRetrieverModel] = Field(
+        default=None,
+        description="Optional instructed retrieval with query decomposition and RRF merging.",
+    )
+    verifier: Optional[VerifierModel] = Field(
+        default=None,
+        description="Optional result verification with structured feedback for retry.",
+    )
 
     @model_validator(mode="after")
     def set_default_columns(self) -> Self:
@@ -1831,9 +2187,13 @@ class RetrieverModel(BaseModel):
 
     @model_validator(mode="after")
     def set_default_reranker(self) -> Self:
-        """Convert bool to ReRankParametersModel with defaults."""
+        """Convert bool to ReRankParametersModel with defaults.
+
+        When rerank: true is used, sets the default FlashRank model
+        (ms-marco-MiniLM-L-12-v2) to enable reranking.
+        """
         if isinstance(self.rerank, bool) and self.rerank:
-            self.rerank = RerankParametersModel()
+            self.rerank = RerankParametersModel(model="ms-marco-MiniLM-L-12-v2")
         return self
 
 
@@ -2091,9 +2451,47 @@ class McpFunctionModel(BaseFunctionModel, IsDatabricksResource):
             return f"{workspace_host}/api/2.0/mcp/sql"
 
         # Databricks App - MCP endpoint is at {app_url}/mcp
+        # Try McpFunctionModel's workspace_client first (which may have credentials),
+        # then fall back to DatabricksAppModel.url property (which uses its own workspace_client)
         if self.app:
-            app_url = self.app.url.rstrip("/")
-            return f"{app_url}/mcp"
+            from databricks.sdk.service.apps import App
+
+            app_url: str | None = None
+
+            # First, try using McpFunctionModel's workspace_client
+            try:
+                app: App = self.workspace_client.apps.get(self.app.name)
+                app_url = app.url
+                logger.trace(
+                    "Got app URL using McpFunctionModel workspace_client",
+                    app_name=self.app.name,
+                    url=app_url,
+                )
+            except Exception as e:
+                logger.debug(
+                    "Failed to get app URL using McpFunctionModel workspace_client, "
+                    "trying DatabricksAppModel.url property",
+                    app_name=self.app.name,
+                    error=str(e),
+                )
+
+            # Fall back to DatabricksAppModel.url property
+            if not app_url:
+                try:
+                    app_url = self.app.url
+                    logger.trace(
+                        "Got app URL using DatabricksAppModel.url property",
+                        app_name=self.app.name,
+                        url=app_url,
+                    )
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Databricks App '{self.app.name}' does not have a URL. "
+                        "The app may not be deployed yet, or credentials may be invalid. "
+                        f"Error: {e}"
+                    ) from e
+
+            return f"{app_url.rstrip('/')}/mcp"
 
         # Vector Search
         if self.vector_search:
@@ -2608,7 +3006,6 @@ class SupervisorModel(BaseModel):
 
 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,
@@ -2622,11 +3019,17 @@ class SwarmModel(BaseModel):
 class OrchestrationModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
     supervisor: Optional[SupervisorModel] = None
-    swarm: Optional[SwarmModel] = None
+    swarm: Optional[SwarmModel | Literal[True]] = None
     memory: Optional[MemoryModel] = None
 
     @model_validator(mode="after")
-    def validate_mutually_exclusive(self) -> Self:
+    def validate_and_normalize(self) -> Self:
+        """Validate orchestration and normalize swarm shorthand."""
+        # Convert swarm: true to SwarmModel()
+        if self.swarm is True:
+            self.swarm = SwarmModel()
+
+        # Validate mutually exclusive
         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:
@@ -2898,9 +3301,7 @@ class AppModel(BaseModel):
             elif len(self.agents) == 1:
                 default_agent: AgentModel = self.agents[0]
                 self.orchestration = OrchestrationModel(
-                    swarm=SwarmModel(
-                        model=default_agent.model, default_agent=default_agent
-                    )
+                    swarm=SwarmModel(default_agent=default_agent)
                 )
             else:
                 raise ValueError("At least one agent must be specified")
@@ -2940,8 +3341,24 @@ class GuidelineModel(BaseModel):
 
 
 class EvaluationModel(BaseModel):
+    """
+    Configuration for MLflow GenAI evaluation.
+
+    Attributes:
+        model: LLM model used as the judge for LLM-based scorers (e.g., Guidelines, Safety).
+               This model evaluates agent responses during evaluation.
+        table: Table to store evaluation results.
+        num_evals: Number of evaluation samples to generate.
+        agent_description: Description of the agent for evaluation data generation.
+        question_guidelines: Guidelines for generating evaluation questions.
+        custom_inputs: Custom inputs to pass to the agent during evaluation.
+        guidelines: List of guideline configurations for Guidelines scorers.
+    """
+
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    model: LLMModel
+    model: LLMModel = Field(
+        ..., description="LLM model used as the judge for LLM-based evaluation scorers"
+    )
     table: TableModel
     num_evals: int
     agent_description: Optional[str] = None
@@ -2949,6 +3366,16 @@ class EvaluationModel(BaseModel):
     custom_inputs: dict[str, Any] = Field(default_factory=dict)
     guidelines: list[GuidelineModel] = Field(default_factory=list)
 
+    @property
+    def judge_model_endpoint(self) -> str:
+        """
+        Get the judge model endpoint string for MLflow scorers.
+
+        Returns:
+            Endpoint string in format 'databricks:/model-name'
+        """
+        return f"databricks:/{self.model.name}"
+
 
 class EvaluationDatasetExpectationsModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")