Update app.py

app.py

@@ -206,27 +206,49 @@ def initialize_chroma():
 initialize_chroma()
 
 def query_rag(query_text: str):
-    embedding_function = get_embedding_function()
-    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
-    results = db.similarity_search_with_score(query_text, k=5)
-    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
-    
-    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
-    prompt = prompt_template.format(context=context_text, question=query_text)
-
-    response = ''
-    for chunk in AI71(AI71_API_KEY).chat.completions.create(
-            model="tiiuae/falcon-180b-chat",
-            messages=[
-                {"role": "system", "content": "You are the best agricultural assistant. Remember to give a response in not more than 2 sentences."},
-                {"role": "user", "content": f'''Answer the following query based on the given context: {prompt}'''},
-            ],
-            stream=True,
-    ):
-        if chunk.choices[0].delta.content:
-            response += chunk.choices[0].delta.content
+    try:
+        # Ensure query_text is a string
+        if not isinstance(query_text, str):
+            raise ValueError("Query text must be a string.")
+        
+        # Initialize the embedding function and Chroma DB
+        embedding_function = get_embedding_function()
+        db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
+        
+        # Perform similarity search
+        results = db.similarity_search_with_score(query_text, k=5)
+        
+        # Extract and clean context text
+        context_texts = [doc.page_content for doc, _score in results]
+        if not all(isinstance(text, str) for text in context_texts):
+            raise ValueError("All context texts must be strings.")
+        
+        context_text = "\n\n---\n\n".join(context_texts)
+        
+        # Create prompt
+        prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
+        prompt = prompt_template.format(context=context_text, question=query_text)
+        
+        # Generate response using AI71
+        response = ''
+        for chunk in AI71(AI71_API_KEY).chat.completions.create(
+                model="tiiuae/falcon-180b-chat",
+                messages=[
+                    {"role": "system", "content": "You are the best agricultural assistant. Remember to give a response in not more than 2 sentences."},
+                    {"role": "user", "content": f'Answer the following query based on the given context: {prompt}'},
+                ],
+                stream=True,
+        ):
+            if chunk.choices[0].delta.content:
+                response += chunk.choices[0].delta.content
+        
+        # Return cleaned response
+        return response.replace("###", '').replace('\nUser:', '')
     
-    return response.replace("###", '').replace('\nUser:', '')
+    except Exception as e:
+        # Log the error and return a user-friendly message
+        print(f"Error in query_rag: {e}")
+        return "Sorry, there was an error processing your query."
     
 def download_file(url, extension):
     try: