Update app.py

app.py

@@ -206,49 +206,41 @@ def initialize_chroma():
 initialize_chroma()
 
 def query_rag(query_text: str):
-    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
+    embedding_function = get_embedding_function()
+    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
+
+    # Check if the query is related to a PDF
+    if "from pdf" in query_text.lower() or "in pdf" in query_text.lower():
+        # Provide some context about handling PDFs
+        response_text = "I see you're asking about a PDF-related query. Let me check the context from the PDF."
+    else:
+        # Regular RAG functionality
+        response_text = "Your query is not related to PDFs. Please make sure your question is clear."
+
+    results = db.similarity_search_with_score(query_text, k=5)
+    
+    if not results:
+        response_text = "Sorry, I couldn't find any relevant information."
+    else:
+        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)
         
-        # 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}'},
+                    {"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:', '')
-    
-    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."
+
+        response_text = response.replace("###", '').replace('\nUser:', '')
+
+    return response_text
     
 def download_file(url, extension):
     try: