Essential Techniques for Production Vector Search Systems Part 1 - Hybrid Search

"""Generic Hybrid Search Implementation for Qdrant"""

from qdrant_client import QdrantClient
from dotenv import load_dotenv
import os
from typing import List, Dict, Any, Optional
from collections import Counter
import re

# Cache the embedding model globally
_embedding_model = None

def get_embedding_model():
    """Get or create the embedding model (cached)."""
    global _embedding_model
    if _embedding_model is None:
        try:
            from sentence_transformers import SentenceTransformer
            _embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        except ImportError:
            raise ImportError(
                "sentence-transformers not installed. Install it with: pip install sentence-transformers"
            )
    return _embedding_model


def get_qdrant_client() -> QdrantClient:
    """Initialize and return Qdrant client."""
    load_dotenv()
    return QdrantClient(
        url=os.getenv("QDRANT_URL"),
        api_key=os.getenv("QDRANT_API_KEY"),
    )


def tokenize(text: str) -> List[str]:
    """Simple tokenization - split on whitespace and punctuation."""
    text = re.sub(r'[^\w\s]', ' ', text.lower())
    return [word for word in text.split() if len(word) > 2]


def calculate_tf_score(query_words: List[str], text: str, boost_exact: bool = True) -> float:
    """Calculate Term Frequency score for keyword matching."""
    text_words = tokenize(text)
    if not text_words or not query_words:
        return 0.0
    
    word_counts = Counter(text_words)
    query_word_counts = Counter(query_words)
    
    score = 0.0
    total_words = len(text_words)
    matched_words = 0
    
    for word, query_freq in query_word_counts.items():
        text_freq = word_counts.get(word, 0)
        if text_freq > 0:
            matched_words += 1
            tf = text_freq / total_words
            boost = 2.0 if boost_exact else 1.0
            score += tf * query_freq * boost
    
    match_ratio = matched_words / len(set(query_words)) if query_words else 0.0
    
    if match_ratio == 1.0:
        score *= 1.5
    
    base_score = score / len(set(query_words)) if query_words else 0.0
    return base_score * match_ratio


def hybrid_search(
    collection_name: str,
    query: str,
    client: Optional[QdrantClient] = None,
    vector_weight: float = 0.7,
    keyword_weight: float = 0.3,
    limit: int = 10,
    keyword_fields: Optional[List[str]] = None
) -> List[Dict[str, Any]]:
    """Perform hybrid search combining vector and keyword search."""
    if client is None:
        client = get_qdrant_client()
    
    if keyword_fields is None:
        keyword_fields = ['text', 'description', 'name', 'title', 'content']
    
    query_words = tokenize(query)
    vector_results = []
    try:
        model = get_embedding_model()
        query_vector = model.encode(query).tolist()
        
        vector_search_response = client.query_points(
            collection_name=collection_name,
            query=query_vector,
            limit=limit * 3,
            with_payload=True
        )
        vector_search = vector_search_response.points
        
        if vector_search:
            scores = [r.score for r in vector_search]
            min_score = min(scores)
            max_score = max(scores)
            score_range = max_score - min_score if max_score > min_score else 1.0
            
            vector_results = [
                {
                    "id": r.id,
                    "score": (r.score - min_score) / score_range if score_range > 0 else r.score,
                    "payload": r.payload,
                    "raw_score": r.score
                }
                for r in vector_search
            ]
    except Exception as e:
        print(f"Vector search error: {e}")
    
    keyword_results = []
    try:
        collection_info = client.get_collection(collection_name)
        total_points = collection_info.points_count
        scroll_limit = min(total_points, 1000)
        
        all_points = client.scroll(
            collection_name=collection_name,
            limit=scroll_limit,
            with_payload=True,
            with_vectors=False
        )
        
        for point in all_points[0]:
            keyword_score = 0.0
            exact_match_found = False
            query_lower = query.lower().strip()
            query_upper = query.upper().strip()
            
            for field in keyword_fields:
                field_value = point.payload.get(field, "")
                if isinstance(field_value, str):
                    field_lower = field_value.lower().strip()
                    field_upper = field_value.upper().strip()
                    
                    if query_lower == field_lower or query_upper == field_upper:
                        exact_match_found = True
                        keyword_score = 1.0
                        break
            
            if not exact_match_found:
                for field in keyword_fields:
                    field_value = point.payload.get(field, "")
                    if isinstance(field_value, str):
                        if query_lower in field_value.lower():
                            exact_match_found = True
                        
                        score = calculate_tf_score(query_words, field_value, boost_exact=True)
                        if field in ['name', 'title', 'id', 'part_id', 'part_name']:
                            keyword_score += score * 3.0
                        else:
                            keyword_score += score * 1.0
                
                if exact_match_found:
                    keyword_score *= 5.0
                    keyword_score = min(keyword_score / 10.0, 1.0)
                else:
                    keyword_score = min(keyword_score / 10.0, 1.0)
            
            if keyword_score > 0:
                keyword_results.append({
                    "id": point.id,
                    "score": keyword_score,
                    "payload": point.payload
                })
        
        keyword_results.sort(key=lambda x: x["score"], reverse=True)
        keyword_results = keyword_results[:limit * 3]
        
    except Exception as e:
        print(f"Keyword search error: {e}")
    
    fused_results = {}
    
    for result in vector_results:
        point_id = result["id"]
        fused_results[point_id] = {
            "id": point_id,
            "payload": result["payload"],
            "vector_score": result["score"],
            "keyword_score": 0.0,
            "combined_score": 0.0,
            "raw_vector_score": result.get("raw_score", 0.0)
        }
    
    for result in keyword_results:
        point_id = result["id"]
        if point_id not in fused_results:
            fused_results[point_id] = {
                "id": point_id,
                "payload": result["payload"],
                "vector_score": 0.0,
                "keyword_score": 0.0,
                "combined_score": 0.0,
                "raw_vector_score": 0.0
            }
        fused_results[point_id]["keyword_score"] = result["score"]
    
    for point_id, result in fused_results.items():
        if result["keyword_score"] >= 0.99:
            result["combined_score"] = 1.0
        elif result["vector_score"] > 0 and result["keyword_score"] > 0:
            result["combined_score"] = (
                vector_weight * result["vector_score"] +
                keyword_weight * result["keyword_score"]
            )
        elif result["vector_score"] > 0:
            result["combined_score"] = vector_weight * result["vector_score"]
        elif result["keyword_score"] > 0:
            result["combined_score"] = keyword_weight * result["keyword_score"]
    
    sorted_results = sorted(
        fused_results.values(),
        key=lambda x: x["combined_score"],
        reverse=True
    )
    
    return sorted_results[:limit]


def display_results(results: List[Dict[str, Any]], query: str, show_fields: Optional[List[str]] = None):
    """Display search results."""
    if show_fields is None:
        show_fields = ['name', 'title', 'description', 'text']
    
    print(f"\nHybrid Search Results for: '{query}'")
    print("=" * 80)
    
    if not results:
        print("No results found.")
        return
    
    for i, result in enumerate(results, 1):
        payload = result["payload"]
        display_name = "Result"
        for field in ['name', 'title', 'part_name', 'id', 'part_id']:
            if field in payload:
                display_name = str(payload[field])
                break
        
        print(f"\n{i}. {display_name}")
        
        for field in show_fields:
            if field in payload and field not in ['name', 'title']:
                value = payload[field]
                if isinstance(value, str):
                    print(f"   {field.capitalize()}: {value[:100]}{'...' if len(value) > 100 else ''}")
        
        print(f"   Scores: Vector={result['vector_score']:.3f}, "
              f"Keyword={result['keyword_score']:.3f}, "
              f"Combined={result['combined_score']:.3f}")
        
        if result.get('raw_vector_score'):
            print(f"   Raw Vector Score: {result['raw_vector_score']:.4f}")
        
        print("-" * 80)


if __name__ == "__main__":
    collection_name = os.getenv("QDRANT_COLLECTION", "your_collection")
    query1 = "example query for exact match"
    
    try:
        client = get_qdrant_client()
        results1 = hybrid_search(
            collection_name=collection_name,
            query=query1,
            client=client,
            vector_weight=0.7,
            keyword_weight=0.3,
            limit=3,
            keyword_fields=['name', 'description', 'text']
        )
        display_results(results1, query1, show_fields=['name', 'description'])
    except Exception as e:
        print(f"Error: {e}")
        print("\nTo use this script:")
        print("1. Set QDRANT_URL and QDRANT_API_KEY in your .env file")
        print("2. Set QDRANT_COLLECTION environment variable or update collection_name in code")
        print("3. Adjust keyword_fields to match your collection's payload structure")