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

dao_ai/tools/vector_search.py

@@ -1,282 +1,392 @@
-from typing import Annotated, Any, Callable, List, Optional, Sequence
+"""
+Vector search tool for retrieving documents from Databricks Vector Search.
+
+This module provides a tool factory for creating semantic search tools
+with dynamic filter schemas based on table columns and FlashRank reranking support.
+"""
+
+import json
+import os
+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_core.documents import Document
-from langchain_core.tools import InjectedToolCallId, tool
-from langgraph.prebuilt import InjectedState
+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.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_dynamic_input_schema(
+    index_name: str, workspace_client: WorkspaceClient
+) -> type[BaseModel]:
+    """
+    Create dynamic input schema with column information from the table.
+
+    Args:
+        index_name: Full name of the vector search index
+        workspace_client: Workspace client to query table metadata
+
+    Returns:
+        Pydantic model class for tool input
+    """
+
+    # 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(
+            "Could not retrieve column information for dynamic schema",
+            index=index_name,
+        )
+
+    # 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. "
+    )
+
+    if column_descriptions:
+        filter_description += (
+            f"Available columns for filtering: {', '.join(column_descriptions)}. "
+        )
+
+    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"}]'
+    )
+
+    # 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"),
+    )
+
+    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,
+    )
+
+    # Prepare passages for reranking
+    passages: list[dict[str, Any]] = [
+        {"text": doc.page_content, "meta": doc.metadata} for doc in documents
+    ]
+
+    # Create reranking request
+    rerank_request: RerankRequest = RerankRequest(query=query, passages=passages)
+
+    # Perform reranking
+    results: list[dict[str, Any]] = ranker.rerank(rerank_request)
+
+    # 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)
+
+    logger.debug(
+        "Documents reranked",
+        input_count=len(documents),
+        output_count=len(reranked_docs),
+        model=rerank_config.model,
+    )
+
+    return reranked_docs
 
 
 def create_vector_search_tool(
-    retriever: RetrieverModel | dict[str, Any],
+    retriever: Optional[RetrieverModel | dict[str, Any]] = None,
+    vector_store: Optional[VectorStoreModel | dict[str, Any]] = None,
     name: Optional[str] = None,
     description: Optional[str] = None,
-) -> Callable:
+) -> StructuredTool:
     """
-    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.
+    Create a Vector Search tool with dynamic schema and optional reranking.
 
     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
+        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
 
     Returns:
-        A LangChain tool that performs vector search with optional reranking
+        A LangChain StructuredTool with proper schema (additionalProperties: false)
     """
 
-    if isinstance(retriever, dict):
-        retriever = RetrieverModel(**retriever)
+    # 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"
+        )
 
-    vector_store_config: VectorStoreModel = retriever.vector_store
+    # 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)
 
-    # Index is required for vector search
-    if vector_store_config.index is None:
-        raise ValueError("vector_store.index is required for vector search")
+    vector_store: VectorStoreModel = retriever.vector_store
 
-    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
+    # Index is required
+    if vector_store.index is None:
+        raise ValueError("vector_store.index is required for vector search")
 
-    # Extract reranker configuration
-    reranker_config: Optional[RerankParametersModel] = retriever.rerank
+    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
 
-    # Initialize FlashRank ranker once if reranking is enabled
-    # This is expensive (loads model weights), so we do it once and reuse across invocations
+    # Initialize FlashRank ranker if configured
     ranker: Optional[Ranker] = None
-    if reranker_config:
+    if rerank_config and rerank_config.model:
         logger.debug(
-            f"Creating vector search tool with reranking: '{name}' "
-            f"(model: {reranker_config.model}, top_n: {reranker_config.top_n or 'auto'})"
+            "Initializing FlashRank ranker",
+            model=rerank_config.model,
+            top_n=rerank_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})"
-            )
+            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(
-                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:
-        logger.debug(
-            f"Creating vector search tool without reranking: '{name}' (standard similarity search only)"
-        )
+            logger.warning("Failed to initialize FlashRank ranker", error=str(e))
+            rerank_config = None
+
+    # 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),
+        ]
+    )
+
+    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
+
+    logger.debug(
+        "Creating vector search tool",
+        name=name,
+        index=index_name,
+        client_args_keys=list(client_args.keys()) if client_args else [],
+    )
 
