""" CDMS RAG Search Module Searches Qdrant vector database for CDMS label information with page citations """ import sys from pathlib import Path from typing import Dict, List, Optional # Add project root to path project_root = Path(__file__).parent.parent.parent sys.path.insert(0, str(project_root)) from src.rag.vector_store import QdrantVectorStore, get_shared_vector_store from src.rag.embeddings import OpenAIEmbeddingService from src.config.credentials import CredentialsManager from qdrant_client.models import Filter, FieldCondition, MatchValue class CDMSRAGSearch: """ RAG search for CDMS pesticide labels Searches Qdrant vector database for relevant label information with accurate page number tracking Usage: searcher = CDMSRAGSearch() results = searcher.search("What's the application rate for Roundup?", product_name="Roundup") """ def __init__(self, vector_store=None): """Initialize RAG search components. vector_store: optional shared QdrantVectorStore. The embedded on-disk Qdrant permits one client per path per process, so online mode passes a single store shared with the DocumentLoader instead of opening a second. """ # Reuse the provided store, else the process-wide singleton (never open a # second client to the embedded store — see get_shared_vector_store). if vector_store is not None: self.vector_store = vector_store else: try: self.vector_store = get_shared_vector_store() except Exception as e: print(f"⚠️ Warning: Could not initialize Qdrant: {e}") self.vector_store = None # Initialize embedding service try: creds = CredentialsManager() openai_key = creds.get_api_key("openai") self.embedding_service = OpenAIEmbeddingService(api_key=openai_key) except Exception as e: print(f"⚠️ Warning: Could not initialize OpenAI embeddings: {e}") self.embedding_service = None def search( self, query: str, product_name: Optional[str] = None, limit: int = 5, score_threshold: float = 0.4 ) -> List[Dict]: """ Search CDMS documents in Qdrant Args: query: User's question or search query product_name: Optional product name to filter results (e.g., "Roundup") limit: Maximum number of results to return score_threshold: Minimum similarity score (0-1) Returns: List of dicts with: { "content": str (chunk text), "page_number": int (exact page number), "source_file": str (PDF filename), "score": float (similarity score 0-1), "document_id": str, "chunk_index": int } """ if not self.vector_store or not self.embedding_service: return [] # Abstention gate: if the caller asked for a specific product that we # have not indexed (or whose text failed to extract), return nothing so # the response layer can say "not found" instead of the vector search # falling back to the product that dominates the index. if product_name: try: from src.cdms.product_catalog import get_catalog if not get_catalog().is_available(product_name): print(f"ℹ️ '{product_name}' is not in the indexed catalog - abstaining.") return [] except Exception as e: print(f"⚠️ Catalog availability check skipped: {e}") try: # Generate query embedding query_embedding = self.embedding_service.generate_embedding(query) if not query_embedding: print("⚠️ Warning: Failed to generate query embedding") return [] # Over-fetch so we have room to diversify (general query) or keep the # best chunks within a product (scoped query). search_limit = limit * 3 if product_name: # Scope retrieval to the requested product AT the vector level. # This is the robust fix for the ISA "wrong herbicide" bug: a # minority product (e.g. boron = ~2% of the index) is no longer # crowded out of a global, Roundup-dominated top-k before the # Python post-filter ever runs. from src.cdms.product_catalog import normalize_filename canonical = normalize_filename(product_name) # The Qdrant filter already guarantees we only see the requested # product's chunks, so the 0.4 wrong-product gate no longer applies # here. Use a low threshold so an explicitly-requested, indexed # product still returns its label even when a generic query matches # its (e.g. SDS) text only weakly -- Actagro's best chunk is ~0.34, # which the 0.4 gate wrongly turned into a "0 labels" abstention. product_threshold = min(score_threshold, 0.2) product_filter = Filter( must=[FieldCondition(key="product", match=MatchValue(value=canonical))] ) results = self.vector_store.search_documents( query_embedding=query_embedding, limit=search_limit, score_threshold=product_threshold, query_filter=product_filter, ) if results: results = results[:limit] else: # Fallback for indexes built before the `product` payload field # existed (or a filter miss): global search + substring # post-filter — the previous behavior. Abstention has already # guaranteed the product is in the catalog. global_results = self.vector_store.search_documents( query_embedding=query_embedding, limit=search_limit, score_threshold=product_threshold, ) product_lower = product_name.lower() results = [ r for r in global_results if product_lower in r.get("source_file", "").lower() or product_lower in r.get("metadata", {}).get("document_name", "").lower() ][:limit] else: # No product filter (general query): diversify so the product # that dominates the index doesn't monopolise every slot. results = self.vector_store.search_documents( query_embedding=query_embedding, limit=search_limit, score_threshold=score_threshold, ) from src.cdms.product_catalog import diversify_by_product results = diversify_by_product(results, limit=limit, max_per_product=2) # Format results with page numbers and PDF URLs formatted_results = [] for result in results: # PHASE 1 FIX: Extract pdf_url from result with multiple fallback strategies metadata = result.get("metadata", {}) # Strategy 1: Direct pdf_url field (preferred - from vector_store extraction) pdf_url = result.get("pdf_url", "") # Strategy 2: From metadata dict if not pdf_url: pdf_url = metadata.get("pdf_url", "") # Strategy 3: From payload directly (if metadata is the payload) if not pdf_url and isinstance(metadata, dict): pdf_url = metadata.get("pdf_url", "") # Also extract url_hash for matching url_hash = result.get("url_hash", "") or metadata.get("url_hash", "") # PHASE 2 FIX: Extract page_number with multiple fallback strategies page_number = result.get("page_number", 0) # Strategy 1: Direct page_number field (preferred - from vector_store extraction) if page_number <= 0: # Strategy 2: From metadata dict page_number = metadata.get("page_number", 0) # Strategy 3: Validate and fix if still invalid if page_number <= 0: # Fallback: Estimate page number based on chunk_index chunk_index = metadata.get("chunk_index", 0) if chunk_index > 0: # Rough estimate: 3 chunks per page page_number = (chunk_index // 3) + 1 else: # Last resort: use page 1 page_number = 1 print(f"⚠️ Warning: Invalid or missing page_number for chunk, using estimated value: {page_number}") formatted_results.append({ "content": result.get("content", ""), "page_number": page_number, # PHASE 2 FIX: Validated page number "source_file": result.get("source_file", "Unknown"), "score": result.get("score", 0.0), "document_id": result.get("document_id", ""), "chunk_index": metadata.get("chunk_index", 0), "document_name": metadata.get("document_name", ""), "pdf_url": pdf_url, # PHASE 1 FIX: Include PDF URL from metadata "url_hash": url_hash # PHASE 1 FIX: Include URL hash for matching }) return formatted_results except Exception as e: print(f"⚠️ Warning: RAG search failed: {e}") import traceback traceback.print_exc() return [] def search_by_product( self, product_name: str, query: Optional[str] = None, limit: int = 5 ) -> List[Dict]: """ Search for specific product information Args: product_name: Product name (e.g., "Roundup") query: Optional specific question (if None, returns general product info) limit: Maximum results Returns: List of relevant chunks with page numbers """ if query: return self.search(query=query, product_name=product_name, limit=limit) else: # General product search return self.search( query=f"{product_name} pesticide label information", product_name=product_name, limit=limit ) def get_collection_stats(self) -> Dict: """ Get statistics about the CDMS documents collection Returns: Dict with collection information """ if not self.vector_store: return {"error": "Vector store not initialized"} try: info = self.vector_store.get_collection_info() return info except Exception as e: return {"error": str(e)} # Test function if __name__ == "__main__": print("=" * 80) print("Testing CDMS RAG Search") print("=" * 80) try: searcher = CDMSRAGSearch() # Check collection stats print("\n📊 Collection Stats:") stats = searcher.get_collection_stats() if "error" not in stats: if "cdms_documents" in stats: doc_info = stats["cdms_documents"] print(f" Documents in Qdrant: {doc_info.get('points_count', 0)}") print(f" Vectors: {doc_info.get('vectors_count', 0)}") else: print(" ⚠️ No documents indexed yet") else: print(f" ⚠️ {stats['error']}") # Test search (if documents exist) if stats.get("cdms_documents", {}).get("points_count", 0) > 0: print("\n🔍 Testing search: 'application rate'") results = searcher.search("application rate", limit=3) if results: print(f" ✅ Found {len(results)} result(s)") for i, result in enumerate(results, 1): print(f"\n {i}. Score: {result['score']:.3f}") print(f" Page: {result['page_number']}") print(f" File: {result['source_file']}") print(f" Content: {result['content'][:100]}...") else: print(" ⚠️ No results found") else: print("\n💡 To test search:") print(" 1. Download PDFs (Phase 1)") print(" 2. Process and index PDFs in Qdrant") print(" 3. Run this test again") print("\n" + "=" * 80) except Exception as e: print(f"❌ Error: {e}") import traceback traceback.print_exc()