Update app.py

app.py

@@ -18,10 +18,10 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 import shutil
-import time
-from tenacity import retry, stop_after_attempt, wait_exponential
 import requests
-import numpy as np
+from typing import List
+from langchain.embeddings.base import Embeddings
+
 
 
 # Set up basic configuration for logging
@@ -70,61 +70,35 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
     else:
         raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
 
-class HuggingFaceEmbeddings:
-    def __init__(self, api_token):
+class HuggingFaceInferenceAPIEmbeddings(Embeddings):
+    def __init__(self, api_key: str):
         self.api_url = "https://api-inference.huggingface.co/models/dunzhang/stella_en_1.5B_v5"
-        self.headers = {"Authorization": f"Bearer {api_token}"}
-
-    @retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
-    def query(self, payload):
-        response = requests.post(self.api_url, headers=self.headers, json=payload)
-        result = response.json()
-        if 'error' in result and 'is currently loading' in result['error']:
-            raise Exception("Model is still loading")
-        return result
+        self.headers = {"Authorization": f"Bearer {api_key}"}
 
-    def embed_documents(self, texts):
-        max_chunk_length = 512
-        chunked_texts = []
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        embeddings = []
         for text in texts:
-            if len(text) > max_chunk_length:
-                chunks = [text[i:i+max_chunk_length] for i in range(0, len(text), max_chunk_length)]
-                chunked_texts.extend(chunks)
-            else:
-                chunked_texts.append(text)
-
-        batch_size = 8
-        all_embeddings = []
-        for i in range(0, len(chunked_texts), batch_size):
-            batch = chunked_texts[i:i+batch_size]
-            payload = {
-                "inputs": batch,
-                "task": "sentence-similarity"
+            embedding = self.embed_query(text)
+            embeddings.append(embedding)
+        return embeddings
+
+    def embed_query(self, text: str) -> List[float]:
+        payload = {
+            "inputs": {
+                "source_sentence": text,
+                "sentences": [text]
             }
-            response = self.query(payload)
-            if isinstance(response, list):
-                all_embeddings.extend(response)
-            elif isinstance(response, dict) and 'error' in response:
-                raise ValueError(f"API Error: {response['error']}")
-            else:
-                raise ValueError(f"Unexpected response format: {response}")
-
-        final_embeddings = []
-        i = 0
-        for text in texts:
-            num_chunks = max(1, len(text) // max_chunk_length)
-            text_embeddings = all_embeddings[i:i+num_chunks]
-            avg_embedding = np.mean(text_embeddings, axis=0)
-            final_embeddings.append(avg_embedding)
-            i += num_chunks
-
-        return final_embeddings
-
-    def embed_query(self, text):
-        return self.embed_documents([text])[0]
+        }
+        response = requests.post(self.api_url, headers=self.headers, json=payload)
+        result = response.json()
+        if isinstance(result, list) and len(result) > 0:
+            return result[0]
+        else:
+            raise ValueError(f"Unexpected response format: {result}")
 
 def get_embeddings():
-    return HuggingFaceEmbeddings(api_token=huggingface_token)
+    huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
+    return HuggingFaceInferenceAPIEmbeddings(api_key=huggingface_token)
 
 # Add this at the beginning of your script, after imports
 DOCUMENTS_FILE = "uploaded_documents.json"
@@ -151,12 +125,7 @@ def update_vectors(files, parser):
         logging.warning("No files provided for update_vectors")
         return "Please upload at least one PDF file.", display_documents()
     
-    try:
-        embed = get_embeddings()
-    except Exception as e:
-        logging.error(f"Error initializing embeddings: {str(e)}")
-        return f"Error initializing embeddings: {str(e)}", display_documents()
-
+    embed = get_embeddings()
     total_chunks = 0
     
     all_data = []
@@ -187,13 +156,11 @@ def update_vectors(files, parser):
     try:
         if os.path.exists("faiss_database"):
             logging.info("Updating existing FAISS database")
-            database = FAISS.load_local("faiss_database", embed.embed_query, allow_dangerous_deserialization=True)
-            embeddings = embed.embed_documents([doc.page_content for doc in all_data])
-            database.add_embeddings(embeddings, all_data)
+            database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+            database.add_documents(all_data)
         else:
             logging.info("Creating new FAISS database")
-            embeddings = embed.embed_documents([doc.page_content for doc in all_data])
-            database = FAISS.from_embeddings(embeddings, all_data, embed)
+            database = FAISS.from_documents(all_data, embed)
         
         database.save_local("faiss_database")
         logging.info("FAISS database saved")
@@ -483,42 +450,12 @@ After writing the document, please provide a list of sources used in your respon
     if not full_response:
         yield "I apologize, but I couldn't generate a response at this time. Please try again later."
 
-def create_web_search_vectors(search_results):
-    embed = get_embeddings()
-    
-    documents = []
-    for result in search_results:
-        if 'body' in result:
-            content = f"{result['title']}\n{result['body']}\nSource: {result['href']}"
-            documents.append(Document(page_content=content, metadata={"source": result['href']}))
-    
-    return FAISS.from_documents(documents, embed)
-
-def create_web_search_vectors(search_results):
-    embed = get_embeddings()
-    
-    documents = []
-    for result in search_results:
-        if 'body' in result:
-            content = f"{result['title']}\n{result['body']}\nSource: {result['href']}"
-            documents.append(Document(page_content=content, metadata={"source": result['href']}))
-    
-    return FAISS.from_documents(documents, embed)
-
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
-    web_search_database = create_web_search_vectors(search_results)
-    
-    if not web_search_database:
-        yield "No web search results available. Please try again.", ""
-        return
+    context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
+                        for result in search_results if 'body' in result)
     
-    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
-    relevant_docs = retriever.get_relevant_documents(query)
-    
-    context = "\n".join([doc.page_content for doc in relevant_docs])
-    
-    prompt = f"""Using the following context from web search results:
+    prompt = f"""Using the following context:
 {context}
 Write a detailed and complete research document that fulfills the following user request: '{query}'
 After writing the document, please provide a list of sources used in your response."""
@@ -544,7 +481,6 @@ After writing the document, please provide a list of sources used in your respon
                     main_content += chunk
                     yield main_content, ""  # Yield partial main content without sources
 
-
 def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=0.2):
     logging.info(f"Entering get_response_from_pdf with query: {query}, model: {model}, selected_docs: {selected_docs}")