helo-ayush/Diarization_VoiceFingerprinted

1# ==============================================================================
2# VECTOR SEARCH PROCESSOR
3# Connects Gemini query rewriting with MongoDB Atlas Vector Search indexing.
4# ==============================================================================
5import os
6import asyncio
7import time
8from langchain_google_genai import ChatGoogleGenerativeAI
9from langchain_core.messages import HumanMessage
10from database import db
11from utils.embedding_processor import generate_embedding
12
13
14# Gemini for query optimization
15llm = ChatGoogleGenerativeAI(
16    model="gemini-3-flash-preview",
17    google_api_key=os.getenv("GEMINI_API_KEY"),
18    temperature=0.1,
19)
20
21# MongoDB collection
22transcriptions = db["transcriptions"]
23
24
25async def optimize_query(user_query: str) -> str:
26    """
27    Use Gemini to convert a natural language query into a
28    search-optimized text that will produce a better embedding
29    for similarity matching against our stored summaries.
30    """
31    prompt = f"""You are a search query optimizer for a call transcription database. 
32Each record has a summary describing: the customer issue, the resolution, technical entities, and satisfaction level.
33
34Convert this user's natural language query into a dense, keyword-rich search phrase that will match well against those summaries. 
35
36Rules:
37- Output ONLY the optimized query text, nothing else
38- Include relevant synonyms and related terms
39- Keep it under 50 words
40- Focus on the core intent
41
42User query: "{user_query}"
43
44Optimized search query:"""
45
46    try:
47        message = HumanMessage(content=prompt)
48        response = await asyncio.wait_for(
49            asyncio.to_thread(llm.invoke, [message]),
50            timeout=15
51        )
52        optimized = response.content.strip().strip('"')
53        print(f"   🔍 Optimized query: {optimized}")
54        return optimized
55    except Exception as e:
56        print(f"   ⚠️ Query optimization failed: {e}, using original")
57        return user_query
58
59
60async def vector_search(query: str, limit: int = 5, all_entries: bool = False) -> dict:
61    """
62    Full search pipeline:
63    1. Optimize query with Gemini
64    2. Generate embedding
65    3. Run MongoDB Atlas Vector Search
66    4. Return results
67    """
68    start_time = time.time()
69
70    # Step 1: Optimize query
71    print(f"🔎 Search query: \"{query}\" (all_entries={all_entries})")
72    print("   🧠 Optimizing query with Gemini...")
73    optimized_query = await optimize_query(query)
74
75    # Step 2: Generate embedding from optimized query
76    query_embedding = await generate_embedding(optimized_query)
77
78    if not query_embedding:
79        return {
80            "query": query,
81            "optimizedQuery": optimized_query,
82            "results": [],
83            "totalResults": 0,
84            "error": "Failed to generate query embedding"
85        }
86
87    # Step 3: MongoDB Atlas Vector Search
88    # Atlas Vector Search requires a 'limit'. To "bypass" it for 'all_entries', 
89    # we use a significantly higher limit. 1000 is a safe threshold for a single
90    # request response. In a massive production system, you'd use pagination.
91    MAX_ALL_ENTRIES = 1000 
92    effective_limit = MAX_ALL_ENTRIES if all_entries else limit
93    
94    # numCandidates should be larger than limit to ensure accuracy, 
95    # but we cap it to avoid performance degradation on large datasets.
96    num_candidates = min(effective_limit * 10, 5000) if not all_entries else 2000
97
98    print(f"   📊 Running vector search (limit={effective_limit}, candidates={num_candidates})...")
99    
100    pipeline = [
101        {
102            "$vectorSearch": {
103                "index": "vector_index",
104                "path": "embedding",
105                "queryVector": query_embedding,
106                "numCandidates": num_candidates,
107                "limit": effective_limit,
108            }
109        },
110        {
111            "$project": {
112                "_id": {"$toString": "$_id"},
113                "filename": 1,
114                "summary": 1,
115                "transcript": 1,
116                "satisfactionScore": 1,
117                "tags": 1,
118                "detectedRoles": 1,
119                "speakerCount": 1,
120                "createdAt": 1,
121                "score": {"$meta": "vectorSearchScore"},
122            }
123        }
124    ]
125
126    results = []
127    async for doc in transcriptions.aggregate(pipeline):
128        print(f"      📎 {doc.get('filename', '?')} → score: {doc.get('score', 0):.4f}")
129        results.append(doc)
130
131    # Filter out low-relevance results
132    MIN_SCORE = 0.80
133    filtered = [r for r in results if r.get("score", 0) >= MIN_SCORE]
134    dropped = len(results) - len(filtered)
135    if dropped:
136        print(f"   🚫 Filtered out {dropped} low-relevance results (score < {MIN_SCORE})")
137
138    elapsed = int((time.time() - start_time) * 1000)
139    print(f"   ✅ Found {len(filtered)} relevant results in {elapsed}ms")
140
141    return {
142        "query": query,
143        "optimizedQuery": optimized_query,
144        "results": filtered,
145        "totalResults": len(filtered),
146        "processingMs": elapsed,
147    }
148