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

dao_ai/tools/genie.py

@@ -6,7 +6,7 @@ interact with Databricks Genie.
 
 For the core Genie service and cache implementations, see:
 - dao_ai.genie: GenieService, GenieServiceBase
-- dao_ai.genie.cache: LRUCacheService, SemanticCacheService
+- dao_ai.genie.cache: LRUCacheService, PostgresContextAwareGenieService, InMemoryContextAwareGenieService
 """
 
 import json
@@ -25,13 +25,19 @@ from pydantic import BaseModel
 from dao_ai.config import (
     AnyVariable,
     CompositeVariableModel,
+    GenieContextAwareCacheParametersModel,
+    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,
+    InMemoryContextAwareGenieService,
+    LRUCacheService,
+    PostgresContextAwareGenieService,
+)
 from dao_ai.state import AgentState, Context, SessionState
 
 
@@ -64,7 +70,10 @@ def create_genie_tool(
     persist_conversation: bool = True,
     truncate_results: bool = False,
     lru_cache_parameters: GenieLRUCacheParametersModel | dict[str, Any] | None = None,
-    semantic_cache_parameters: GenieSemanticCacheParametersModel
+    semantic_cache_parameters: GenieContextAwareCacheParametersModel
+    | dict[str, Any]
+    | None = None,
+    in_memory_semantic_cache_parameters: GenieInMemorySemanticCacheParametersModel
     | dict[str, Any]
     | None = None,
 ) -> Callable[..., Command]:
@@ -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):
@@ -106,10 +118,15 @@ def create_genie_tool(
         lru_cache_parameters = GenieLRUCacheParametersModel(**lru_cache_parameters)
 
     if isinstance(semantic_cache_parameters, dict):
-        semantic_cache_parameters = GenieSemanticCacheParametersModel(
+        semantic_cache_parameters = GenieContextAwareCacheParametersModel(
             **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"
     )
@@ -139,29 +156,61 @@ Returns:
 GenieResponse: A response object containing the conversation ID and result from Genie."""
     tool_description = tool_description + function_docs
 
-    genie: Genie = Genie(
-        space_id=space_id,
-        client=genie_room.workspace_client,
-        truncate_results=truncate_results,
-    )
+    # Cache for genie service - created lazily on first call
+    # This allows us to use workspace_client_from with runtime context for OBO
+    _cached_genie_service: GenieServiceBase | None = None
+
+    def _get_genie_service(context: Context | None) -> GenieServiceBase:
+        """Get or create the Genie service, using context for OBO auth if available."""
+        nonlocal _cached_genie_service
+
+        # Use cached service if available (for non-OBO or after first call)
+        # For OBO, we need fresh workspace client each time to use the user's token
+        if _cached_genie_service is not None and not genie_room.on_behalf_of_user:
+            return _cached_genie_service
 
-    genie_service: GenieServiceBase = GenieService(genie)
-
-    # Wrap with semantic cache first (checked second due to decorator pattern)
-    if semantic_cache_parameters is not None:
-        genie_service = SemanticCacheService(
-            impl=genie_service,
-            parameters=semantic_cache_parameters,
-            workspace_client=genie_room.workspace_client,  # Pass workspace client for conversation history
-        ).initialize()  # Eagerly initialize to fail fast and create table
-
-    # Wrap with LRU cache last (checked first - fast O(1) exact match)
-    if lru_cache_parameters is not None:
-        genie_service = LRUCacheService(
-            impl=genie_service,
-            parameters=lru_cache_parameters,
+        # Get workspace client using context for OBO support
+        from databricks.sdk import WorkspaceClient
+
+        workspace_client: WorkspaceClient = genie_room.workspace_client_from(context)
+
+        genie: Genie = Genie(
+            space_id=space_id,
+            client=workspace_client,
+            truncate_results=truncate_results,
         )
 
+        genie_service: GenieServiceBase = GenieService(genie)
+
+        # Wrap with context-aware cache first (checked second/third due to decorator pattern)
+        if semantic_cache_parameters is not None:
+            genie_service = PostgresContextAwareGenieService(
+                impl=genie_service,
+                parameters=semantic_cache_parameters,
+                workspace_client=workspace_client,
+            ).initialize()
+
+        # Wrap with in-memory context-aware cache (alternative to PostgreSQL context-aware cache)
+        if in_memory_semantic_cache_parameters is not None:
+            genie_service = InMemoryContextAwareGenieService(
+                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(
+                impl=genie_service,
+                parameters=lru_cache_parameters,
+            )
+
+        # Cache for non-OBO scenarios
+        if not genie_room.on_behalf_of_user:
+            _cached_genie_service = genie_service
+
+        return genie_service
+
     @tool(
         name_or_callable=tool_name,
         description=tool_description,
@@ -177,6 +226,10 @@ GenieResponse: A response object containing the conversation ID and result from
         # Access state through runtime
         state: AgentState = runtime.state
         tool_call_id: str = runtime.tool_call_id
+        context: Context | None = runtime.context
+
+        # Get genie service with OBO support via context
+        genie_service: GenieServiceBase = _get_genie_service(context)
 
         # Ensure space_id is a string for state keys
         space_id_str: str = str(space_id)
@@ -194,6 +247,14 @@ GenieResponse: A response object containing the conversation ID and result from
             conversation_id=existing_conversation_id,
         )
 
+        # Log the prompt being sent to Genie
+        logger.trace(
+            "Sending prompt to Genie",
+            space_id=space_id_str,
+            conversation_id=existing_conversation_id,
+            prompt=question[:500] + "..." if len(question) > 500 else question,
+        )
+
         # Call ask_question which always returns CacheResult with cache metadata
         cache_result: CacheResult = genie_service.ask_question(
             question, conversation_id=existing_conversation_id
@@ -211,6 +272,22 @@ GenieResponse: A response object containing the conversation ID and result from
             cache_key=cache_key,
         )
 
+        # Log truncated response for debugging
+        result_preview: str = str(genie_response.result)
+        if len(result_preview) > 500:
+            result_preview = result_preview[:500] + "..."
+        logger.trace(
+            "Genie response content",
+            question=question[:100] + "..." if len(question) > 100 else question,
+            query=genie_response.query,
+            description=(
+                genie_response.description[:200] + "..."
+                if genie_response.description and len(genie_response.description) > 200
+                else genie_response.description
+            ),
+            result_preview=result_preview,
+        )
+
         # Update session state with cache information
         if persist_conversation:
             session.genie.update_space(

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