Update app.py

app.py

@@ -78,11 +78,35 @@ def embed_pdf(file, collection_name):
     os.remove(file_path)
     return {"message": f"Documents embedded in Weaviate collection '{collection_name}'"}
 
+# Initialize Cohere client
+co = cohere.Client(api_key=cohere_api_key)
+
 def retrieve_info():
     query = request.json.get("query")
     llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
     qa = RetrievalQA.from_chain_type(llm, retriever=vectorstore.as_retriever())
-    result = qa({"query": query})
+    
+    # Retrieve initial results
+    initial_results = qa({"query": query})
+
+    # Assuming initial_results are in the desired format, extract the top 25 documents
+    # Adjust this part according to the actual format of your initial_results
+    top_docs = initial_results[:25]  # Adjust this if your result format is different
+
+    # Rerank the top 25 results
+    reranked_results = co.rerank(query=query, documents=top_docs, top_n=3, model='rerank-english-v2.0')
+
+    # Format the reranked results
+    formatted_results = []
+    for idx, r in enumerate(reranked_results):
+        formatted_result = {
+            "Document Rank": idx + 1,
+            "Document Index": r.index,
+            "Document": r.document['text'],
+            "Relevance Score": f"{r.relevance_score:.2f}"
+        }
+        formatted_results.append(formatted_result)
+        
     return {"results": result}
 
 # Gradio interface
@@ -106,4 +130,4 @@ iface.add_endpoint(
     outputs="text"
 )
 
-iface.launch()
+iface.launch()