""" RAG (Retrieval-Augmented Generation) service for Silver Table Assistant. Handles document loading, vector storage, and similarity search using Supabase vector store. """ import os import logging from pathlib import Path from typing import List, Optional, Dict, Any from uuid import uuid4 import asyncio from langchain_openai import OpenAIEmbeddings from langchain_community.document_loaders import PyPDFLoader, UnstructuredMarkdownLoader from langchain_community.vectorstores import SupabaseVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter from supabase import create_client, Client from langchain_core.documents import Document from huggingface_hub import snapshot_download from cache import DocumentCache, document_cache, cache_result # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class RAGService: """RAG service for document management and similarity search.""" def __init__(self): """Initialize RAG service with OpenAI embeddings (via LiteLLM) and Supabase vector store.""" # Environment variables self.supabase_url = os.getenv("SUPABASE_URL") self.supabase_service_key = os.getenv("SUPABASE_SERVICE_ROLE_KEY") self.openai_api_key = os.getenv("OPENAI_API_KEY") or os.getenv("LITELLM_API_KEY", "sk-eT_04m428oAPUD5kUmIhVA") self.openai_base_url = os.getenv("OPENAI_BASE_URL") or os.getenv("LITELLM_BASE_URL", "https://litellm-ekkks8gsocw.dgx-coolify.apmic.ai/") if not all([self.supabase_url, self.supabase_service_key, self.openai_api_key]): raise ValueError("Missing required environment variables: SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY, OPENAI_API_KEY or LITELLM_API_KEY") # Initialize OpenAI embeddings (works with LiteLLM compatible endpoints) embed_kwargs = { "model": "azure-text-embedding-3-large", "openai_api_key": self.openai_api_key } if self.openai_base_url: embed_kwargs["openai_api_base"] = self.openai_base_url self.embeddings = OpenAIEmbeddings(**embed_kwargs) logger.info(f"Initialized OpenAIEmbeddings with base_url: {self.openai_base_url}") # Initialize Supabase client self.supabase_client: Client = create_client( self.supabase_url, self.supabase_service_key ) # Initialize Supabase vector store self.vector_store = SupabaseVectorStore( client=self.supabase_client, embedding=self.embeddings, table_name="documents", query_name="match_documents" ) # Text splitter for document chunking self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200, length_function=len, is_separator_regex=False, ) async def load_knowledge_base(self, data_dir: str = "backend/data") -> Dict[str, Any]: """ Load and process documents from the data directory. If local directory is empty, download from Hugging Face Dataset. Args: data_dir: Path to directory containing documents Returns: Dictionary with loading statistics """ data_path = Path(data_dir) results = { "total_files": 0, "processed_files": 0, "failed_files": 0, "total_chunks": 0, "errors": [] } # 如果本地資料夾不存在或裡面沒有 PDF/MD 檔案，就從 HF Dataset 下載 if not data_path.exists() or not any(data_path.glob("*.pdf")) and not any(data_path.glob("*.md")): logger.info("Local knowledge base empty or missing. Downloading from Hugging Face Dataset...") data_path.mkdir(parents=True, exist_ok=True) try: snapshot_download( repo_id="pcreem/dietinstruction", # ← 這裡一定要正確！ local_dir=data_dir, local_dir_use_symlinks=False, repo_type="dataset", revision="main", allow_patterns=["*.pdf", "*.md", "*.txt"], # 只下載我們需要的檔案 tqdm_class=None # 避免日誌衝突 ) logger.info(f"Successfully downloaded knowledge base to {data_dir}") except Exception as e: error_msg = f"Failed to download from Hugging Face Dataset: {str(e)}" logger.error(error_msg) results["errors"].append(error_msg) # 如果下載失敗，至少確保資料夾存在 data_path.mkdir(parents=True, exist_ok=True) else: logger.info(f"Using existing local knowledge base at {data_dir}") # ===== 以下是原本的檔案載入邏輯（不需改動太多）===== documents: List[Document] = [] # Supported file types pdf_files = list(data_path.glob("*.pdf")) md_files = list(data_path.glob("*.md")) txt_files = list(data_path.glob("*.txt")) all_files = pdf_files + md_files + txt_files results["total_files"] = len(all_files) if not all_files: logger.warning("No documents found in knowledge base directory") return results for file_path in all_files: try: logger.info(f"Processing file: {file_path.name}") if file_path.suffix == ".pdf": loader = PyPDFLoader(str(file_path)) elif file_path.suffix == ".md": loader = UnstructuredMarkdownLoader(str(file_path)) elif file_path.suffix == ".txt": # Simple text loader with open(file_path, "r", encoding="utf-8") as f: content = f.read() documents.append(Document( page_content=content, metadata={"file_name": file_path.name, "source": str(file_path)} )) results["processed_files"] += 1 continue else: continue docs = loader.load() for doc in docs: doc.metadata.update({ "file_name": file_path.name, "source": str(file_path) }) documents.extend(docs) results["processed_files"] += 1 except Exception as e: error_msg = f"Error processing {file_path.name}: {str(e)}" logger.error(error_msg) results["errors"].append(error_msg) results["failed_files"] += 1 # Split documents into chunks if documents: chunks = self.text_splitter.split_documents(documents) results["total_chunks"] = len(chunks) logger.info(f"Created {len(chunks)} document chunks") # Add to vector store (with upsert) try: self.vector_store.add_documents(chunks) logger.info(f"Successfully added {len(chunks)} chunks to vector store") except Exception as e: error_msg = f"Error adding documents to vector store: {str(e)}" logger.error(error_msg) results["errors"].append(error_msg) else: logger.warning("No documents were successfully loaded") return results async def _process_file(self, file_path: Path, results: Dict[str, Any]) -> None: """ Process a single file and add to vector store. Args: file_path: Path to the file results: Results dictionary to update """ logger.info(f"Processing file: {file_path}") # Load document based on file type if file_path.suffix.lower() == ".pdf": loader = PyPDFLoader(str(file_path)) documents = loader.load() elif file_path.suffix.lower() == ".md": # Try UnstructuredMarkdownLoader first, fallback to simple text loading try: loader = UnstructuredMarkdownLoader(str(file_path)) documents = loader.load() except Exception: # Fallback to simple text loading with open(file_path, 'r', encoding='utf-8') as f: content = f.read() documents = [Document(page_content=content, metadata={"source": str(file_path)})] else: raise ValueError(f"Unsupported file type: {file_path.suffix}") # Split documents into chunks chunks = self.text_splitter.split_documents(documents) # Add metadata to chunks for chunk in chunks: chunk.metadata["source"] = str(file_path) chunk.metadata["file_name"] = file_path.name chunk.metadata["chunk_id"] = str(uuid4()) # Add chunks to vector store if chunks: await self.vector_store.aadd_documents(chunks) results["processed_files"] += 1 results["total_chunks"] += len(chunks) logger.info(f"Added {len(chunks)} chunks from {file_path}") else: logger.warning(f"No chunks generated from {file_path}") @cache_result(document_cache, "rag_documents", ttl=1800) async def get_relevant_documents(self, query: str, k: int = 8) -> List[Document]: """ Perform similarity search to find relevant documents with caching. Args: query: Search query k: Number of documents to return (default: 8) Returns: List of relevant Document objects """ logger.info(f"Searching for relevant documents with query: '{query}' (k={k})") try: # Check cache first cached_results = DocumentCache.get_relevant_documents(query, k) if cached_results is not None: logger.info(f"Returning cached results for query: '{query}'") return cached_results # Perform similarity search try: results = await self.vector_store.asimilarity_search(query, k=k) except Exception as e: if "'SyncRPCFilterRequestBuilder' object has no attribute 'params'" in str(e) or "'AsyncRPCFilterRequestBuilder' object has no attribute 'params'" in str(e): logger.warning(f"SupabaseVectorStore incompatibility detected, using manual RPC: {str(e)}") # Manual RPC fallback embedding = await self.embeddings.aembed_query(query) res = self.supabase_client.rpc( "match_documents", { "query_embedding": embedding, "match_threshold": 0.1, "match_count": k, } ).execute() results = [] for row in res.data: results.append(Document( page_content=row["content"], metadata=row["metadata"] )) else: raise e # Cache the results DocumentCache.set_relevant_documents(query, k, results) logger.info(f"Found {len(results)} relevant documents") return results except Exception as e: logger.error(f"Error during similarity search: {str(e)}") return [] @cache_result(document_cache, "rag_documents_scored", ttl=1800) async def get_relevant_documents_with_scores(self, query: str, k: int = 8, score_threshold: float = 0.7) -> List[Document]: """ Perform similarity search with score threshold and pagination support. Args: query: Search query k: Number of documents to return score_threshold: Minimum similarity score Returns: List of relevant Document objects above threshold """ logger.info(f"Searching for relevant documents with query: '{query}' (k={k}, threshold={score_threshold})") try: # Check cache first cached_results = DocumentCache.get_relevant_documents(query, k, score_threshold) if cached_results is not None: logger.info(f"Returning cached scored results for query: '{query}'") return cached_results # Perform similarity search with scores try: results = await self.vector_store.asimilarity_search_with_score(query, k=k*2) # Get more to filter filtered_results = [doc for doc, score in results if score >= score_threshold][:k] except Exception as e: if "'SyncRPCFilterRequestBuilder' object has no attribute 'params'" in str(e) or "'AsyncRPCFilterRequestBuilder' object has no attribute 'params'" in str(e): logger.warning(f"SupabaseVectorStore incompatibility detected in scored search, using manual RPC: {str(e)}") # Manual RPC fallback embedding = await self.embeddings.aembed_query(query) res = self.supabase_client.rpc( "match_documents", { "query_embedding": embedding, "match_threshold": score_threshold, "match_count": k, } ).execute() filtered_results = [] for row in res.data: filtered_results.append(Document( page_content=row["content"], metadata=row["metadata"] )) else: raise e # Cache the results DocumentCache.set_relevant_documents(query, k, filtered_results, score_threshold) logger.info(f"Found {len(filtered_results)} relevant documents above threshold") return filtered_results except Exception as e: logger.error(f"Error during similarity search with scores: {str(e)}") return [] async def get_relevant_documents_paginated( self, query: str, page: int = 1, page_size: int = 10, score_threshold: Optional[float] = None ) -> Dict[str, Any]: """ Perform paginated similarity search. Args: query: Search query page: Page number (1-indexed) page_size: Number of documents per page score_threshold: Minimum similarity score (optional) Returns: Dictionary with documents, pagination info, and metadata """ logger.info(f"Paginated search for query: '{query}' (page={page}, page_size={page_size})") try: # Calculate total number needed (get more for pagination) total_needed = page * page_size # Perform search with more results for pagination if score_threshold: results = await self.get_relevant_documents_with_scores(query, k=total_needed, score_threshold=score_threshold) else: results = await self.get_relevant_documents(query, k=total_needed) # Paginate results start_idx = (page - 1) * page_size end_idx = start_idx + page_size paginated_results = results[start_idx:end_idx] # Calculate pagination metadata total_results = len(results) total_pages = (total_results + page_size - 1) // page_size # Ceiling division has_next = page < total_pages has_prev = page > 1 return { "documents": paginated_results, "pagination": { "page": page, "page_size": page_size, "total_results": total_results, "total_pages": total_pages, "has_next": has_next, "has_prev": has_prev, "start_index": start_idx, "end_index": end_idx }, "query": query, "score_threshold": score_threshold } except Exception as e: logger.error(f"Error during paginated search: {str(e)}") return { "documents": [], "pagination": { "page": page, "page_size": page_size, "total_results": 0, "total_pages": 0, "has_next": False, "has_prev": False, "start_index": 0, "end_index": 0 }, "query": query, "score_threshold": score_threshold, "error": str(e) } async def get_document_count(self) -> int: """ Get the total number of documents in the vector store. Returns: Total number of documents """ try: # This is a simplified approach - in a real implementation, # you'd want to query the Supabase table directly result = self.supabase_client.table("documents").select("id", count="exact").execute() return result.count if result.count else 0 except Exception as e: logger.error(f"Error getting document count: {str(e)}") return 0 async def clear_knowledge_base(self) -> bool: """ Clear all documents from the vector store. Returns: True if successful, False otherwise """ try: # Delete all documents from the table result = self.supabase_client.table("documents").delete().gte("id", "").execute() logger.info("Knowledge base cleared successfully") return True except Exception as e: logger.error(f"Error clearing knowledge base: {str(e)}") return False # Global RAG service instance rag_service: Optional[RAGService] = None def get_rag_service() -> RAGService: """ Get or create the global RAG service instance. Returns: RAGService instance """ global rag_service if rag_service is None: rag_service = RAGService() return rag_service # Convenience functions for backward compatibility async def load_knowledge_base(data_dir: str = "backend/data") -> Dict[str, Any]: """Load knowledge base documents.""" service = get_rag_service() return await service.load_knowledge_base(data_dir) async def get_relevant_documents(query: str, k: int = 8) -> List[Document]: """Get relevant documents for a query.""" service = get_rag_service() return await service.get_relevant_documents(query, k) async def get_relevant_documents_with_scores(query: str, k: int = 8, score_threshold: float = 0.7) -> List[Document]: """Get relevant documents with similarity scores.""" service = get_rag_service() return await service.get_relevant_documents_with_scores(query, k, score_threshold) if __name__ == "__main__": """ Main block for testing and manual knowledge base loading. """ async def main(): """Main function for testing.""" print("Loading knowledge base...") try: # Initialize and load knowledge base service = get_rag_service() results = await service.load_knowledge_base() print(f"Knowledge base loading results:") print(f"- Total files: {results['total_files']}") print(f"- Processed files: {results['processed_files']}") print(f"- Failed files: {results['failed_files']}") print(f"- Total chunks: {results['total_chunks']}") if results['errors']: print(f"- Errors: {len(results['errors'])}") for error in results['errors']: print(f" * {error}") # Test search functionality test_query = "高血壓飲食建議" print(f"\nTesting search with query: '{test_query}'") documents = await service.get_relevant_documents(test_query) print(f"Found {len(documents)} relevant documents:") for i, doc in enumerate(documents, 1): print(f"{i}. {doc.metadata.get('file_name', 'Unknown')} - {doc.page_content[:100]}...") print(f"\nTotal documents in vector store: {await service.get_document_count()}") except Exception as e: print(f"Error: {str(e)}") raise # Run the main function asyncio.run(main())