Update app.py

app.py

@@ -1,180 +1,161 @@
 import os
 import logging
-import asyncio
+import json
 import gradio as gr
 from huggingface_hub import InferenceClient
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.schema import Document
 from duckduckgo_search import DDGS
+from dotenv import load_dotenv
+from functools import lru_cache
+
+# Load environment variables
+load_dotenv()
 
 # Configure logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 # Environment variables and configurations
-huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
-logging.info("Environment variable for HuggingFace token retrieved.")
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+logger.info(f"Using Hugging Face token: {HUGGINGFACE_TOKEN[:4]}...{HUGGINGFACE_TOKEN[-4:] if HUGGINGFACE_TOKEN else 'Not Set'}")
 
 MODELS = [
     "mistralai/Mistral-7B-Instruct-v0.3",
     "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "mistralai/Mistral-Nemo-Instruct-2407",
     "meta-llama/Meta-Llama-3.1-8B-Instruct",
-    "meta-llama/Meta-Llama-3.1-70B-Instruct"
+    "meta-llama/Meta-Llama-3.1-70B-Instruct",
+    "google/gemma-2-9b-it",
+    "google/gemma-2-27b-it"
 ]
-logging.info(f"Models list initialized with {len(MODELS)} models.")
 
+FALLBACK_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
+
+DEFAULT_SYSTEM_PROMPT = """You are a world-class financial AI assistant, capable of complex reasoning and reflection. 
+Reason through the query inside <thinking> tags, and then provide your final response inside <output> tags.
+Providing comprehensive and accurate information based on web search results is essential. 
+Your goal is to synthesize the given context into a coherent and detailed response that directly addresses the user's query. 
+Please ensure that your response is well-structured and factual.
+If you detect that you made a mistake in your reasoning at any point, correct yourself inside <reflection> tags."""
+
+class WebSearcher:
+    def __init__(self):
+        self.ddgs = DDGS()
+
+    @lru_cache(maxsize=100)
+    def search(self, query, max_results=5):
+        try:
+            results = list(self.ddgs.text(query, max_results=max_results))
+            logger.info(f"Search completed for query: {query}")
+            return results
+        except Exception as e:
+            logger.error(f"Error during DuckDuckGo search: {str(e)}")
+            return []
+
+@lru_cache(maxsize=1)
 def get_embeddings():
-    logging.info("Loading HuggingFace embeddings model.")
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
-def duckduckgo_search(query):
-    logging.info(f"Initiating DuckDuckGo search for query: {query}")
-    try:
-        with DDGS() as ddgs:
-            results = ddgs.text(query, max_results=10)
-        logging.info(f"Search completed, found {len(results)} results.")
-        return results
-    except Exception as e:
-        logging.error(f"Error during DuckDuckGo search: {str(e)}")
-        return []
-
-async def rephrase_query(query, context, model):
-    # Log the original query for debugging
-    logging.info(f"Original query: {query}")
-
-    prompt = f"""You are a highly intelligent conversational chatbot. Your task is to analyze the given context and new query, then decide whether to rephrase the query with or without incorporating the context. Follow these steps:
-        1. Determine if the new query is a continuation of the previous conversation or an entirely new topic.
-        2. If it's a continuation, rephrase the query by incorporating relevant information from the context to make it more specific and contextual.
-        3. If it's a new topic, rephrase the query to make it more appropriate for a web search, focusing on clarity and accuracy without using the previous context.
-        4. Provide ONLY the rephrased query without any additional explanation or reasoning.
-        
-        Context: {context}
-        
-        New query: {query}
-        
-        Rephrased query:"""
-
-    client = InferenceClient(model, token=huggingface_token)
-
-    try:
-        response = await asyncio.to_thread(
-            client.text_generation,
-            prompt=prompt,
-            max_new_tokens=100,
-            temperature=0.2,
-        )
-
-        # Extract the rephrased query
-        rephrased_query = response.strip()
-
-        # Log the rephrased query
-        logging.info(f"Rephrased query: {rephrased_query}")
-
-        return rephrased_query
-
-    except Exception as e:
-        logging.error(f"Error in rephrase_query: {str(e)}")
-        return query  # Fallback to the original query if there's an error
-
 def create_web_search_vectors(search_results):