-    # 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
-    vector_store: DatabricksVectorSearch = DatabricksVectorSearch(
+    # 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,  # Let DatabricksVectorSearch determine this from the index
+        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,
-        workspace_client=vector_store_config.workspace_client,
+        reranker=(
+            DatabricksReranker(columns_to_rerank=rerank_config.columns)
+            if rerank_config and rerank_config.columns
+            else None
+        ),
     )
 
-    # 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),
+    # Create dynamic input schema
+    input_schema: type[BaseModel] = _create_dynamic_input_schema(
+        index_name, vector_store.workspace_client
     )
 
-    # 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
+    # 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:
-                item_dict = dict(item)
-                filters_dict[item_dict["key"]] = item_dict["value"]
+                filters_dict[item.key] = item.value
 
-        # Merge with any configured filters
+        # Merge with configured filters
         combined_filters: dict[str, Any] = {
             **filters_dict,
-            **search_parameters.get("filters", {}),
+            **(search_parameters.filters or {}),
         }
 
-        # Perform similarity search
-        num_results: int = search_parameters.get("num_results", 10)
-        query_type: str = search_parameters.get("query_type", "ANN")
-
-        logger.debug(
-            f"Performing vector search: query='{query[:50]}...', k={num_results}, filters={combined_filters}"
+        # 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",
         )
 
-        # Build similarity search kwargs
-        search_kwargs = {
-            "query": query,
-            "k": num_results,
-            "filter": combined_filters if combined_filters else None,
-            "query_type": query_type,
-        }
-
-        # Add DatabricksReranker if configured with columns
-        if reranker_config and reranker_config.columns:
-            search_kwargs["reranker"] = DatabricksReranker(
-                columns_to_rerank=reranker_config.columns
+        # 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: List[Document] = vector_store.similarity_search(**search_kwargs)
-
-        logger.debug(f"Retrieved {len(documents)} documents from vector search")
-        return documents
-
-    # 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.
+        # Return as JSON string
+        return json.dumps(serialized_docs)
 
-        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
-
-        logger.debug(
-            f"Starting reranking for {len(documents)} documents using model '{reranker_config.model}'"
-        )
-
-        # Prepare passages for reranking
-        passages: List[dict[str, Any]] = [
-            {"text": doc.page_content, "meta": doc.metadata} for doc in documents
-        ]
-
-        # Create reranking request
-        rerank_request: RerankRequest = RerankRequest(query=query, passages=passages)
-
-        # Perform reranking
-        logger.debug(f"Reranking {len(passages)} passages for query: '{query[:50]}...'")
-        results: List[dict[str, Any]] = ranker.rerank(rerank_request)
-
-        # Apply top_n filtering
-        top_n: int = reranker_config.top_n or len(documents)
-        results = results[:top_n]
-        logger.debug(
-            f"Reranking complete. Filtered to top {top_n} results from {len(documents)} candidates"
-        )
-
-        # 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)
-
-        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"
-        )
-
-        return reranked_docs
-
-    # Create the main vector search tool using @tool decorator
-    # Note: args_schema provides descriptions for query and filters,
-    # so Annotated is only needed for injected LangGraph parameters
-    @tool(
-        name_or_callable=name or index_name,
-        description=description or "Search for documents using vector similarity",
-        args_schema=VectorSearchRetrieverToolInput,
+    # 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,
     )
-    def vector_search_tool(
-        query: str,
-        filters: Optional[List[FilterItem]] = None,
-        state: Annotated[dict, InjectedState] = None,
-        tool_call_id: Annotated[str, InjectedToolCallId] = 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
-
-        Both stages are traced in MLflow for observability.
-
-        Returns:
-            Command with ToolMessage containing the retrieved documents
-        """
-        logger.debug(
-            f"Vector search tool called: query='{query[:50]}...', reranking={reranker_config is not None}"
-        )
-
-        # 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'})"
-            )
-            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
+    logger.success("Vector search tool created", name=tool.name, index=index_name)
 
-    return vector_search_tool
+    return tool