Update app.py

app.py

@@ -192,11 +192,23 @@ def query_rag(query_text: str):
     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)
-    model = Ollama(model="llama2")
-    response_text = model.invoke(prompt)
-    return response_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
+    
+    return response.replace("###", '').replace('\nUser:', '')
 
 def save_pdf_and_update_database(media_url):
     response = requests.get(media_url)