Update app.py

app.py

@@ -1,48 +1,16 @@
 import os
 import json
-import re
-import gradio as gr
-import requests
-from duckduckgo_search import DDGS
+import logging
+import shutil
+from tempfile import NamedTemporaryFile
 from typing import List
 from pydantic import BaseModel, Field
-from tempfile import NamedTemporaryFile
-from langchain_community.vectorstores import FAISS
-from langchain_core.vectorstores import VectorStore
 from langchain_core.documents import Document
 from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from llama_parse import LlamaParse
-from langchain_core.documents import Document
-from huggingface_hub import InferenceClient
-import inspect
-import logging
-import shutil
-
-
-# Set up basic configuration for logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-# Environment variables and configurations
-huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
-llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
-ACCOUNT_ID = os.environ.get("CLOUDFARE_ACCOUNT_ID")
-API_TOKEN = os.environ.get("CLOUDFLARE_AUTH_TOKEN")
-API_BASE_URL = "https://api.cloudflare.com/client/v4/accounts/a17f03e0f049ccae0c15cdcf3b9737ce/ai/run/"
-
-print(f"ACCOUNT_ID: {ACCOUNT_ID}")
-print(f"CLOUDFLARE_AUTH_TOKEN: {API_TOKEN[:5]}..." if API_TOKEN else "Not set")
-
-MODELS = [
-    "mistralai/Mistral-7B-Instruct-v0.3",
-    "mistralai/Mixtral-8x7B-Instruct-v0.1",
-    "@cf/meta/llama-3.1-8b-instruct",
-    "mistralai/Mistral-Nemo-Instruct-2407",
-    "gpt-4o-mini",
-    "claude-3-haiku",
-    "llama-3.1-70b",
-    "mixtral-8x7b"
-]
+import gradio as gr
 
 # Initialize LlamaParse
 llama_parser = LlamaParse(
@@ -73,7 +41,6 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="avsolatorio/GIST-Embedding-v0")
 
-# Add this at the beginning of your script, after imports
 DOCUMENTS_FILE = "uploaded_documents.json"
 
 def load_documents():
@@ -89,7 +56,6 @@ def save_documents(documents):
 # Replace the global uploaded_documents with this
 uploaded_documents = load_documents()
 
-# Modify the update_vectors function
 def update_vectors(files, parser):
     global uploaded_documents
     logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
@@ -185,316 +151,6 @@ def delete_documents(selected_docs):
     
     return f"Deleted documents: {', '.join(deleted_docs)}", display_documents()
 