-    logging.info(f"Creating web search vectors from {len(search_results)} 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']}))
-    logging.info(f"{len(documents)} documents created for FAISS vectorization.")
+    documents = [
+        Document(
+            page_content=f"{result['title']}\n{result['body']}\nSource: {result['href']}",
+            metadata={"source": result['href']}
+        )
+        for result in search_results if 'body' in result
+    ]
+    logger.info(f"Created vectors for {len(documents)} search results.")
     return FAISS.from_documents(documents, embed)
 
-async def get_response_with_search(query, model, use_embeddings, num_calls=3, temperature=0.2):
-    logging.info(f"Performing web search for query: {query}")
-    search_results = duckduckgo_search(query)
+def get_response_with_search(query, system_prompt, model, use_embeddings, num_calls=3, temperature=0.2):
+    searcher = WebSearcher()
+    search_results = searcher.search(query)
     
     if not search_results:
-        logging.warning("No web search results found.")
-        yield "No web search results available. Please try again.", ""
-        return
+        logger.warning(f"No web search results found for query: {query}")
+        return "No web search results available. Please try again.", ""
+
+    sources = [result['href'] for result in search_results if 'href' in result]
+    source_list_str = "\n".join(sources)
 
     if use_embeddings:
-        logging.info("Using embeddings to retrieve relevant documents.")
         web_search_database = create_web_search_vectors(search_results)
         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])
     else:
-        logging.info("Using raw search results for context.")
-        context = "\n".join([f"{result['title']}\n{result['body']}\nSource: {result['href']}" for result in search_results])
+        context = "\n".join([f"{result['title']}\n{result['body']}" for result in search_results])
 
-    system_message = """ You are a world-class AI system, capable of complex reasoning and reflection. 
-    Reason through the query inside <thinking> tags, and then provide your final response inside <output> tags.
-    Providing comprehensive and accurate information based on web search results is essential. 
-    Your goal is to synthesize the given context into a coherent and detailed response that directly addresses the user's query. 
-    Please ensure that your response is well-structured, factual, and cites sources where appropriate.
-    If you detect that you made a mistake in your reasoning at any point, correct yourself inside <reflection> tags."""
+    logger.info(f"Context created for query: {query}")
 
     user_message = f"""Using the following context from web search results:
 {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."""
