dao-ai 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

dao_ai/prompts/instructed_retriever_decomposition.yaml

@@ -0,0 +1,58 @@
+name: instructed_retriever_decomposition
+description: Decomposes user queries into multiple search queries with metadata filters
+
+template: |
+  You are a search query decomposition expert. Your task is to break down a user query into one or more focused search queries with appropriate metadata filters. Respond with a JSON object.
+
+  ## Current Time
+  {current_time}
+
+  ## Database Schema
+  {schema_description}
+
+  ## Constraints
+  {constraints}
+
+  ## Few-Shot Examples
+  {examples}
+
+  ## Instructions
+  1. Analyze the user query and identify distinct search intents
+  2. For each intent, create a focused search query text
+  3. Extract metadata filters from the query using the exact filter syntax above
+  4. Resolve relative time references (e.g., "last month", "past year") using the current time
+  5. Generate at most {max_subqueries} search queries
+  6. If no filters apply, set filters to null
+
+  ## User Query
+  {query}
+
+  Generate search queries that together capture all aspects of the user's information need.
+
+variables:
+  - current_time
+  - schema_description
+  - constraints
+  - examples
+  - max_subqueries
+  - query
+
+output_format: |
+  The output must be a JSON object with a "queries" field containing an array of search query objects.
+  Each search query object has:
+  - "text": The search query string
+  - "filters": An array of filter objects, each with "key" (column + optional operator) and "value", or null if no filters
+  
+  Supported filter operators (append to column name):
+  - Equality: {"key": "column", "value": "val"} or {"key": "column", "value": ["val1", "val2"]}
+  - Exclusion: {"key": "column NOT", "value": "val"}
+  - Comparison: {"key": "column <", "value": 100}, also <=, >, >=
+  - Token match: {"key": "column LIKE", "value": "word"}
+  - Exclude token: {"key": "column NOT LIKE", "value": "word"}
+  
+  Examples:
+  - [{"key": "brand_name", "value": "MILWAUKEE"}]
+  - [{"key": "price <", "value": 100}]
+  - [{"key": "brand_name NOT", "value": "DEWALT"}]
+  - [{"key": "brand_name", "value": ["MILWAUKEE", "DEWALT"]}]
+  - [{"key": "description LIKE", "value": "cordless"}]

dao_ai/prompts/instruction_reranker.yaml

@@ -0,0 +1,14 @@
+name: instruction_aware_reranking
+version: "1.1"
+description: Rerank documents based on user instructions and constraints
+
+template: |
+  Rerank these search results for the query "{query}".
+
+  {instructions}
+
+  ## Documents
+
+  {documents}
+
+  Score each document 0.0-1.0 based on relevance to the query and instructions. Return results sorted by score (highest first). Only include documents scoring > 0.1.

dao_ai/prompts/router.yaml

@@ -0,0 +1,37 @@
+name: router_query_classification
+version: "1.0"
+description: Classify query to determine execution mode (standard vs instructed)
+
+template: |
+  You are a query classification system. Your task is to determine the best execution mode for a search query.
+
+  ## Execution Modes
+
+  **standard**: Use for simple keyword or product searches without specific constraints.
+  - General questions about products
+  - Simple keyword searches
+  - Broad category browsing
+
+  **instructed**: Use for queries with explicit constraints that require metadata filtering.
+  - Price constraints ("under $100", "between $50 and $200")
+  - Brand preferences ("Milwaukee", "not DeWalt", "excluding Makita")
+  - Category filters ("power tools", "paint supplies")
+  - Time/recency constraints ("recent", "from last month", "updated this year")
+  - Comparison queries ("compare X and Y")
+  - Multiple combined constraints
+
+  ## Available Schema for Filtering
+
+  {schema_description}
+
+  ## Query to Classify
+
+  "{query}"
+
+  ## Instructions
+
+  Analyze the query and determine:
+  1. Does it contain explicit constraints that can be translated to metadata filters?
+  2. Would the query benefit from being decomposed into subqueries?
+
+  Return your classification as a JSON object with a single field "mode" set to either "standard" or "instructed".

