Update src/rag_engine.py

src/rag_engine.py

@@ -5,6 +5,7 @@ from typing import List, Literal, Tuple
 
 # --- LANGCHAIN & DB IMPORTS ---
 from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_openai import OpenAIEmbeddings
 from langchain_core.documents import Document
 from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
 from sentence_transformers import CrossEncoder
@@ -30,14 +31,29 @@ logger = logging.getLogger(__name__)
 _embedding_func = None
 _rerank_model = None
 
-def get_embedding_func():
-    """Lazy loads the embedding model to save startup resources."""
-    global _embedding_func
-    if _embedding_func is None:
-        logger.info(f"⏳ Loading Embedding Model: {EMBED_MODEL_NAME}...")
-        _embedding_func = HuggingFaceEmbeddings(model_name=EMBED_MODEL_NAME)
-        logger.info("✅ Embedding Model Loaded.")
-    return _embedding_func
+def get_embedding_func(model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
+    """
+    Dynamically loads the correct embedding model based on the selection.
+    """
+    try:
+        # 1. OpenAI Models
+        if "openai" in model_name.lower():
+            if not os.getenv("OPENAI_API_KEY"):
+                raise ValueError("OpenAI API Key not found.")
+            
+            # Map friendly names to actual API model names if needed
+            # But usually we just pass the exact string like "text-embedding-3-small"
+            return OpenAIEmbeddings(model=model_name)
+            
+        # 2. Hugging Face Models (Local / CPU-friendly)
+        else:
+            # Default to all-MiniLM if something weird is passed, or use the specific HF model
+            return HuggingFaceEmbeddings(model_name=model_name)
+            
+    except Exception as e:
+        logger.error(f"Failed to load embedding model '{model_name}': {e}")
+        # Fallback to the safe default if everything explodes
+        return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 def get_rerank_model():
     """Lazy loads the Cross-Encoder model."""
@@ -150,7 +166,7 @@ def save_uploaded_file(uploaded_file, username: str = "default") -> str:
         logger.error(f"Error saving file: {e}")
         return None
 
-def process_and_add_text(text: str, source_name: str, username: str, index_name: str) -> Tuple[bool, str]:
+def process_and_add_text(text: str, source_name: str, username: str, embed_model_name: str, index_name: str) -> Tuple[bool, str]:
     """Ingests raw text (Flattener) -> Saves Backup to Disk -> Uploads to Pinecone."""
     if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
     
@@ -165,7 +181,7 @@ def process_and_add_text(text: str, source_name: str, username: str, index_name:
 
         # 2. UPLOAD TO PINECONE
         pm = PineconeManager(PINECONE_KEY)
-        emb_fn = get_embedding_func()
+        emb_fn = get_embedding_func(embed_model_name)
         
         # Create Document
         doc = Document(
@@ -182,8 +198,7 @@ def process_and_add_text(text: str, source_name: str, username: str, index_name:
         logger.error(f"Error indexing text: {e}")
         return False, str(e)
 
-def ingest_file(file_path: str, username: str, index_name: str, strategy: str = "paragraph") -> Tuple[bool, str]:
-    """Chunks File -> Scans Acronyms -> Uploads to Pinecone."""
+def ingest_file(file_path: str, username: str, index_name: str, embed_model_name: str, strategy: str = "paragraph") -> Tuple[bool, str]:    """Chunks File -> Scans Acronyms -> Uploads to Pinecone."""
     if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
 
     try:
@@ -198,7 +213,7 @@ def ingest_file(file_path: str, username: str, index_name: str, strategy: str =
 
         # 3. Pinecone Safety Check (Dynamic)
         pm = PineconeManager(PINECONE_KEY)
-        emb_fn = get_embedding_func()
+        emb_fn = get_embedding_func(embed_model_name)
         
         # DYNAMIC CHECK: Generate a test embedding to see true dimension
         # This allows you to swap models in CONFIGURATION later without breaking code
@@ -209,9 +224,10 @@ def ingest_file(file_path: str, username: str, index_name: str, strategy: str =
             return False, f"Dimension Mismatch! Index '{index_name}' expects {model_dim}d vectors (based on current model), but found incompatible dimensions."
 
         # 4. Upload
-        emb_fn = get_embedding_func()
+        emb_fn = get_embedding_func(embed_model_name)
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
-        vstore.add_documents(docs)
+        custom_ids = [f"{doc.metadata.get('source', 'doc')}_{i}" for i, doc in enumerate(docs)]
+        vstore.add_documents(docs, ids=custom_ids)
         
         return True, f"Successfully indexed {len(docs)} chunks."
 
@@ -219,7 +235,7 @@ def ingest_file(file_path: str, username: str, index_name: str, strategy: str =
         logger.error(f"Ingestion failed: {e}")
         return False, str(e)
 
-def search_knowledge_base(query: str, username: str, index_name: str, k: int = 10, final_k: int = 4) -> List[Document]:
+def search_knowledge_base(query: str, username: str, index_name: str, embed_model_name: str, k: int = 10, final_k: int = 4) -> List[Document]:
     """Retrieves from Pinecone -> Reranks."""
     if not PINECONE_KEY or not index_name: return []
     
@@ -230,7 +246,7 @@ def search_knowledge_base(query: str, username: str, index_name: str, k: int = 1
 
         # 2. Vector Search
         pm = PineconeManager(PINECONE_KEY)
-        emb_fn = get_embedding_func()
+        emb_fn = get_embedding_func(embed_model_name)
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
         
         results = vstore.similarity_search(expanded_query, k=k)