dao-ai 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

dao_ai/tools/vector_search.py

@@ -2,317 +2,391 @@
 Vector search tool for retrieving documents from Databricks Vector Search.
 
 This module provides a tool factory for creating semantic search tools
-using ToolRuntime[Context, AgentState] for type-safe runtime access.
+with dynamic filter schemas based on table columns and FlashRank reranking support.
 """
 
+import json
 import os
-from typing import Any, Callable, List, Optional, Sequence
+from typing import Any, Optional
 
 import mlflow
+from databricks.sdk import WorkspaceClient
 from databricks.vector_search.reranker import DatabricksReranker
-from databricks_ai_bridge.vector_search_retriever_tool import (
-    FilterItem,
-    VectorSearchRetrieverToolInput,
-)
-from databricks_langchain.vectorstores import DatabricksVectorSearch
+from databricks_langchain import DatabricksVectorSearch
 from flashrank import Ranker, RerankRequest
-from langchain.tools import ToolRuntime, tool
 from langchain_core.documents import Document
+from langchain_core.tools import StructuredTool
 from loguru import logger
 from mlflow.entities import SpanType
+from pydantic import BaseModel, ConfigDict, Field, create_model
 
 from dao_ai.config import (
     RerankParametersModel,
     RetrieverModel,
+    SearchParametersModel,
     VectorStoreModel,
+    value_of,
 )
-from dao_ai.state import AgentState, Context
 from dao_ai.utils import normalize_host
 
+# Create FilterItem model at module level so it can be used in type hints
+FilterItem = create_model(
+    "FilterItem",
+    key=(
+        str,
+        Field(
+            description="The filter key, which includes the column name and can include operators like 'NOT', '<', '>=', 'LIKE', 'OR'"
+        ),
+    ),
+    value=(
+        Any,
+        Field(
+            description="The filter value, which can be a single value or an array of values"
+        ),
+    ),
+    __config__=ConfigDict(extra="forbid"),
+)
 
-def create_vector_search_tool(
-    retriever: RetrieverModel | dict[str, Any],
-    name: Optional[str] = None,
-    description: Optional[str] = None,
-) -> Callable[..., list[dict[str, Any]]]:
-    """
-    Create a Vector Search tool for retrieving documents from a Databricks Vector Search index.
 
-    This function creates a tool that enables semantic search over product information,
-    documentation, or other content using the @tool decorator pattern. It supports optional
-    reranking of results using FlashRank for improved relevance.
+def _create_dynamic_input_schema(
+    index_name: str, workspace_client: WorkspaceClient
+) -> type[BaseModel]:
+    """
+    Create dynamic input schema with column information from the table.
 
     Args:
-        retriever: Configuration details for the vector search retriever, including:
-            - name: Name of the tool
-            - description: Description of the tool's purpose
-            - primary_key: Primary key column for the vector store
-            - text_column: Text column used for vector search
-            - doc_uri: URI for documentation or additional context
-            - vector_store: Dictionary with 'endpoint_name' and 'index' for vector search
-            - columns: List of columns to retrieve from the vector store
-            - search_parameters: Additional parameters for customizing the search behavior
-            - rerank: Optional rerank configuration for result reranking
-        name: Optional custom name for the tool
-        description: Optional custom description for the tool
+        index_name: Full name of the vector search index
+        workspace_client: Workspace client to query table metadata
 
     Returns:
-        A LangChain tool that performs vector search with optional reranking
+        Pydantic model class for tool input
     """
 
