Update RAG/Retriever.py

RAG/Retriever.py

@@ -1,8 +1,57 @@
 from langchain_chroma import Chroma
 from langchain_core.vectorstores import VectorStore
-from task1 import LangchainGeminiWrapper #This is from your old task1 file
+#from task1 import LangchainGeminiWrapper #This is from your old task1 file
 import chromadb
 
+from llama_index.embeddings.gemini import GeminiEmbedding
+from typing import List, Dict
+
+import chromadb
+import os
+import pickle
+
+
+# Retrieve API keys from environment variables
+userdata = {
+    "GEMINI_API_KEY":os.getenv("GEMINI_API_KEY"),
+}
+gemini_key = userdata.get("GEMINI_API_KEY")
+parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+# sync
+# Load docs later
+with open('split_docs.pkl', 'rb') as f:
+   docs = pickle.load(f)
+client = chromadb.PersistentClient(path=parent_dir)
+
+# For all subsequent usage:
+class LangchainGeminiWrapper:
+    """
+    Wrapper class to make GeminiEmbedding compatible with Langchain Chroma's interface
+    """
+    def __init__(self, api_key: str, model_name: str = "models/embedding-001"):
+        self.model = GeminiEmbedding(
+            api_key=api_key,
+            model_name=model_name
+        )
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        Embed multiple documents
+        """
+        return [self.model.get_text_embedding(text) for text in texts]
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Embed a single query
+        """
+        return self.model.get_text_embedding(text)
+
+
+
+
+
+
 def load_vector_store(gemini_key: str, persist_directory: str) -> VectorStore:
     gemini_embedder = LangchainGeminiWrapper(api_key=gemini_key)
     return Chroma(