Create Retriever.py

RAG/Retriever.py

@@ -0,0 +1,22 @@
+from langchain_chroma import Chroma
+from langchain_core.vectorstores import VectorStore
+from task1 import LangchainGeminiWrapper #This is from your old task1 file
+import chromadb
+
+def load_vector_store(gemini_key: str, persist_directory: str) -> VectorStore:
+    gemini_embedder = LangchainGeminiWrapper(api_key=gemini_key)
+    return Chroma(
+        collection_name="example_collection",
+        embedding_function=gemini_embedder,
+        persist_directory=persist_directory
+    )
+
+class Retriever:
+    def __init__(self, vectordb: VectorStore):
+        self.vectordb = vectordb
+
+    def retrieve_documents(self, query: str, k: int = 7) -> str:
+        docs = self.vectordb.similarity_search(query, k=k)
+        return "\nRetrieved documents:\n" + "".join(
+            [f"===== Document {str(i)} =====\n" + doc.page_content for i, doc in enumerate(docs)]
+        )