add chroma await

app.py

@@ -19,6 +19,9 @@ import torch.nn.functional as F
 from dotenv import load_dotenv
 from utils import load_env_variables, parse_and_route , escape_special_characters 
 from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name , metadata_prompt 
+import time  
+import httpx  
+  
 
 
 load_dotenv()
@@ -121,12 +124,34 @@ def load_documents(file_path: str, mode: str = "elements"):
     loader = UnstructuredFileLoader(file_path, mode=mode)  
     docs = loader.load()  
     return [doc.page_content for doc in docs]  
-  
+
+def wait_for_chroma_server(host, port, retries=5, delay=5):  
+    url = f"http://{host}:{port}/"  
+    for _ in range(retries):  
+        try:  
+            response = httpx.get(url)  
+            if response.status_code == 200:  
+                print("Chroma server is up and running!")  
+                return True  
+        except httpx.HTTPStatusError as e:  
+            print(f"Attempt to connect to Chroma server failed with status code: {e.response.status_code}")  
+        except httpx.RequestError as e:  
+            print(f"Attempt to connect to Chroma server failed: {e}")  
+        time.sleep(delay)  
+    print("Failed to connect to Chroma server after multiple attempts.")  
+    return False  
+    
+
 def initialize_chroma(collection_name: str, embedding_function: MyEmbeddingFunction):  
-    client = HttpClient(host='localhost', port=8000, settings = Settings(allow_reset=True, anonymized_telemetry=False)) 
-    client.reset()  # resets the database   
+    host = 'localhost'  
+    port = 8000  
+    if not wait_for_chroma_server(host, port):  
+        raise ConnectionError("Could not connect to Chroma server. Ensure it is running.")  
+  
+    client = HttpClient(host=host, port=port, settings=Settings(allow_reset=True, anonymized_telemetry=False))  
+    client.reset() 
     collection = client.create_collection(collection_name)  
-    return client, collection  
+    return client, collection
   
 def add_documents_to_chroma(client, collection, documents: list, embedding_function: MyEmbeddingFunction):  
     for doc in documents:  
@@ -224,5 +249,5 @@ with gr.Blocks() as demo:
             query_button.click(query_documents, inputs=query_input, outputs=query_output)  
   
 if __name__ == "__main__":  
-    os.system("chroma run --host localhost --port 8000")
+    os.system("chroma run --host localhost --port 8000 &")
     demo.launch()