-    if isinstance(retriever, dict):
-        retriever = RetrieverModel(**retriever)
-
-    vector_store_config: VectorStoreModel = retriever.vector_store
-
-    # Index is required for vector search
-    if vector_store_config.index is None:
-        raise ValueError("vector_store.index is required for vector search")
-
-    index_name: str = vector_store_config.index.full_name
-    columns: Sequence[str] = retriever.columns or []
-    search_parameters: dict[str, Any] = retriever.search_parameters.model_dump()
-    primary_key: str = vector_store_config.primary_key or ""
-    doc_uri: str = vector_store_config.doc_uri or ""
-    text_column: str = vector_store_config.embedding_source_column
-
-    # Extract reranker configuration
-    reranker_config: Optional[RerankParametersModel] = retriever.rerank
-
-    # Initialize FlashRank ranker once if reranking is enabled
-    # This is expensive (loads model weights), so we do it once and reuse across invocations
-    ranker: Optional[Ranker] = None
-    if reranker_config:
-        logger.debug(
-            f"Creating vector search tool with reranking: '{name}' "
-            f"(model: {reranker_config.model}, top_n: {reranker_config.top_n or 'auto'})"
-        )
-        try:
-            ranker = Ranker(
-                model_name=reranker_config.model, cache_dir=reranker_config.cache_dir
-            )
-            logger.info(
-                f"FlashRank ranker initialized successfully (model: {reranker_config.model})"
-            )
-        except Exception as e:
-            logger.warning(
-                f"Failed to initialize FlashRank ranker during tool creation: {e}. "
-                "Reranking will be disabled for this tool."
-            )
-            # Set reranker_config to None so we don't attempt reranking
-            reranker_config = None
-    else:
+    # Try to get column information
+    column_descriptions = []
+    try:
+        table_info = workspace_client.tables.get(full_name=index_name)
+        for column_info in table_info.columns:
+            name = column_info.name
+            col_type = column_info.type_name.name
+            if not name.startswith("__"):
+                column_descriptions.append(f"{name} ({col_type})")
+    except Exception:
         logger.debug(
-            f"Creating vector search tool without reranking: '{name}' (standard similarity search only)"
+            "Could not retrieve column information for dynamic schema",
+            index=index_name,
         )
 
-    # Initialize the vector store
-    # Note: text_column is only required for self-managed embeddings
-    # For Databricks-managed embeddings, it's automatically determined from the index
+    # Build filter description matching VectorSearchRetrieverTool format
+    filter_description = (
+        "Optional filters to refine vector search results as an array of key-value pairs. "
+        "IMPORTANT: If unsure about filter values, try searching WITHOUT filters first to get broad results, "
+        "then optionally add filters to narrow down if needed. This ensures you don't miss relevant results due to incorrect filter values. "
+    )
 