dao_ai/prompts/verifier.yaml

@@ -0,0 +1,46 @@
+name: result_verification
+version: "1.0"
+description: Verify search results satisfy user constraints
+
+template: |
+  You are a result verification system. Your task is to determine whether search results satisfy the user's query constraints.
+
+  ## User Query
+
+  "{query}"
+
+  ## Schema Information
+
+  {schema_description}
+
+  ## Constraints to Verify
+
+  {constraints}
+
+  ## Retrieved Results (Top {num_results})
+
+  {results_summary}
+
+  ## Previous Attempt Feedback (if retry)
+
+  {previous_feedback}
+
+  ## Instructions
+
+  Analyze whether the results satisfy the user's explicit and implicit constraints:
+
+  1. **Intent Match**: Do the results address what the user is looking for?
+  2. **Explicit Constraints**: Are price, brand, category, date constraints met?
+  3. **Relevance**: Are the results actually useful for the user's needs?
+
+  If results do NOT satisfy constraints, suggest specific filter relaxations:
+  - Use "REMOVE" to drop a filter entirely
+  - Use "BROADEN" to widen a range (e.g., price < 100 -> price < 150)
+  - Use specific values to change a filter
+
+  Return a JSON object with:
+  - passed: boolean (true if results are satisfactory)
+  - confidence: float (0.0-1.0, your confidence in the assessment)
+  - feedback: string (brief explanation of issues, if any)
+  - suggested_filter_relaxation: object (filter changes for retry, e.g., {{"brand_name": "REMOVE"}})
+  - unmet_constraints: array of strings (list of constraints not satisfied)

dao_ai/providers/databricks.py

@@ -397,6 +397,8 @@ class DatabricksProvider(ServiceProvider):
 
             pip_requirements += get_installed_packages()
 
+        code_paths = list(dict.fromkeys(code_paths))
+
         logger.trace("Pip requirements prepared", count=len(pip_requirements))
         logger.trace("Code paths prepared", count=len(code_paths))
 
@@ -434,19 +436,38 @@ class DatabricksProvider(ServiceProvider):
             pip_packages_count=len(pip_requirements),
         )
 
