import os import logging from dotenv import load_dotenv from motor.motor_asyncio import AsyncIOMotorClient # IMPORT AsyncMongoClient from pythainlp.tokenize import word_tokenize # Moved import here import models # Keep standard import import asyncio from typing import Optional, Dict from utils import get_device # import time # No longer needed for reranker # import numpy as np # No longer needed for reranker # import onnxruntime as ort # No longer needed for reranker # from transformers import AutoTokenizer # No longer needed for reranker # Load environment variables if get_device() == "cpu": load_dotenv(override=True) if get_device() == "mps": load_dotenv(override=True) # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # MongoDB Configuration DATABASE_URL = os.getenv("MONGO_URL") # DATABASE_URL = "mongodb://rabbit_reward:rabbit_password@host.docker.internal:27017/?directConnection=true" DB_NAME = "homeshopping" DEFAULT_VECTOR_INDEX = "default" # Example: Make configurable DEFAULT_KEYWORD_INDEX = "default" # Example: Make configurable class MongoHybridSearch: def __init__(self, database_name=DB_NAME, mongo_uri=DATABASE_URL): """ Initialize MongoDB connection and embedder. """ try: # self.client = AsyncIOMotorClient(mongo_uri) # self.database = self.client[database_name] # # Consider making collection name configurable # self.collection = self.database["homeshopping"] # # self.collection_fact = self.database["SCG_financial_report_jai"] # self.llm_analyzer = models.LLMFinanceAnalyzer() # self.embedder = models.Embedder() # Instantiate Embedder class from models logger.info("MongoHybridSearch initialized successfully.") except Exception as e: logger.error(f"Failed to initialize MongoHybridSearch: {e}") raise # Re-raise exception to prevent app from starting with bad config async def search_documents(self, query: str) -> list[str]: """ Find relevant data for each (subquery, original_query, quarter, year). Args: query_list (list): List of tuples (subquery, original_query, quarter, year). Returns: list: List of lists, where each inner list contains relevant document content strings. Returns empty list if an error occurs during the overall search process. """ try: all_docs_content = [] # for subquery, subkeyword, quarter, year in query_list: # Unpack the tuple # Pass configured index names result_content = await self.atlas_hybrid_search(collection_name = self.collection, query=query, top_k=100, # Consider making configurable exact_top_k=17, # Consider making configurable vector_index_name=DEFAULT_VECTOR_INDEX, keyword_index_name=DEFAULT_KEYWORD_INDEX, ) all_docs_content.append(result_content) return result_content except Exception as e: logger.error(f"Error in search_documents: {e}") return [] # Return empty list on failure async def atlas_hybrid_search(self, collection_name :str, query: str, top_k: int, exact_top_k: int, vector_index_name: str, keyword_index_name: str, ) -> list[str]: """ Perform hybrid search using Atlas Vector Search & Keyword Search. Returns a list of document content strings. """ try: # Ensure quarter and year are strings for MongoDB query # quarter_str = [str(quarter)] # year_str = [str(year)] # if collection_name == "fact": # collection = self.collection_fact # elif collection_name == "report": # collection = self.collection_report # top_k = 15 # For report collection, we might want fewer results # exact_top_k = 7 # else: # pass query_vector = await self.embedder.embed(query, "query") print(len(query_vector)) # query_vector = query_vector[0] if not query_vector: logger.error(f"Failed to get embedding for query: {query}") return [] # Perform vector search vector_pipeline = [ { "$vectorSearch": { "queryVector": query_vector, "path": "embedding", # Ensure 'embedding' is the correct field name "numCandidates": 10000, # Consider making configurable "limit": top_k, "index": vector_index_name, # "filter": { # "$and": [ # {"quarter": {"$in": quarter_str}}, # {"year": {"$in": year_str}} # ] # } } }, {"$project": {"_id": 1, "content": 1, "score": {"$meta": "vectorSearchScore"}}} ] vector_results_cursor = self.collection.aggregate(vector_pipeline) vector_results = await vector_results_cursor.to_list(length=top_k) logger.info(f"Vector search found {len(vector_results)} results for query: '{query}'") # Tokenize query for keyword search using PyThaiNLP query_tokens = word_tokenize(query, engine="newmm", keep_whitespace=False) logger.info(f"Keyword search tokens: {query_tokens}") # Perform keyword search (Atlas Search) keyword_pipeline = [ { "$search": { "index": keyword_index_name, "text": { "query": query_tokens, "path": "content_tokenized" } } }, # { # "$match": { # "$and": [ # {"quarter": {"$in": quarter_str}}, # {"year": {"$in": year_str}} # ] # } # }, { "$project": { "_id": 1, "content": 1, "score": {"$meta": "searchScore"} } }, {"$limit": top_k} ] keyword_results_cursor = self.collection.aggregate(keyword_pipeline) keyword_results = await keyword_results_cursor.to_list(length=top_k) # Using length for explicit limit from cursor logger.info(f"Keyword search found {len(keyword_results)} results for query: '{query}'") # Apply Weighted Reciprocal Rank Fusion (WRRF) # Prepare results in the expected format for WRRF: list of dicts with _id and content print(f"Vector results: {len(vector_results)}, Keyword results: {len(keyword_results)}") vec_docs = [{"_id": str(doc["_id"]), "content": doc.get("content", "")} for doc in vector_results] key_docs = [{"_id": str(doc["_id"]), "content": doc.get("content", "")} for doc in keyword_results] # Handle potential missing 'content' key more robustly # Ensure content is string for doc_list in [vec_docs, key_docs]: for doc in doc_list: if not isinstance(doc["content"], str): logger.warning(f"Document content is not a string (ID: {doc['_id']}), converting.") doc["content"] = str(doc["content"]) fused_documents = self.weighted_reciprocal_rank([vec_docs, key_docs], top_k) if len(fused_documents) < exact_top_k: exact_top_k = len(fused_documents) fused_documents = fused_documents[:exact_top_k] # async def check_and_get_relevant(doc: Dict) -> Optional[Dict]: # # Use a helper to run the classification and return the doc if relevant # is_relevant = await self.llm_analyzer.classify_relevance(query=query, document_content=doc.get("content", "")) # if is_relevant: # return doc # return None # tasks = [check_and_get_relevant(doc) for doc in fused_documents] # relevance_results = await asyncio.gather(*tasks) # # Filter out None values (non-relevant docs) # relevant_docs = [doc for doc in relevance_results if doc is not None] # logger.info(f"Found {len(relevant_docs)} relevant documents after LLM classification (out of {len(fused_documents)}).") # # if len(relevant_docs) < exact_top_k: # # exact_top_k = len(relevant_docs) # # Return only the content strings, limited to exact_top_k # return [doc["content"] for doc in relevant_docs] if not fused_documents: logger.info("No documents to rank after fusion.") return [] # 1. Format documents for the LLM # docs_for_selection = { # idx: doc.get("content", "") # for idx, doc in enumerate(fused_documents) # } # # 2. Call the LLM to get indices of relevant documents # selected_indices = await self.llm_analyzer.select_relevant_documents( # query=query, # documents=docs_for_selection # ) # # 3. Filter the original fused_documents list based on the selected indices # relevant_docs = [] # if selected_indices: # # Create a set for efficient lookup and filter out-of-bounds indices # valid_indices = set(idx for idx in selected_indices if 0 <= idx < len(fused_documents)) # relevant_docs = [fused_documents[i] for i in sorted(list(valid_indices))] # Sort to maintain some order # return [doc["content"] for doc in relevant_docs] # else: # return [e["content"] for e in fused_documents] # If no indices selected, return all content # --- END OF NEW LOGIC --- return [e["content"] for e in fused_documents] except Exception as e: logger.error(f"Error in atlas_hybrid_search for query '{query}': {e}", exc_info=True) return [] def weighted_reciprocal_rank(self, doc_lists: list[list[dict]], top_k: int) -> list[dict]: """ Apply Weighted Reciprocal Rank Fusion (WRRF) to rank results. Args: doc_lists: List of lists of documents. Each inner list is from one search method. Each document is a dict with at least '_id' and 'content'. top_k: The maximum number of documents to return after fusion. Returns: List of fused documents, sorted by RRF score, limited by top_k. """ try: # Ensure doc_lists is not empty and contains lists if not doc_lists or not all(isinstance(dl, list) for dl in doc_lists): logger.warning("WRRF called with invalid doc_lists.") return [] # Configuration for WRRF c = 60 # Constant for rank penalty, tunable weights = [1.0, 1.0] # Vector search weight, keyword search weight - Tunable if len(doc_lists) != len(weights): # Fallback if weights don't match lists (e.g., one search returned nothing) # This basic handling might need refinement based on desired behavior weights = [1.0] * len(doc_lists) logger.warning(f"Number of doc lists ({len(doc_lists)}) != number of weights ({len(weights)}). Using equal weights.") # raise ValueError("Number of rank lists must be equal to the number of weights.") # Use a dictionary to map unique content to its document dict and accumulate scores # This handles cases where the same doc appears in multiple lists or multiple times rrf_scores = {} # content -> {'score': float, 'doc': dict} for doc_list, weight in zip(doc_lists, weights): processed_ids_in_list = set() # Track IDs within the current list to handle duplicates from the *same* source for rank, doc in enumerate(doc_list, start=1): doc_id = doc.get("_id") content = doc.get("content") # Basic validation if not doc_id or content is None: logger.warning(f"Skipping doc with missing ID or content in WRRF: {doc}") continue if not isinstance(content, str): # Ensure content is string for keying content = str(content) doc["content"] = content # Update doc dict too # Only score the first occurrence of a document *within the same list* if doc_id in processed_ids_in_list: continue processed_ids_in_list.add(doc_id) # Calculate RRF score contribution rank_score = weight * (1.0 / (rank + c)) # Accumulate score or add new entry if content in rrf_scores: rrf_scores[content]['score'] += rank_score else: # Store the first encountered 'doc' dict for this content rrf_scores[content] = {'score': rank_score, 'doc': doc} # Sort documents based on accumulated RRF score # We sort the items (content, score_data) by score sorted_items = sorted(rrf_scores.items(), key=lambda item: item[1]['score'], reverse=True) # Return the document dictionaries from the sorted items, limited by top_k return [item[1]['doc'] for item in sorted_items[:top_k]] except Exception as e: logger.error(f"Error in weighted_reciprocal_rank: {e}", exc_info=True) return [] # Example usage (optional, for testing) if __name__ == "__main__": # To test async code, you need an asyncio event loop async def main_test(): print("Testing MongoHybridSearch...") try: search_engine = MongoHybridSearch() query_example = 'มี product ไรบ้าง' results = await search_engine.search_documents(query_example) # Await here print("\nSearch Results:") if results: print(results) else: print("Search failed or returned no results.") except Exception as e: print(f"An error occurred during testing: {e}") # Run the async test function asyncio.run(main_test())