-    # Build client_args for VectorSearchClient from environment variables
-    # This is needed because during MLflow model validation, credentials must be
-    # explicitly passed to VectorSearchClient via client_args.
-    # The workspace_client parameter in DatabricksVectorSearch is only used to detect
-    # model serving mode - it doesn't pass credentials to VectorSearchClient.
-    client_args: dict[str, Any] = {}
-    databricks_host = normalize_host(os.environ.get("DATABRICKS_HOST"))
-    if databricks_host:
-        client_args["workspace_url"] = databricks_host
-    if os.environ.get("DATABRICKS_TOKEN"):
-        client_args["personal_access_token"] = os.environ.get("DATABRICKS_TOKEN")
-    if os.environ.get("DATABRICKS_CLIENT_ID"):
-        client_args["service_principal_client_id"] = os.environ.get(
-            "DATABRICKS_CLIENT_ID"
-        )
-    if os.environ.get("DATABRICKS_CLIENT_SECRET"):
-        client_args["service_principal_client_secret"] = os.environ.get(
-            "DATABRICKS_CLIENT_SECRET"
+    if column_descriptions:
+        filter_description += (
+            f"Available columns for filtering: {', '.join(column_descriptions)}. "
         )
 
-    logger.debug(
-        f"Creating DatabricksVectorSearch with client_args keys: {list(client_args.keys())}"
+    filter_description += (
+        "Supports the following operators:\n\n"
+        '- Inclusion: [{"key": "column", "value": value}] or [{"key": "column", "value": [value1, value2]}] (matches if the column equals any of the provided values)\n'
+        '- Exclusion: [{"key": "column NOT", "value": value}]\n'
+        '- Comparisons: [{"key": "column <", "value": value}], [{"key": "column >=", "value": value}], etc.\n'
+        '- Pattern match: [{"key": "column LIKE", "value": "word"}] (matches full tokens separated by whitespace)\n'
+        '- OR logic: [{"key": "column1 OR column2", "value": [value1, value2]}] '
+        "(matches if column1 equals value1 or column2 equals value2; matches are position-specific)\n\n"
+        "Examples:\n"
+        '- Filter by category: [{"key": "category", "value": "electronics"}]\n'
+        '- Filter by price range: [{"key": "price >=", "value": 100}, {"key": "price <", "value": 500}]\n'
+        '- Exclude specific status: [{"key": "status NOT", "value": "archived"}]\n'
+        '- Pattern matching: [{"key": "description LIKE", "value": "wireless"}]'
     )
 
-    # Pass both workspace_client (for model serving detection) and client_args (for credentials)
-    vector_store: DatabricksVectorSearch = DatabricksVectorSearch(
-        index_name=index_name,
-        text_column=None,  # Let DatabricksVectorSearch determine this from the index
-        columns=columns,
-        include_score=True,
-        workspace_client=vector_store_config.workspace_client,
-        client_args=client_args if client_args else None,
+    # Create the input model
+    VectorSearchInput = create_model(
+        "VectorSearchInput",
+        query=(
+            str,
+            Field(description="The search query string to find relevant documents"),
+        ),
+        filters=(
+            Optional[list[FilterItem]],
+            Field(default=None, description=filter_description),
+        ),
+        __config__=ConfigDict(extra="forbid"),
     )
 
-    # Register the retriever schema with MLflow for model serving integration
-    mlflow.models.set_retriever_schema(
-        name=name or "retriever",
-        primary_key=primary_key,
-        text_column=text_column,
-        doc_uri=doc_uri,
-        other_columns=list(columns),
+    return VectorSearchInput
+
+
+@mlflow.trace(name="rerank_documents", span_type=SpanType.RETRIEVER)
+def _rerank_documents(
+    query: str,
+    documents: list[Document],
+    ranker: Ranker,
+    rerank_config: RerankParametersModel,
+) -> list[Document]:
+    """
+    Rerank documents using FlashRank cross-encoder model.
+
+    Args:
+        query: The search query string
+        documents: List of documents to rerank
+        ranker: The FlashRank Ranker instance
+        rerank_config: Reranking configuration
+
+    Returns:
+        Reranked list of documents with reranker_score in metadata
+    """
+    logger.trace(
+        "Starting reranking",
+        documents_count=len(documents),
+        model=rerank_config.model,
     )
 
-    # Helper function to perform vector similarity search
-    @mlflow.trace(name="find_documents", span_type=SpanType.RETRIEVER)
-    def _find_documents(
-        query: str, filters: Optional[List[FilterItem]] = None
-    ) -> List[Document]:
-        """Perform vector similarity search."""
-        # Convert filters to dict format
-        filters_dict: dict[str, Any] = {}
-        if filters:
-            for item in filters:
-                item_dict = dict(item)
-                filters_dict[item_dict["key"]] = item_dict["value"]
+    # Prepare passages for reranking
+    passages: list[dict[str, Any]] = [
+        {"text": doc.page_content, "meta": doc.metadata} for doc in documents
+    ]
 
-        # Merge with any configured filters
-        combined_filters: dict[str, Any] = {
-            **filters_dict,
-            **search_parameters.get("filters", {}),
-        }
+    # Create reranking request
+    rerank_request: RerankRequest = RerankRequest(query=query, passages=passages)
 
-        # Perform similarity search
-        num_results: int = search_parameters.get("num_results", 10)
-        query_type: str = search_parameters.get("query_type", "ANN")
+    # Perform reranking
+    results: list[dict[str, Any]] = ranker.rerank(rerank_request)
 
-        logger.debug(
-            f"Performing vector search: query='{query[:50]}...', k={num_results}, filters={combined_filters}"
+    # Apply top_n filtering
+    top_n: int = rerank_config.top_n or len(documents)
+    results = results[:top_n]
+    logger.debug("Reranking complete", top_n=top_n, candidates_count=len(documents))
+
+    # Convert back to Document objects with reranking scores
+    reranked_docs: list[Document] = []
+    for result in results:
+        orig_doc: Optional[Document] = next(
+            (doc for doc in documents if doc.page_content == result["text"]), None
         )
+        if orig_doc:
+            reranked_doc: Document = Document(
+                page_content=orig_doc.page_content,
+                metadata={
+                    **orig_doc.metadata,
+                    "reranker_score": result["score"],
+                },
+            )
+            reranked_docs.append(reranked_doc)
 
-        # Build similarity search kwargs
-        search_kwargs = {
-            "query": query,
-            "k": num_results,
-            "filter": combined_filters if combined_filters else None,
-            "query_type": query_type,
-        }
+    logger.debug(
+        "Documents reranked",
+        input_count=len(documents),
+        output_count=len(reranked_docs),
+        model=rerank_config.model,
+    )
 
-        # Add DatabricksReranker if configured with columns
-        if reranker_config and reranker_config.columns:
-            search_kwargs["reranker"] = DatabricksReranker(
-                columns_to_rerank=reranker_config.columns
-            )
+    return reranked_docs
 
-        documents: List[Document] = vector_store.similarity_search(**search_kwargs)
 
-        logger.debug(f"Retrieved {len(documents)} documents from vector search")
-        return documents
+def create_vector_search_tool(
+    retriever: Optional[RetrieverModel | dict[str, Any]] = None,
+    vector_store: Optional[VectorStoreModel | dict[str, Any]] = None,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+) -> StructuredTool:
+    """
+    Create a Vector Search tool with dynamic schema and optional reranking.
 
-    # Helper function to rerank documents
-    @mlflow.trace(name="rerank_documents", span_type=SpanType.RETRIEVER)
-    def _rerank_documents(query: str, documents: List[Document]) -> List[Document]:
-        """Rerank documents using FlashRank.
+    Args:
+        retriever: Full retriever configuration with search parameters and reranking
+        vector_store: Direct vector store reference (uses default search parameters)
+        name: Optional custom name for the tool
+        description: Optional custom description for the tool
 
-        Uses the ranker instance initialized at tool creation time (captured in closure).
-        This avoids expensive model loading on every invocation.
-        """
-        if not reranker_config or ranker is None:
-            return documents
+    Returns:
+        A LangChain StructuredTool with proper schema (additionalProperties: false)
+    """
 
-        logger.debug(
-            f"Starting reranking for {len(documents)} documents using model '{reranker_config.model}'"
+    # Validate mutually exclusive parameters
+    if retriever is None and vector_store is None:
+        raise ValueError("Must provide either 'retriever' or 'vector_store' parameter")
+    if retriever is not None and vector_store is not None:
+        raise ValueError(
+            "Cannot provide both 'retriever' and 'vector_store' parameters"
         )
 
-        # Prepare passages for reranking
-        passages: List[dict[str, Any]] = [
-            {"text": doc.page_content, "meta": doc.metadata} for doc in documents
-        ]
+    # Handle vector_store parameter
+    if vector_store is not None:
+        if isinstance(vector_store, dict):
+            vector_store = VectorStoreModel(**vector_store)
+        retriever = RetrieverModel(vector_store=vector_store)
+    else:
+        if isinstance(retriever, dict):
+            retriever = RetrieverModel(**retriever)
+
+    vector_store: VectorStoreModel = retriever.vector_store
 
-        # Create reranking request
-        rerank_request: RerankRequest = RerankRequest(query=query, passages=passages)
+    # Index is required
+    if vector_store.index is None:
+        raise ValueError("vector_store.index is required for vector search")
 
-        # Perform reranking
-        logger.debug(f"Reranking {len(passages)} passages for query: '{query[:50]}...'")
-        results: List[dict[str, Any]] = ranker.rerank(rerank_request)
+    index_name: str = vector_store.index.full_name
+    columns: list[str] = list(retriever.columns or [])
+    search_parameters: SearchParametersModel = retriever.search_parameters
+    rerank_config: Optional[RerankParametersModel] = retriever.rerank
 
-        # Apply top_n filtering
-        top_n: int = reranker_config.top_n or len(documents)
-        results = results[:top_n]
+    # Initialize FlashRank ranker if configured
+    ranker: Optional[Ranker] = None
+    if rerank_config and rerank_config.model:
         logger.debug(
-            f"Reranking complete. Filtered to top {top_n} results from {len(documents)} candidates"
+            "Initializing FlashRank ranker",
+            model=rerank_config.model,
+            top_n=rerank_config.top_n or "auto",
         )
+        try:
+            cache_dir = os.path.expanduser(rerank_config.cache_dir)
+            ranker = Ranker(model_name=rerank_config.model, cache_dir=cache_dir)
+            logger.success("FlashRank ranker initialized", model=rerank_config.model)
+        except Exception as e:
+            logger.warning("Failed to initialize FlashRank ranker", error=str(e))
+            rerank_config = None
 
-        # Convert back to Document objects with reranking scores
-        reranked_docs: List[Document] = []
-        for result in results:
-            # Find original document by matching text
-            orig_doc: Optional[Document] = next(
-                (doc for doc in documents if doc.page_content == result["text"]), None
-            )
-            if orig_doc:
-                # Add reranking score to metadata
-                reranked_doc: Document = Document(
-                    page_content=orig_doc.page_content,
-                    metadata={
-                        **orig_doc.metadata,
-                        "reranker_score": result["score"],
-                    },
-                )
-                reranked_docs.append(reranked_doc)
+    # Build client_args for VectorSearchClient
+    # Use getattr to safely access attributes that may not exist (e.g., in mocks)
+    client_args: dict[str, Any] = {}
+    has_explicit_auth = any(
+        [
+            os.environ.get("DATABRICKS_TOKEN"),
+            os.environ.get("DATABRICKS_CLIENT_ID"),
+            getattr(vector_store, "pat", None),
+            getattr(vector_store, "client_id", None),
+            getattr(vector_store, "on_behalf_of_user", None),
+        ]
+    )
 
-        logger.debug(
-            f"Reranked {len(documents)} documents → {len(reranked_docs)} results "
-            f"(model: {reranker_config.model}, top score: {reranked_docs[0].metadata.get('reranker_score', 0):.4f})"
-            if reranked_docs
-            else f"Reranking completed with {len(reranked_docs)} results"
-        )
+    if has_explicit_auth:
+        databricks_host = os.environ.get("DATABRICKS_HOST")
+        if (
+            not databricks_host
+            and getattr(vector_store, "_workspace_client", None) is not None
+        ):
+            databricks_host = vector_store.workspace_client.config.host
+        if databricks_host:
+            client_args["workspace_url"] = normalize_host(databricks_host)
+
+        token = os.environ.get("DATABRICKS_TOKEN")
+        if not token and getattr(vector_store, "pat", None):
+            token = value_of(vector_store.pat)
+        if token:
+            client_args["personal_access_token"] = token
+
+        client_id = os.environ.get("DATABRICKS_CLIENT_ID")
+        if not client_id and getattr(vector_store, "client_id", None):
+            client_id = value_of(vector_store.client_id)
+        if client_id:
+            client_args["service_principal_client_id"] = client_id
+
+        client_secret = os.environ.get("DATABRICKS_CLIENT_SECRET")
+        if not client_secret and getattr(vector_store, "client_secret", None):
+            client_secret = value_of(vector_store.client_secret)
+        if client_secret:
+            client_args["service_principal_client_secret"] = client_secret
 
-        return reranked_docs
+    logger.debug(
+        "Creating vector search tool",
+        name=name,
+        index=index_name,
+        client_args_keys=list(client_args.keys()) if client_args else [],
+    )
 
-    # Create the main vector search tool using @tool decorator
-    # Uses ToolRuntime[Context, AgentState] for type-safe runtime access
-    @tool(
-        name_or_callable=name or index_name,
-        description=description or "Search for documents using vector similarity",
-        args_schema=VectorSearchRetrieverToolInput,
+    # Create DatabricksVectorSearch
+    # Note: text_column should be None for Databricks-managed embeddings
+    # (it's automatically determined from the index)
+    vector_search: DatabricksVectorSearch = DatabricksVectorSearch(
+        index_name=index_name,
+        text_column=None,
+        columns=columns,
+        workspace_client=vector_store.workspace_client,
+        client_args=client_args if client_args else None,
+        primary_key=vector_store.primary_key,
+        doc_uri=vector_store.doc_uri,
+        include_score=True,
+        reranker=(
+            DatabricksReranker(columns_to_rerank=rerank_config.columns)
+            if rerank_config and rerank_config.columns
+            else None
+        ),
     )
-    def vector_search_tool(
-        query: str,
-        filters: Optional[List[FilterItem]] = None,
-        runtime: ToolRuntime[Context, AgentState] = None,
-    ) -> list[dict[str, Any]]:
-        """
-        Search for documents using vector similarity with optional reranking.
 
-        This tool performs a two-stage retrieval process:
-        1. Vector similarity search to find candidate documents
-        2. Optional reranking using cross-encoder model for improved relevance
+    # Create dynamic input schema
+    input_schema: type[BaseModel] = _create_dynamic_input_schema(
+        index_name, vector_store.workspace_client
+    )
 
-        Both stages are traced in MLflow for observability.
+    # Define the tool function
+    def vector_search_func(
+        query: str, filters: Optional[list[FilterItem]] = None
+    ) -> str:
+        """Search for relevant documents using vector similarity."""
+        # Convert FilterItem Pydantic models to dict format for DatabricksVectorSearch
+        filters_dict: dict[str, Any] = {}
+        if filters:
+            for item in filters:
+                filters_dict[item.key] = item.value
 
-        Uses ToolRuntime[Context, AgentState] for type-safe runtime access.
+        # Merge with configured filters
+        combined_filters: dict[str, Any] = {
+            **filters_dict,
+            **(search_parameters.filters or {}),
+        }
 
-        Returns:
-            List of serialized documents with page_content and metadata
-        """
-        logger.debug(
-            f"Vector search tool called: query='{query[:50]}...', reranking={reranker_config is not None}"
+        # Perform vector search
+        logger.trace("Performing vector search", query_preview=query[:50])
+        documents: list[Document] = vector_search.similarity_search(
+            query=query,
+            k=search_parameters.num_results or 5,
+            filter=combined_filters if combined_filters else None,
+            query_type=search_parameters.query_type or "ANN",
         )
 
-        # Step 1: Perform vector similarity search
-        documents: List[Document] = _find_documents(query, filters)
-
-        # Step 2: If reranking is enabled, rerank the documents
-        if reranker_config:
-            logger.debug(
-                f"Reranking enabled (model: '{reranker_config.model}', top_n: {reranker_config.top_n or 'all'})"
+        # Apply FlashRank reranking if configured
+        if ranker and rerank_config:
+            logger.debug("Applying FlashRank reranking")
+            documents = _rerank_documents(query, documents, ranker, rerank_config)
+
+        # Serialize documents to JSON format for LLM consumption
+        # Convert Document objects to dicts with page_content and metadata
+        # Need to handle numpy types in metadata (e.g., float32, int64)
+        serialized_docs: list[dict[str, Any]] = []
+        for doc in documents:
+            doc: Document
+            # Convert metadata values to JSON-serializable types
+            metadata_serializable: dict[str, Any] = {}
+            for key, value in doc.metadata.items():
+                # Handle numpy types
+                if hasattr(value, "item"):  # numpy scalar
+                    metadata_serializable[key] = value.item()
+                else:
+                    metadata_serializable[key] = value
+
+            serialized_docs.append(
+                {
+                    "page_content": doc.page_content,
+                    "metadata": metadata_serializable,
+                }
             )
-            documents = _rerank_documents(query, documents)
-            logger.debug(f"Returning {len(documents)} reranked documents")
-        else:
-            logger.debug("Reranking disabled, returning original vector search results")
-
-        # Return Command with ToolMessage containing the documents
-        # Serialize documents to dicts for proper ToolMessage handling
-        serialized_docs: list[dict[str, Any]] = [
-            {
-                "page_content": doc.page_content,
-                "metadata": doc.metadata,
-            }
-            for doc in documents
-        ]
 
-        return serialized_docs
+        # Return as JSON string
+        return json.dumps(serialized_docs)
+
+    # Create the StructuredTool
+    tool: StructuredTool = StructuredTool.from_function(
+        func=vector_search_func,
+        name=name or f"vector_search_{vector_store.index.name}",
+        description=description or f"Search documents in {index_name}",
+        args_schema=input_schema,
+    )
+
+    logger.success("Vector search tool created", name=tool.name, index=index_name)
 
-    return vector_search_tool
+    return tool