-        with mlflow.start_run(run_name=run_name):
-            mlflow.set_tag("type", "agent")
-            mlflow.set_tag("dao_ai", dao_ai_version())
-            logged_agent_info: ModelInfo = mlflow.pyfunc.log_model(
-                python_model=model_path.as_posix(),
-                code_paths=code_paths,
-                model_config=config.model_dump(mode="json", by_alias=True),
-                name="agent",
-                conda_env=conda_env,
-                input_example=input_example,
-                # resources=all_resources,
-                auth_policy=auth_policy,
+        # End any stale runs before starting to ensure clean state on retry
+        if mlflow.active_run():
+            logger.warning(
+                "Ending stale MLflow run before creating new agent",
+                run_id=mlflow.active_run().info.run_id,
+            )
+            mlflow.end_run()
+
+        try:
+            with mlflow.start_run(run_name=run_name):
+                mlflow.set_tag("type", "agent")
+                mlflow.set_tag("dao_ai", dao_ai_version())
+                logged_agent_info: ModelInfo = mlflow.pyfunc.log_model(
+                    python_model=model_path.as_posix(),
+                    code_paths=code_paths,
+                    model_config=config.model_dump(mode="json", by_alias=True),
+                    name="agent",
+                    conda_env=conda_env,
+                    input_example=input_example,
+                    # resources=all_resources,
+                    auth_policy=auth_policy,
+                )
+        except Exception as e:
+            # Ensure run is ended on failure to prevent stale state on retry
+            if mlflow.active_run():
+                mlflow.end_run(status="FAILED")
+            logger.error(
+                "Failed to log model",
+                run_name=run_name,
+                error=str(e),
             )
+            raise
 
         registered_model_name: str = config.app.registered_model.full_name
 

dao_ai/tools/genie.py

@@ -25,13 +25,19 @@ from pydantic import BaseModel
 from dao_ai.config import (
     AnyVariable,
     CompositeVariableModel,
+    GenieInMemorySemanticCacheParametersModel,
     GenieLRUCacheParametersModel,
     GenieRoomModel,
     GenieSemanticCacheParametersModel,
     value_of,
 )
 from dao_ai.genie import GenieService, GenieServiceBase
-from dao_ai.genie.cache import CacheResult, LRUCacheService, SemanticCacheService
+from dao_ai.genie.cache import (
+    CacheResult,
+    InMemorySemanticCacheService,
+    LRUCacheService,
+    SemanticCacheService,
+)
 from dao_ai.state import AgentState, Context, SessionState
 
 
@@ -67,6 +73,9 @@ def create_genie_tool(
     semantic_cache_parameters: GenieSemanticCacheParametersModel
     | dict[str, Any]
     | None = None,
+    in_memory_semantic_cache_parameters: GenieInMemorySemanticCacheParametersModel
+    | dict[str, Any]
+    | None = None,
 ) -> Callable[..., Command]:
     """
     Create a tool for interacting with Databricks Genie for natural language queries to databases.
@@ -84,7 +93,9 @@ def create_genie_tool(
         truncate_results: Whether to truncate large query results to fit token limits
         lru_cache_parameters: Optional LRU cache configuration for SQL query caching
         semantic_cache_parameters: Optional semantic cache configuration using pg_vector
-            for similarity-based query matching
+            for similarity-based query matching (requires PostgreSQL/Lakebase)
+        in_memory_semantic_cache_parameters: Optional in-memory semantic cache configuration
+            for similarity-based query matching (no database required)
 
     Returns:
         A LangGraph tool that processes natural language queries through Genie
@@ -97,6 +108,7 @@ def create_genie_tool(
         name=name,
         has_lru_cache=lru_cache_parameters is not None,
         has_semantic_cache=semantic_cache_parameters is not None,
+        has_in_memory_semantic_cache=in_memory_semantic_cache_parameters is not None,
     )
 
     if isinstance(genie_room, dict):
@@ -110,6 +122,11 @@ def create_genie_tool(
             **semantic_cache_parameters
         )
 
+    if isinstance(in_memory_semantic_cache_parameters, dict):
+        in_memory_semantic_cache_parameters = GenieInMemorySemanticCacheParametersModel(
+            **in_memory_semantic_cache_parameters
+        )
+
     space_id: AnyVariable = genie_room.space_id or os.environ.get(
         "DATABRICKS_GENIE_SPACE_ID"
     )
@@ -165,7 +182,7 @@ GenieResponse: A response object containing the conversation ID and result from
 
         genie_service: GenieServiceBase = GenieService(genie)
 
-        # Wrap with semantic cache first (checked second due to decorator pattern)
+        # Wrap with semantic cache first (checked second/third due to decorator pattern)
         if semantic_cache_parameters is not None:
             genie_service = SemanticCacheService(
                 impl=genie_service,
@@ -173,6 +190,14 @@ GenieResponse: A response object containing the conversation ID and result from
                 workspace_client=workspace_client,
             ).initialize()
 
+        # Wrap with in-memory semantic cache (alternative to PostgreSQL semantic cache)
+        if in_memory_semantic_cache_parameters is not None:
+            genie_service = InMemorySemanticCacheService(
+                impl=genie_service,
+                parameters=in_memory_semantic_cache_parameters,
+                workspace_client=workspace_client,
+            ).initialize()
+
         # Wrap with LRU cache last (checked first - fast O(1) exact match)
         if lru_cache_parameters is not None:
             genie_service = LRUCacheService(

dao_ai/tools/instructed_retriever.py

@@ -0,0 +1,366 @@
+"""
+Instructed retriever for query decomposition and result fusion.
+
+This module provides functions for decomposing user queries into multiple
+subqueries with metadata filters and merging results using Reciprocal Rank Fusion.
+"""
+
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import mlflow
+import yaml
+from langchain_core.documents import Document
+from langchain_core.language_models import BaseChatModel
+from langchain_core.runnables import Runnable
+from loguru import logger
+from mlflow.entities import SpanType
+from pydantic import BaseModel, ConfigDict, Field
+
+from dao_ai.config import (
+    ColumnInfo,
+    DecomposedQueries,
+    FilterItem,
+    LLMModel,
+    SearchQuery,
+)
+
+# Module-level cache for LLM clients
+_llm_cache: dict[str, BaseChatModel] = {}
+
+# Load prompt template
+_PROMPT_PATH = (
+    Path(__file__).parent.parent / "prompts" / "instructed_retriever_decomposition.yaml"
+)
+
+
+def _load_prompt_template() -> dict[str, Any]:
+    """Load the decomposition prompt template from YAML."""
+    with open(_PROMPT_PATH) as f:
+        return yaml.safe_load(f)
+
+
+def _get_cached_llm(model_config: LLMModel) -> BaseChatModel:
+    """
+    Get or create cached LLM client for decomposition.
+
+    Uses full config as cache key to avoid collisions when same model name
+    has different parameters (temperature, API keys, etc.).
+    """
+    cache_key = model_config.model_dump_json()
+    if cache_key not in _llm_cache:
+        _llm_cache[cache_key] = model_config.as_chat_model()
+        logger.debug(
+            "Created new LLM client for decomposition", model=model_config.name
+        )
+    return _llm_cache[cache_key]
+
+
+def _format_constraints(constraints: list[str] | None) -> str:
+    """Format constraints list for prompt injection."""
+    if not constraints:
+        return "No additional constraints."
+    return "\n".join(f"- {c}" for c in constraints)
+
+
+def _format_examples(examples: list[dict[str, Any]] | None) -> str:
+    """Format few-shot examples for prompt injection.
+
+    Converts dict-style filters from config to FilterItem array format
+    to match the expected JSON schema output.
+    """
+    if not examples:
+        return "No examples provided."
+
+    formatted = []
+    for i, ex in enumerate(examples, 1):
+        query = ex.get("query", "")
+        filters = ex.get("filters", {})
+        # Convert dict to FilterItem array format
+        filter_items = [{"key": k, "value": v} for k, v in filters.items()]
+        formatted.append(
+            f'Example {i}:\n  Query: "{query}"\n  Filters: {json.dumps(filter_items)}'
+        )
+    return "\n".join(formatted)
+
+
+def create_decomposition_schema(
+    columns: list[ColumnInfo] | None = None,
+) -> type[BaseModel]:
+    """Create schema-aware DecomposedQueries model with dynamic descriptions.
+
+    When columns are provided, the column names and valid operators are embedded
+    directly into the JSON schema that with_structured_output sends to the LLM.
+    This improves accuracy by making valid filter keys explicit in the schema.
+
+    Args:
+        columns: List of column metadata for dynamic schema generation
+
+    Returns:
+        A DecomposedQueries-compatible Pydantic model class
+    """
+    if not columns:
+        # Fall back to generic models
+        return DecomposedQueries
+
+    # Build column info with types for the schema description
+    column_info = ", ".join(f"{c.name} ({c.type})" for c in columns)
+
+    # Build operator list from column definitions (union of all column operators)
+    all_operators: set[str] = set()
+    for col in columns:
+        all_operators.update(col.operators)
+    # Remove empty string (equality) and sort for consistent output
+    named_operators = sorted(all_operators - {""})
+    operator_list = ", ".join(named_operators) if named_operators else "equality only"
+
+    # Build valid key examples with operators
+    key_examples: list[str] = []
+    for col in columns[:3]:  # Show examples for first 3 columns
+        key_examples.append(f"'{col.name}'")
+        if "<" in col.operators:
+            key_examples.append(f"'{col.name} <'")
+        if "NOT" in col.operators:
+            key_examples.append(f"'{col.name} NOT'")
+
+    # Create dynamic FilterItem with schema-aware description
+    class SchemaFilterItem(BaseModel):
+        """A metadata filter for vector search with schema-specific columns."""
+
+        model_config = ConfigDict(extra="forbid")
+        key: str = Field(
+            description=(
+                f"Column name with optional operator suffix. "
+                f"Valid columns: {column_info}. "
+                f"Operators: (none) for equality, {operator_list}. "
+                f"Examples: {', '.join(key_examples[:5])}"
+            )
+        )
+        value: Union[str, int, float, bool, list[Union[str, int, float, bool]]] = Field(
+            description="The filter value matching the column type."
+        )
+
+    # Create dynamic SearchQuery using SchemaFilterItem
+    class SchemaSearchQuery(BaseModel):
+        """A search query with schema-aware filters."""
+
+        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[SchemaFilterItem]] = Field(
+            default=None,
+            description=(
+                f"Metadata filters to constrain search results. "
+                f"Valid filter columns: {column_info}. "
+                f"Set to null if no filters apply."
+            ),
+        )
+
+    # Create dynamic DecomposedQueries using SchemaSearchQuery
+    class SchemaDecomposedQueries(BaseModel):
+        """Decomposed search queries with schema-aware filters."""
+
+        model_config = ConfigDict(extra="forbid")
+        queries: list[SchemaSearchQuery] = 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."
+            )
+        )
+
+    return SchemaDecomposedQueries
+
+
+@mlflow.trace(name="decompose_query", span_type=SpanType.LLM)
+def decompose_query(
+    llm: BaseChatModel,
+    query: str,
+    schema_description: str,
+    constraints: list[str] | None = None,
+    max_subqueries: int = 3,
+    examples: list[dict[str, Any]] | None = None,
+    previous_feedback: str | None = None,
+    columns: list[ColumnInfo] | None = None,
+) -> list[SearchQuery]:
+    """
+    Decompose a user query into multiple search queries with filters.
+
+    Uses structured output for reliable parsing and injects current time
+    for resolving relative date references. When columns are provided,
+    schema-aware Pydantic models are used for improved filter accuracy.
+
+    Args:
+        llm: Language model for decomposition
+        query: User's search query
+        schema_description: Column names, types, and valid filter syntax
+        constraints: Default constraints to apply
+        max_subqueries: Maximum number of subqueries to generate
+        examples: Few-shot examples for domain-specific filter translation
+        previous_feedback: Feedback from failed verification (for retry)
+        columns: Structured column info for dynamic schema generation
+
+    Returns:
+        List of SearchQuery objects with text and optional filters
+    """
+    current_time = datetime.now().isoformat()
+
+    # Load and format prompt
+    prompt_config = _load_prompt_template()
+    prompt_template = prompt_config["template"]
+
+    # Add previous feedback section if provided (for retry)
+    feedback_section = ""
+    if previous_feedback:
+        feedback_section = f"\n\n## Previous Attempt Feedback\nThe previous search attempt failed verification: {previous_feedback}\nAdjust your filters to address this feedback."
+
+    prompt = (
+        prompt_template.format(
+            current_time=current_time,
+            schema_description=schema_description,
+            constraints=_format_constraints(constraints),
+            examples=_format_examples(examples),
+            max_subqueries=max_subqueries,
+            query=query,
+        )
+        + feedback_section
+    )
+
+    logger.trace(
+        "Decomposing query",
+        query=query[:100],
+        max_subqueries=max_subqueries,
+        dynamic_schema=columns is not None,
+    )
+
+    # Create schema-aware model when columns are provided
+    DecompositionSchema: type[BaseModel] = create_decomposition_schema(columns)
+
+    # Use LangChain's with_structured_output for automatic strategy selection
+    # (JSON schema vs tool calling based on model capabilities)
+    try:
+        structured_llm: Runnable[str, BaseModel] = llm.with_structured_output(
+            DecompositionSchema
+        )
+        result: BaseModel = structured_llm.invoke(prompt)
+    except Exception as e:
+        logger.warning("Query decomposition failed", error=str(e))
+        raise
+
+    # Extract queries from result (works with both static and dynamic schemas)
+    subqueries: list[SearchQuery] = []
+    for query_obj in result.queries[:max_subqueries]:
+        # Convert dynamic schema objects to SearchQuery for consistent return type
+        filters: list[FilterItem] | None = None
+        if query_obj.filters:
+            filters = [FilterItem(key=f.key, value=f.value) for f in query_obj.filters]
+        subqueries.append(SearchQuery(text=query_obj.text, filters=filters))
+
+    # Log for observability
+    mlflow.set_tag("num_subqueries", len(subqueries))
+    mlflow.log_text(
+        json.dumps([sq.model_dump() for sq in subqueries], indent=2),
+        "decomposition.json",
+    )
+
+    logger.debug(
+        "Query decomposed",
+        num_subqueries=len(subqueries),
+        queries=[sq.text[:50] for sq in subqueries],
+    )
+
+    return subqueries
+
+
+def rrf_merge(
+    results_lists: list[list[Document]],
+    k: int = 60,
+    primary_key: str | None = None,
+) -> list[Document]:
+    """
+    Merge results from multiple queries using Reciprocal Rank Fusion.
+
+    RRF is safer than raw score sorting because Databricks Vector Search
+    scores aren't normalized across query types (HYBRID vs ANN).
+
+    RRF Score = Σ 1 / (k + rank_i) for each result list
+
+    Args:
+        results_lists: List of document lists from different subqueries
+        k: RRF constant (lower values weight top ranks more heavily)
+        primary_key: Metadata key for document identity (for deduplication)
+
+    Returns:
+        Merged and deduplicated documents sorted by RRF score
+    """
+    if not results_lists:
+        return []
+
+    # Filter empty lists first
+    non_empty = [r for r in results_lists if r]
+    if not non_empty:
+        return []
+
+    # Single list optimization (still add RRF scores for consistency)
+    if len(non_empty) == 1:
+        docs_with_scores: list[Document] = []
+        for rank, doc in enumerate(non_empty[0]):
+            rrf_score = 1.0 / (k + rank + 1)
+            docs_with_scores.append(
+                Document(
+                    page_content=doc.page_content,
+                    metadata={**doc.metadata, "rrf_score": rrf_score},
+                )
+            )
+        return docs_with_scores
+
+    # Calculate RRF scores
+    # Key: document identifier, Value: (total_rrf_score, Document)
+    doc_scores: dict[str, tuple[float, Document]] = {}
+
+    def get_doc_id(doc: Document) -> str:
+        """Get unique identifier for document."""
+        if primary_key and primary_key in doc.metadata:
+            return str(doc.metadata[primary_key])
+        # Fallback to content hash
+        return str(hash(doc.page_content))
+
+    for result_list in non_empty:
+        for rank, doc in enumerate(result_list):
+            doc_id = get_doc_id(doc)
+            rrf_score = 1.0 / (k + rank + 1)  # rank is 0-indexed
+
+            if doc_id in doc_scores:
+                # Accumulate RRF score for duplicates
+                existing_score, existing_doc = doc_scores[doc_id]
+                doc_scores[doc_id] = (existing_score + rrf_score, existing_doc)
+            else:
+                doc_scores[doc_id] = (rrf_score, doc)
+
+    # Sort by RRF score descending
+    sorted_docs = sorted(doc_scores.values(), key=lambda x: x[0], reverse=True)
+
+    # Add RRF score to metadata
+    merged_docs: list[Document] = []
+    for rrf_score, doc in sorted_docs:
+        merged_doc = Document(
+            page_content=doc.page_content,
+            metadata={**doc.metadata, "rrf_score": rrf_score},
+        )
+        merged_docs.append(merged_doc)
+
+    logger.debug(
+        "RRF merge complete",
+        input_lists=len(results_lists),
+        total_docs=sum(len(r) for r in results_lists),
+        unique_docs=len(merged_docs),
+    )
+
+    return merged_docs