Update initialize.py

initialize.py

@@ -7,9 +7,6 @@ from langchain_ollama import OllamaEmbeddings
 import ollama
 import chromadb
 
-import time
-import httpx
-
 
 GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY')
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
@@ -24,49 +21,13 @@ embedding_model = OpenAIEmbeddings(api_key=OPENAI_API_KEY, model="text-embedding
 # Create Embeddings for Searching the Splits
 persist_directory = './chroma/'
 
-# def initialize():
-#     # splits = gen_splits.gen_splits()
-#     # vectordb = Chroma.from_documents(documents=splits, persist_directory=persist_directory, embedding=embedding_model)
-#     # vectordb.persist()
-
-#     splits = gen_splits.gen_splits()
-#     client = chromadb.Client()
-#     collection = client.create_collection(name="docs")
-#     print(splits)
-#     # store each document in a vector embedding database
-#     for i, d in enumerate(splits):
-#         response = ollama.embeddings(model="mxbai-embed-large", prompt=d.page_content)
-#         embedding = response["embedding"]
-#         collection.add(ids=[str(i)],embeddings=[embedding], documents=[d])
-    
-#     return collection
-
-
 def initialize():
     splits = gen_splits.gen_splits()
-    client = chromadb.Client()
-    collection = client.create_collection(name="docs")
-
-    
-    # store each document in a vector embedding database
-    for i, d in enumerate(splits):
-        success = False
-        attempts = 0
-        max_attempts = 5
-        
-        while not success and attempts < max_attempts:
-            try:
-                response = ollama.embeddings(model="mxbai-embed-large", prompt=d.page_content)
-                embedding = response["embedding"]
-                collection.add(ids=[str(i)], embeddings=[embedding], documents=[d])
-                success = True
-            except httpx.ConnectError as e:
-                attempts += 1
-                print(f"Connection failed (attempt {attempts}): {e}")
-                time.sleep(2)  # retry after waiting for 2 seconds
-    
+    vectordb = Chroma.from_documents(documents=splits, persist_directory=persist_directory, embedding=embedding_model)
+    vectordb.persist()
     return collection
 
 
+
 if __name__ == "__main__":
     vectordb = initialize()