Update app.py

app.py

@@ -2,13 +2,15 @@ import weaviate
 import langchain
 import gradio as gr
 from langchain.embeddings import CohereEmbeddings
-from langchain.document_loaders import UnstructuredFileLoader, PyPDFLoader
-from langchain.vectorstores import Qdrant
+from langchain.document_loaders import UnstructuredFileLoader
+from langchain.vectorstores import Weaviate
+from langchain.chain_types import OpenAI, RetrievalQA
 import os
 import urllib.request
 import ssl
 import mimetypes
 from dotenv import load_dotenv
+import cohere
 
 # Load environment variables
 load_dotenv()
@@ -17,19 +19,18 @@ cohere_api_key = os.getenv('COHERE')
 weaviate_api_key = os.getenv('WEAVIATE')
 weaviate_url = os.getenv('WEAVIATE_URL')
 
-
 # Weaviate connection
 auth_config = weaviate.auth.AuthApiKey(api_key=weaviate_api_key)
-client = weaviate.client(url=weaviate_url, auth_client_secret=auth_config, 
+client = weaviate.Client(url=weaviate_url, auth_client_secret=auth_config, 
                          additional_headers={"X-Cohere-Api-Key": cohere_api_key})
+
 # Initialize vectorstore
 vectorstore = Weaviate(client, index_name="Articles", text_key="text")
 vectorstore._query_attrs = ["text", "title", "url", "views", "lang", "_additional {distance}"]
 vectorstore.embedding = CohereEmbeddings(model="embed-multilingual-v2.0", cohere_api_key=cohere_api_key)
 
-# Initialize OpenAI and RetrievalQA
-llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
-qa = RetrievalQA.from_chain_type(llm, retriever=vectorstore.as_retriever())
+# Initialize Cohere client
+co = cohere.Client(api_key=cohere_api_key)
 
 def embed_pdf(file, collection_name):
     # Save the uploaded file
@@ -43,37 +44,8 @@ def embed_pdf(file, collection_name):
     loader = UnstructuredFileLoader(file_path)
     docs = loader.load()
 
-    # Generate embeddings
-    embeddings = CohereEmbeddings(model="embed-multilingual-v2.0", cohere_api_key=cohere_api_key)
-    
-    # Store documents in vectorstore (Qdrant)
-    for doc in docs:
-        embedding = embeddings.embed([doc['text']])
-        vectorstore_document = {
-            "text": doc['text'],
-            "embedding": embedding
-        }
-    collection_name = request.json.get("collection_name")
-    file_url = request.json.get("file_url")
-
-    # Download the file
-    folder_path = f'./'
-    os.makedirs(folder_path, exist_ok=True)
-    filename = file_url.split('/')[-1]
-    file_path = os.path.join(folder_path, filename)
-    
-    ssl._create_default_https_context = ssl._create_unverified_context
-    urllib.request.urlretrieve(file_url, file_path)
-
-    # Check filetype for document parsing
-    mime_type = mimetypes.guess_type(file_path)[0]
-    loader = UnstructuredFileLoader(file_path)
-    docs = loader.load()
-
-    # Generate embeddings
+    # Generate embeddings and store documents in Weaviate
     embeddings = CohereEmbeddings(model="embed-multilingual-v2.0", cohere_api_key=cohere_api_key)
-    
-    # Store documents in Weaviate
     for doc in docs:
         embedding = embeddings.embed([doc['text']])
         weaviate_document = {
@@ -85,22 +57,17 @@ 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")
+def retrieve_info(query):
     llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
     qa = RetrievalQA.from_chain_type(llm, retriever=vectorstore.as_retriever())
     
     # 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
+    # Assuming initial_results are in the desired format, extract the top documents
     top_docs = initial_results[:25]  # Adjust this if your result format is different
 
-    # Rerank the top 25 results
+    # Rerank the top results
     reranked_results = co.rerank(query=query, documents=top_docs, top_n=3, model='rerank-english-v2.0')
 
     # Format the reranked results
@@ -114,14 +81,13 @@ def retrieve_info():
         }
         formatted_results.append(formatted_result)
         
-    return {"results": result}
+    return {"results": formatted_results}
 
 # Gradio interface
 iface = gr.Interface(
     fn=retrieve_info,
     inputs=[
-        gr.inputs.Textbox(label="Query"),
-        gr.inputs.File(label="PDF File", type="file", optional=True)
+        gr.inputs.Textbox(label="Query")
     ],
     outputs="text",
     allow_flagging="never"