-def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temperature=0.2, should_stop=False):
-    print(f"Starting generate_chunked_response with {num_calls} calls")
-    full_response = ""
-    messages = [{"role": "user", "content": prompt}]
-    
-    if model == "@cf/meta/llama-3.1-8b-instruct":
-        # Cloudflare API
-        for i in range(num_calls):
-            print(f"Starting Cloudflare API call {i+1}")
-            if should_stop:
-                print("Stop clicked, breaking loop")
-                break
-            try:
-                response = requests.post(
-                    f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/meta/llama-3.1-8b-instruct",
-                    headers={"Authorization": f"Bearer {API_TOKEN}"},
-                    json={
-                        "stream": true,
-                        "messages": [
-                            {"role": "system", "content": "You are a friendly assistant"},
-                            {"role": "user", "content": prompt}
-                        ],
-                        "max_tokens": max_tokens,
-                        "temperature": temperature
-                    },
-                    stream=true
-                )
-                
-                for line in response.iter_lines():
-                    if should_stop:
-                        print("Stop clicked during streaming, breaking")
-                        break
-                    if line:
-                        try:
-                            json_data = json.loads(line.decode('utf-8').split('data: ')[1])
-                            chunk = json_data['response']
-                            full_response += chunk
-                        except json.JSONDecodeError:
-                            continue
-                print(f"Cloudflare API call {i+1} completed")
-            except Exception as e:
-                print(f"Error in generating response from Cloudflare: {str(e)}")
-    else:
-        # Original Hugging Face API logic
-        client = InferenceClient(model, token=huggingface_token)
-        
-        for i in range(num_calls):
-            print(f"Starting Hugging Face API call {i+1}")
-            if should_stop:
-                print("Stop clicked, breaking loop")
-                break
-            try:
-                for message in client.chat_completion(
-                    messages=messages,
-                    max_tokens=max_tokens,
-                    temperature=temperature,
-                    stream=True,
-                ):
-                    if should_stop:
-                        print("Stop clicked during streaming, breaking")
-                        break
-                    if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                        chunk = message.choices[0].delta.content
-                        full_response += chunk
-                print(f"Hugging Face API call {i+1} completed")
-            except Exception as e:
-                print(f"Error in generating response from Hugging Face: {str(e)}")
-    
-    # Clean up the response
-    clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', full_response, flags=re.DOTALL)
-    clean_response = clean_response.replace("Using the following context:", "").strip()
-    clean_response = clean_response.replace("Using the following context from the PDF documents:", "").strip()
-    
-    # Remove duplicate paragraphs and sentences
-    paragraphs = clean_response.split('\n\n')
-    unique_paragraphs = []
-    for paragraph in paragraphs:
-        if paragraph not in unique_paragraphs:
-            sentences = paragraph.split('. ')
-            unique_sentences = []
-            for sentence in sentences:
-                if sentence not in unique_sentences:
-                    unique_sentences.append(sentence)
-            unique_paragraphs.append('. '.join(unique_sentences))
-    
-    final_response = '\n\n'.join(unique_paragraphs)
-    
-    print(f"Final clean response: {final_response[:100]}...")
-    return final_response
-
-def duckduckgo_search(query):
-    with DDGS() as ddgs:
-        results = ddgs.text(query, max_results=5)
-    return results
-
-class CitingSources(BaseModel):
-    sources: List[str] = Field(
-        ...,
-        description="List of sources to cite. Should be an URL of the source."
-    )
-def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
-    if not message.strip():
-        return "", history
-
-    history = history + [(message, "")]
-
-    try:
-        for response in respond(message, history, model, temperature, num_calls, use_web_search):
-            history[-1] = (message, response)
-            yield history
-    except gr.CancelledError:
-        yield history
-    except Exception as e:
-        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
-        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
-        yield history
-
-def retry_last_response(history, use_web_search, model, temperature, num_calls):
-    if not history:
-        return history
-    
-    last_user_msg = history[-1][0]
-    history = history[:-1]  # Remove the last response
-    
-    return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
-
-def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
-    logging.info(f"User Query: {message}")
-    logging.info(f"Model Used: {model}")
-    logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
-    logging.info(f"Selected Documents: {selected_docs}")
-
-    try:
-        if use_web_search:
-            for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
-                response = f"{main_content}\n\n{sources}"
-                first_line = response.split('\n')[0] if response else ''
-                logging.info(f"Generated Response (first line): {first_line}")
-                yield response
-        else:
-            embed = get_embeddings()
-            if os.path.exists("faiss_database"):
-                database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-                retriever = database.as_retriever(search_kwargs={"k": 20})
-                
-                # Filter relevant documents based on user selection
-                all_relevant_docs = retriever.get_relevant_documents(message)
-                relevant_docs = [doc for doc in all_relevant_docs if doc.metadata["source"] in selected_docs]
-                
-                if not relevant_docs:
-                    logging.info("No relevant information found in the selected documents.")
-                    yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
-                    return
-
-                context_str = "\n".join([doc.page_content for doc in relevant_docs])
-            else:
-                context_str = "No documents available."
-                logging.info("No documents available.")
-                yield "No documents available. Please upload PDF documents to answer questions."
-                return
-            
-            if model in ["gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"]:
-                # Use DuckDuckGo Chat API
-                logging.info(f"Calling DuckDuckGo Chat API with model: {model}")
-                response = chat(message, model=model, timeout=30)
-                yield response
-            else:
-                # Use Hugging Face API or Cloudflare API
-                for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
-                    first_line = partial_response.split('\n')[0] if partial_response else ''
-                    logging.info(f"Generated Response (first line): {first_line}")
-                    yield partial_response
-    except Exception as e:
-        logging.error(f"Error with {model}: {str(e)}")
-        if "microsoft/Phi-3-mini-4k-instruct" in model:
-            logging.info("Falling back to Mistral model due to Phi-3 error")
-            fallback_model = "mistralai/Mistral-7B-Instruct-v0.3"
-            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search, selected_docs)
-        else:
-            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
-
-from duckduckgo_search import DDGS
-
-def chat(keywords: str, model: str, timeout: int = 30) -> str:
-    """Initiates a chat session with DuckDuckGo AI.
-
-    Args:
-        keywords (str): The initial message or question to send to the AI.
-        model (str): The model to use: "gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b".
-        timeout (int): Timeout value for the HTTP client. Defaults to 30.
-
-    Returns:
-        str: The response from the AI.
-    """
-    logging.info(f"Calling DuckDuckGo Chat API with model: {model}")
-
-    try:
-        with DDGS() as ddgs:
-            return ddgs.chat(keywords, model=model, timeout=timeout)
-    except Exception as e:
-        logging.error(f"Error in DuckDuckGo chat: {str(e)}")
-        return "Error in DuckDuckGo chat. Please try again later."
-
-def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
-    headers = {
-        "Authorization": f"Bearer {API_TOKEN}",
-        "Content-Type": "application/json"
-    }
-    model = "@cf/meta/llama-3.1-8b-instruct"
-
-    if search_type == "pdf":
-        instruction = f"""Using the following context from the PDF documents:
-{context}
-Write a detailed and complete response that answers the following user question: '{query}'"""
-    else:  # web search
-        instruction = 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."""
-
-    inputs = [
-        {"role": "system", "content": instruction},
-        {"role": "user", "content": query}
-    ]
-
-    payload = {
-        "messages": inputs,
-        "stream": True,
-        "temperature": temperature,
-        "max_tokens": 32000
-    }
-
-    full_response = ""
-    for i in range(num_calls):
-        try:
-            with requests.post(f"{API_BASE_URL}{model}", headers=headers, json=payload, stream=True) as response:
-                if response.status_code == 200:
-                    for line in response.iter_lines():
-                        if line:
-                            try:
-                                json_response = json.loads(line.decode('utf-8').split('data: ')[1])
-                                if 'response' in json_response:
-                                    chunk = json_response['response']
-                                    full_response += chunk
-                                    yield full_response
-                            except (json.JSONDecodeError, IndexError) as e:
-                                logging.error(f"Error parsing streaming response: {str(e)}")
-                                continue
-                else:
-                    logging.error(f"HTTP Error: {response.status_code}, Response: {response.text}")
-                    yield f"I apologize, but I encountered an HTTP error: {response.status_code}. Please try again later."
-        except Exception as e:
-            logging.error(f"Error in generating response from Cloudflare: {str(e)}")
-            yield f"I apologize, but an error occurred: {str(e)}. Please try again later."
-    
-    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 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
-    
-    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:
-{context}
-You are an expert AI assistant, write a detailed and complete research document that fulfills the following user request: '{query}'
-Base your entire response strictly on the information retrieved from trusted sources. Importantly, only include information that is directly supported by the retrieved content. 
-If any part of the information cannot be verified from the given sources, clearly state that it could not be confirmed.
-After writing the document, please provide a list of sources used in your response."""
-
-    if model == "@cf/meta/llama-3.1-8b-instruct":
-        # Use Cloudflare API
-        for response in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web"):
-            yield response, ""  # Yield streaming response without sources
-    else:
-        # Use Hugging Face API
-        client = InferenceClient(model, token=huggingface_token)
-        
-        main_content = ""
-        for i in range(num_calls):
-            for message in client.chat_completion(
-                messages=[{"role": "user", "content": prompt}],
-                max_tokens=10000,
-                temperature=temperature,
-                stream=True,
-            ):
-                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                    chunk = message.choices[0].delta.content
-                    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}")
     
