Spaces:
Running
Running
| """ | |
| Vector retrieval module - REAL implementation using FAISS | |
| Retrieves relevant documents from user's vector store | |
| """ | |
| from vector.store_faiss import FaissStore | |
| from core.llm import embed_text | |
| import numpy as np | |
| def retrieve(query: str, k: int = 5, user_id: str = None): | |
| """ | |
| Retrieve relevant documents from FAISS vector store. | |
| IMPROVED: Ensures results come from ALL uploaded files, not just top-k similar. | |
| Args: | |
| query: Search query string | |
| k: Number of results to return (will return k per file for comprehensive coverage) | |
| user_id: User identifier for per-user store | |
| Returns: | |
| List of document dicts with text, source, and score | |
| """ | |
| if not user_id: | |
| print("⚠️ No user_id provided to retrieve()") | |
| return [] | |
| try: | |
| # Load user's FAISS store | |
| store = FaissStore.load_or_create(user_id=user_id, fresh=False) | |
| if store.index.ntotal == 0: | |
| print(f"⚠️ FAISS store for {user_id} is empty (0 vectors)") | |
| return [] | |
| print(f"📚 Searching FAISS store with {store.index.ntotal} vectors for user {user_id}") | |
| # Embed the query | |
| query_embedding = embed_text(query) | |
| query_vector = np.array([query_embedding], dtype="float32") | |
| # IMPROVED: Search for MORE results then balance across files | |
| search_k = min(50, max(k * 5, store.index.ntotal)) # Search 50 or 5x requested | |
| distances, indices = store.index.search(query_vector, search_k) | |
| # Group results by source file | |
| results_by_source = {} | |
| all_results = [] | |
| for i, (dist, idx) in enumerate(zip(distances[0], indices[0])): | |
| if idx < 0 or idx >= len(store.meta): | |
| continue | |
| meta = store.meta[idx] | |
| source = meta.get("source", "Unknown") | |
| # Calculate similarity score (convert L2 distance to similarity) | |
| score = 1 / (1 + dist) | |
| result = { | |
| "text": meta.get("text", ""), | |
| "source": source, | |
| "score": float(score), | |
| "chunk": meta.get("chunk", i), | |
| "metadata": meta | |
| } | |
| # Group by source | |
| if source not in results_by_source: | |
| results_by_source[source] = [] | |
| results_by_source[source].append(result) | |
| all_results.append(result) | |
| # CRITICAL: Ensure we get results from EACH file | |
| num_files = len(results_by_source) | |
| print(f"📁 Found results from {num_files} files: {list(results_by_source.keys())}") | |
| if num_files > 1: | |
| # Balance results across all files | |
| chunks_per_file = max(2, k // num_files + 1) | |
| balanced_results = [] | |
| for source, file_results in results_by_source.items(): | |
| # Take top chunks from each file (sorted by score) | |
| file_results.sort(key=lambda x: x['score'], reverse=True) | |
| balanced_results.extend(file_results[:chunks_per_file]) | |
| print(f" 📄 {source}: {len(file_results[:chunks_per_file])} chunks") | |
| # Sort balanced results by score for final ordering | |
| balanced_results.sort(key=lambda x: x['score'], reverse=True) | |
| results = balanced_results[:k * 2] # Return more for comprehensive coverage | |
| else: | |
| # Single file - just return top k | |
| results = all_results[:k] | |
| print(f"✅ Retrieved {len(results)} documents from FAISS (balanced across {num_files} files)") | |
| return results | |
| except Exception as e: | |
| print(f"⚠️ Error in retrieve(): {e}") | |
| import traceback | |
| traceback.print_exc() | |
| return [] | |
| def hybrid_retrieve(query: str, k: int = 5, user_id: str = None): | |
| """ | |
| Hybrid retrieval combining FAISS vector search with keyword matching. | |
| Falls back to standard retrieve if hybrid index not available. | |
| Args: | |
| query: Search query string | |
| k: Number of results to return | |
| user_id: User identifier | |
| Returns: | |
| List of document dicts | |
| """ | |
| # For now, use standard FAISS retrieval | |
| # TODO: Implement BM25 keyword matching + score fusion | |
| return retrieve(query, k=k, user_id=user_id) | |