Update initialize.py

initialize.py

@@ -22,23 +22,49 @@ 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()
+# 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")
-    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])
+        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
     
     return collection
-    
+
 
 if __name__ == "__main__":
     vectordb = initialize()