+Write a detailed and complete research document that fulfills the following user request: '{query}'."""
 
-    client = InferenceClient(model, token=huggingface_token)
+    client = InferenceClient(model, token=HUGGINGFACE_TOKEN)
     full_response = ""
-    
     try:
         for _ in range(num_calls):
-            logging.info(f"Sending request to model with {num_calls} API calls and temperature {temperature}.")
-            for response in client.chat_completion(
-                messages=[
-                    {"role": "system", "content": system_message},
+            api_params = {
+                "messages": [
+                    {"role": "system", "content": system_prompt},
                     {"role": "user", "content": user_message}
                 ],
-                max_tokens=6000,
-                temperature=temperature,
-                stream=True,
-                top_p=0.8,
-            ):
-                if isinstance(response, dict) and "choices" in response:
-                    for choice in response["choices"]:
-                        if "delta" in choice and "content" in choice["delta"]:
-                            chunk = choice["delta"]["content"]
-                            full_response += chunk
-                            yield full_response, ""
-                else:
-                    logging.error("Unexpected response format or missing attributes in the response object.")
-                    break
+                "max_tokens": 3000,
+                "temperature": temperature,
+                "top_p": 0.8,
+            }
+            logger.info(f"Sending request to API with params: {json.dumps(api_params, indent=2, default=str)}")
+            response = client.chat_completion(**api_params)
+            logger.info(f"Raw response from model: {response}")
+
+            if isinstance(response, dict):
+                if 'generated_text' in response:
+                    full_response += response['generated_text']
+                elif 'choices' in response and len(response['choices']) > 0:
+                    if isinstance(response['choices'][0], dict) and 'message' in response['choices'][0]:
+                        full_response += response['choices'][0]['message'].get('content', '')
+                    elif isinstance(response['choices'][0], str):
+                        full_response += response['choices'][0]
+            elif hasattr(response, 'generated_text'):
+                full_response += response.generated_text
+            elif hasattr(response, 'content'):
+                full_response += response.content
+            else:
+                logger.error(f"Unexpected response format from the model: {type(response)}")
+                return "Unexpected response format from the model. Please try again.", ""
     except Exception as e:
-        logging.error(f"Error in get_response_with_search: {str(e)}")
-        yield f"An error occurred while processing your request: {str(e)}", ""
+        logger.error(f"Error in get_response_with_search: {str(e)}")
+        logger.info(f"Attempting fallback to {FALLBACK_MODEL}")
+        client = InferenceClient(FALLBACK_MODEL, token=HUGGINGFACE_TOKEN)
+        # Retry with fallback model (you can implement retry logic here)
+        return f"An error occurred while processing your request: {str(e)}", ""
 
     if not full_response:
-        logging.warning("No response generated from the model.")
-        yield "No response generated from the model.", ""
+        logger.warning("No response generated from the model")
+        return "No response generated from the model.", ""
+    else:
+        return f"{full_response}\n\nSources:\n{source_list_str}", ""
 
-async def respond(message, history, model, temperature, num_calls, use_embeddings):
-    logging.info(f"User Query: {message}")
-    logging.info(f"Model Used: {model}")
-    logging.info(f"Temperature: {temperature}")
-    logging.info(f"Number of API Calls: {num_calls}")
-    logging.info(f"Use Embeddings: {use_embeddings}")
+def respond(message, system_prompt, history, model, temperature, num_calls, use_embeddings):
+    logger.info(f"Respond function called with message: {message}")
+    logger.info(f"User Query: {message}")
+    logger.info(f"Model Used: {model}")
+    logger.info(f"Temperature: {temperature}")
+    logger.info(f"Number of API Calls: {num_calls}")
+    logger.info(f"Use Embeddings: {use_embeddings}")
+    logger.info(f"System Prompt: {system_prompt}")
 
     try:
-        # Rephrase the query
-        rephrased_query = await rephrase_query(message, history, model)
-        
-        yield f"Rephrased Query: {rephrased_query}\n\nSearching the web...\n\n"
-        
-        async for main_content, sources in get_response_with_search(rephrased_query, model, use_embeddings, num_calls=num_calls, temperature=temperature):
-            response = f"{main_content}\n\n{sources}"
-            yield response
-    except asyncio.CancelledError:
-        logging.warning("Operation cancelled by user.")
-        yield "The operation was cancelled. Please try again."
+        main_content, sources = get_response_with_search(message, system_prompt, model, use_embeddings, num_calls=num_calls, temperature=temperature)
+        return main_content
     except Exception as e:
-        logging.error(f"Error in respond function: {str(e)}")
-        yield f"An error occurred: {str(e)}"
+        logger.error(f"Error in respond function: {str(e)}")
+        return f"An error occurred: {str(e)}"
 
 css = """
 /* Fine-tune chatbox size */
@@ -188,14 +169,14 @@ css = """
 }
 """
 
-# Gradio interface setup
 def create_gradio_interface():
-    logging.info("Setting up Gradio interface.")
     custom_placeholder = "Enter your question here for web search."
 
     demo = gr.ChatInterface(
-        respond,
+        fn=respond,
+        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=True, render=False),
         additional_inputs=[
+            gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, lines=6, label="System Prompt", placeholder="Enter your system prompt here"),
             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"),
             gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
@@ -225,18 +206,17 @@ def create_gradio_interface():
         gr.Markdown("""
         ## How to use
         1. Enter your question in the chat interface.
-        2. Select the model you want to use from the dropdown.
-        3. Adjust the Temperature to control the randomness of the response.
-        4. Set the Number of API Calls to determine how many times the model will be queried.
-        5. Check or uncheck the "Use Embeddings" box to toggle between using embeddings or direct text summarization.
-        6. Press Enter or click the submit button to get your answer.
-        7. Use the provided examples or ask your own questions.
+        2. Optionally, modify the System Prompt to guide the AI's behavior.
+        3. Select the model you want to use from the dropdown.
+        4. Adjust the Temperature to control the randomness of the response.
+        5. Set the Number of API Calls to determine how many times the model will be queried.
+        6. Check or uncheck the "Use Embeddings" box to toggle between using embeddings or direct text summarization.
+        7. Press Enter or click the submit button to get your answer.
+        8. Use the provided examples or ask your own questions.
         """)
 
-    logging.info("Gradio interface ready.")
     return demo
 
 if __name__ == "__main__":
-    logging.info("Launching Gradio application.")
     demo = create_gradio_interface()
     demo.launch(share=True)