@@ -565,25 +221,36 @@ Write a detailed and complete response that answers the following user question:
         
         logging.info("Finished generating response")
 
-def vote(data: gr.LikeData):
-    if data.liked:
-        print(f"You upvoted this response: {data.value}")
-    else:
-        print(f"You downvoted this response: {data.value}")
-
-css = """
-/* Fine-tune chatbox size */
-.chatbot-container {
-    height: 600px !important;
-    width: 100% !important;
-}
-.chatbot-container > div {
-    height: 100%;
-    width: 100%;
-}
-"""
-
-uploaded_documents = []
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, selected_docs):
+    if not message.strip():
+        return "", history
+
+    history = history + [(message, "")]
+
+    try:
+        if use_web_search:
+            for response in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+                history[-1] = (message, response)
+                yield history
+        else:
+            for response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
+                history[-1] = (message, response)
+                yield history
+    except gr.CancelledError:
+        yield history
+    except Exception as e:
+        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
+        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
+        yield history
+
+def retry_last_response(history, use_web_search, model, temperature, num_calls, selected_docs):
+    if not history:
+        return history
+    
+    last_user_msg = history[-1][0]
+    history = history[:-1]  # Remove the last response
+    
+    return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls, selected_docs)
 
 def display_documents():
     return gr.CheckboxGroup(
@@ -601,7 +268,7 @@ def initial_conversation():
                 "4. For any queries feel free to reach out @[email protected] or discord - shreyas094\n\n"
                 "To get started, upload some PDFs or ask me a question!")
     ]
-# Add this new function
+
 def refresh_documents():
     global uploaded_documents
     uploaded_documents = load_documents()
@@ -615,7 +282,7 @@ use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
 demo = gr.ChatInterface(
-    respond,
+    chatbot_interface,
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
@@ -652,13 +319,13 @@ demo = gr.ChatInterface(
     cache_examples=False,
     analytics_enabled=False,
     textbox=gr.Textbox(placeholder=custom_placeholder, container=False, scale=7),
-    chatbot = gr.Chatbot(  
-    show_copy_button=True,
-    likeable=True,
-    layout="bubble",
-    height=400,
-    value=initial_conversation()
-)
+    chatbot=gr.Chatbot(
+        show_copy_button=True,
+        likeable=True,
+        layout="bubble",
+        height=400,
+        value=initial_conversation()
+    )
 )
 
 # Add file upload functionality