Update app.py

app.py

@@ -1,11 +1,16 @@
 #using codes from mistralai official cookbook
 import gradio as gr
-from mistralai.client import MistralClient
-from mistralai.models.chat_completion import ChatMessage
+from llama_index.llms import MistralAI
 import numpy as np
 import PyPDF2
 import faiss
 import os
+from llama_index.core import SimpleDirectoryReader
+from llama_index.embeddings import MistralAIEmbedding
+from llama_index import ServiceContext
+from llama_index.core import VectorStoreIndex, StorageContext
+from llama_index.vector_stores.milvus import MilvusVectorStore
+import textwrap
 
 
 mistral_api_key = os.environ.get("API_KEY")
@@ -36,17 +41,40 @@ def rag_pdf(pdfs: list, question: str) -> str:
     text_retrieved = "\n\n".join(retrieved_chunk)
     return text_retrieved
 
+def load_doc(path_list):
+    documents = SimpleDirectoryReader(input_files=path).load_data()
+    print("Document ID:", documents[0].doc_id)
+    vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
+    return index
+
+
+
 def ask_mistral(message: str, history: list):
     messages = []
-    pdfs = message["files"]
+    docs = message["files"]
     for couple in history:
         if type(couple[0]) is tuple:
-            pdfs += couple[0]
+            docs += couple[0][0]
         else:
             messages.append(ChatMessage(role= "user", content = couple[0]))
             messages.append(ChatMessage(role= "assistant", content = couple[1]))
+    if docs:
+        print(docs)
+        index = load_doc(docs)
+        query_engine = index.as_query_engine()
+        response = query_engine.query(message["text"])
+
+        full_response = ""
+        for text in response.response_gen:
+            full_response += chunk.choices[0].delta.content
+            yield full_response        
+            
+        
+
 
-    if pdfs:
+        
         pdfs_extracted = []
         for pdf in pdfs:
             reader = PyPDF2.PdfReader(pdf)