Update app.py

app.py

@@ -15,6 +15,20 @@ app = Flask(__name__)
 CHROMA_PATH = '/code/chroma_db'
 if not os.path.exists(CHROMA_PATH):
     os.makedirs(CHROMA_PATH)
+def initialize_chroma():
+    try:
+        # Initialize Chroma
+        embedding_function = HuggingFaceEmbeddings()  # Use your desired embedding function
+        db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
+        
+        # Perform an initial operation to ensure the database is correctly initialized
+        db.similarity_search_with_score("test query", k=1)
+        print("Chroma initialized successfully.")
+    except Exception as e:
+        print(f"Error initializing Chroma: {e}")
+
+initialize_chroma()
+
 
 # Set AI71 API key
 AI71_API_KEY = os.environ.get('AI71_API_KEY')
@@ -72,19 +86,28 @@ def generate_response(query, chat_history):
 
 # Query the RAG system
 def query_rag(query_text: str, chat_history):
-    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=HuggingFaceEmbeddings())
-    
-    results = db.similarity_search_with_score(query_text, k=5)
-    
-    if not results:
-        return "Sorry, I couldn't find any relevant information."
-    
-    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
-    
-    prompt = f"Context:\n{context_text}\n\nQuestion:\n{query_text}"
-    response = generate_response(prompt, chat_history)
-    
-    return response
+    try:
+        # Ensure the database is initialized
+        initialize_chroma()
+
+        embedding_function = HuggingFaceEmbeddings()
+        db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
+        
+        results = db.similarity_search_with_score(query_text, k=5)
+        
+        if not results:
+            return "Sorry, I couldn't find any relevant information."
+        
+        context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
+        
+        prompt = f"Context:\n{context_text}\n\nQuestion:\n{query_text}"
+        response = generate_response(prompt, chat_history)
+        
+        return response
+    except Exception as e:
+        print(f"Error querying RAG system: {e}")
+        return "An error occurred while querying the RAG system."
+
 
 # Flask route to handle WhatsApp webhook
 @app.route('/whatsapp', methods=['POST'])