Update app.py

app.py

@@ -18,6 +18,7 @@ import openai
 import numpy as np
 from pinecone.grpc import PineconeGRPC as Pinecone
 import gradio as gr
+import asyncio
 
     
 load_dotenv()
@@ -35,21 +36,7 @@ try:
 except Exception as e:
         print(f"Error connecting to Pinecone: {str(e)}")
 
-# if pinecone_index_name not in pc.list_indexes().names():
-#     pc.create_index(
-#         name=pinecone_index_name, 
-#         dimension=1024, #1024- voyage-law-2,  # '1536' is the dimension for ada-002 embeddings
-#         metric='cosine',
-#         spec=ServerlessSpec(
-#             cloud='aws',
-#             region=pinecone_environment
-#         )
-#     )
-#     print("Pinecone Index provisioned")
-# else:
-#     print("Pinecone Index already provisioned")
-
-#embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
+
 embeddings = VoyageAIEmbeddings(
     voyage_api_key=voyage_api_key, model="voyage-law-2"
 )
@@ -170,6 +157,52 @@ def complete_workflow(query):
     except Exception as e:
         return {"results": [], "total_results": 0}, f"Error in workflow: {str(e)}"
 
+
+async def async_complete_workflow(query):
+    try:
+        search_task = asyncio.to_thread(search_documents, query)
+        context_data = await search_task  # Run search in parallel
+
+        rerank_task = asyncio.to_thread(rerank, query, context_data)
+        reranked = await rerank_task  # Run rerank in parallel
+
+        # Process results
+        context_data = [
+            {
+                'chunk_id': entry['document']['chunk_id'],
+                'doc_id': entry['document']['doc_id'],
+                'title': entry['document']['title'],
+                'text': entry['document']['text'],
+                'page_number': str(entry['document']['page_number']),
+                'score': str(entry['score'])
+            }
+            for entry in reranked.data
+        ]
+
+        # Generate responses
+        output_tasks = [asyncio.to_thread(generate_output, doc["text"], query) for doc in context_data]
+        outputs = await asyncio.gather(*output_tasks)  # Run LLM calls in parallel
+
+        results = {
+            "results": [
+                {
+                    "natural_language_output": outputs[i],
+                    "chunk_id": doc["chunk_id"],
+                    "document_id": doc["doc_id"],
+                    "title": doc["title"],
+                    "text": doc["text"],
+                    "page_number": doc["page_number"],
+                    "score": doc["score"],
+                }
+                for i, doc in enumerate(context_data)
+            ],
+            "total_results": len(context_data)
+        }
+
+        return results
+    except Exception as e:
+        return {"results": [], "total_results": 0}, f"Error in workflow: {str(e)}"
+
 def gradio_app():
     with gr.Blocks(css=".result-output {width: 150%; font-size: 16px; padding: 10px;}") as app:
         gr.Markdown("### Intelligent Document Search Prototype-v0.1.2 ")
@@ -184,7 +217,7 @@ def gradio_app():
             titles_output = gr.Textbox(label="Document Titles", interactive=False)  # New Textbox for Titles
         
         search_btn.click(
-            complete_workflow,
+            async_complete_workflow,
             inputs=user_query,
             outputs=[result_output, titles_output],  
         )