Update app.py

app.py

@@ -1,6 +1,5 @@
 import requests
 import gradio as gr
-from bs4 import BeautifulSoup
 import logging
 from urllib.parse import urlparse
 from requests.adapters import HTTPAdapter
@@ -38,10 +37,11 @@ from collections import Counter
 import numpy as np
 from typing import List, Dict, Tuple
 import datetime
-CURRENT_YEAR = datetime.datetime.now().year
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any
 
 # Automatically get the current year
-current_year = datetime.datetime.now().year
+CURRENT_YEAR = datetime.datetime.now().year
 
 # Load environment variables from a .env file
 load_dotenv()
@@ -51,8 +51,42 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 # SearXNG instance details
-SEARXNG_URL = 'https://shreyas094-searxng-hf.hf.space/search'
-SEARXNG_KEY = 'f9f07f93b37b8483aadb5ba717f556f3a4ac507b281b4ca01e6c6288aa3e3ae5'
+SEARXNG_URL = os.getenv("SEARXNG_URL")
+SEARXNG_KEY = os.getenv("SEARXNG_KEY")
+
+
+logger.info(f"SearXNG URL: {SEARXNG_URL}")
+logger.info(f"SearXNG Key: {SEARXNG_KEY}") 
+
+
+# ... other environment variables ...
+CUSTOM_LLM = os.getenv("CUSTOM_LLM")
+CUSTOM_LLM_DEFAULT_MODEL = os.getenv("CUSTOM_LLM_DEFAULT_MODEL")
+
+logger.info(f"CUSTOM_LLM: {CUSTOM_LLM}")
+logger.info(f"CUSTOM_LLM_DEFAULT_MODEL: {CUSTOM_LLM_DEFAULT_MODEL}")
+
+# Define the fetch_custom_models function here
+def fetch_custom_models():
+    if not CUSTOM_LLM:
+        return []
+    try:
+        response = requests.get(f"{CUSTOM_LLM}/v1/models")
+        response.raise_for_status()
+        models = response.json().get("data", [])
+        return [model["id"] for model in models]
+    except Exception as e:
+        logger.error(f"Error fetching custom models: {e}")
+        return []
+
+# Fetch custom models and determine the default model
+custom_models = fetch_custom_models()
+all_models = ["huggingface", "groq", "mistral"] + custom_models
+
+# Determine the default model
+default_model = CUSTOM_LLM_DEFAULT_MODEL if CUSTOM_LLM_DEFAULT_MODEL in all_models else "groq"
+
+logger.info(f"Default model selected: {default_model}")
 
 # Use the environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -75,7 +109,89 @@ mistral_client = Mistral(api_key=MISTRAL_API_KEY)
 similarity_model = SentenceTransformer('all-MiniLM-L6-v2')
 
 
-def determine_query_type(query: str, chat_history: str, llm_client) -> str:
+
+# Step 1: Create a base class for AI models
+class AIModel(ABC):
+    @abstractmethod
+    def generate_response(self, messages: List[Dict[str, str]], max_tokens: int, temperature: float) -> str:
+        pass
+
+# Step 2: Implement specific classes for each AI model
+class HuggingFaceModel(AIModel):
+    def __init__(self, client):
+        self.client = client
+
+    def generate_response(self, messages: List[Dict[str, str]], max_tokens: int, temperature: float) -> str:
+        response = self.client.chat_completion(
+            messages=messages,
+            max_tokens=max_tokens,
+            temperature=temperature
+        )
+        return response.choices[0].message.content.strip()
+
+class GroqModel(AIModel):
+    def __init__(self, client):
+        self.client = client
+
+    def generate_response(self, messages: List[Dict[str, str]], max_tokens: int, temperature: float) -> str:
+        response = self.client.chat.completions.create(
+            messages=messages,
+            model="llama-3.1-70b-versatile",
+            max_tokens=max_tokens,
+            temperature=temperature
+        )
+        return response.choices[0].message.content.strip()
+
+class MistralModel(AIModel):
+    def __init__(self, client):
+        self.client = client
+
+    def generate_response(self, messages: List[Dict[str, str]], max_tokens: int, temperature: float) -> str:
+        response = self.client.chat.complete(
+            model="open-mistral-nemo",
+            messages=messages,
+            max_tokens=max_tokens,
+            temperature=temperature
+        )
+        return response.choices[0].message.content.strip()
+
+# Step 3: Use a factory pattern to create model instances
+class CustomModel(AIModel):
+    def __init__(self, model_name):
+        self.model_name = model_name
+
+    def generate_response(self, messages: List[Dict[str, str]], max_tokens: int, temperature: float) -> str:
+        try:
+            response = requests.post(
+                f"{CUSTOM_LLM}/v1/chat/completions",
+                json={
+                    "model": self.model_name,
+                    "messages": messages,
+                    "max_tokens": max_tokens,
+                    "temperature": temperature
+                }
+            )
+            response.raise_for_status()
+            return response.json()["choices"][0]["message"]["content"].strip()
+        except Exception as e:
+            logger.error(f"Error generating response from custom model: {e}")
+            return "Error: Unable to generate response from custom model."
+
+class AIModelFactory:
+    @staticmethod
+    def create_model(model_name: str, client: Any = None) -> AIModel:
+        if model_name == "huggingface":
+            return HuggingFaceModel(client)
+        elif model_name == "groq":
+            return GroqModel(client)
+        elif model_name == "mistral":
+            return MistralModel(client)
+        elif CUSTOM_LLM and model_name in fetch_custom_models():
+            return CustomModel(model_name)
+        else:
+            raise ValueError(f"Unsupported model: {model_name}")
+
+def determine_query_type(query: str, chat_history: str, ai_model: AIModel) -> str:
     system_prompt = """You are Sentinel, an intelligent AI agent tasked with determining whether a user query requires a web search or can be answered using your existing knowledge base. Your knowledge cutoff date is 2023, and the current year is 2024. Your task is to analyze the query and decide on the appropriate action.
 
     Instructions for Sentinel:
@@ -120,18 +236,18 @@ def determine_query_type(query: str, chat_history: str, llm_client) -> str:
     ]
 
     try:
-        response = llm_client.chat_completion(
+        response = ai_model.generate_response(
             messages=messages,
             max_tokens=10,
             temperature=0.2
         )
-        decision = response.choices[0].message.content.strip().lower()
+        decision = response.strip().lower()
         return "web_search" if decision == "web_search" else "knowledge_base"
     except Exception as e:
         logger.error(f"Error determining query type: {e}")
         return "web_search"  # Default to web search if there's an error
 
-def generate_ai_response(query: str, chat_history: str, llm_client, model: str) -> str:
+def generate_ai_response(query: str, chat_history: str, ai_model: AIModel, temperature: float) -> str:
     system_prompt = """You are a helpful AI assistant. Provide a concise and informative response to the user's query based on your existing knowledge. Do not make up information or claim to have real-time data."""
 
     user_prompt = f"""
@@ -149,29 +265,12 @@ def generate_ai_response(query: str, chat_history: str, llm_client, model: str)
     ]
 
     try:
-        if model == "groq":
-            response = groq_client.chat.completions.create(
-                messages=messages,
-                model="llama-3.1-70b-versatile",
-                max_tokens=500,
-                temperature=0.7
-            )
-            return response.choices[0].message.content.strip()
-        elif model == "mistral":
-            response = mistral_client.chat.complete(
-                model="open-mistral-nemo",
-                messages=messages,
-                max_tokens=500,
-                temperature=0.7
-            )
-            return response.choices[0].message.content.strip()
-        else:  # huggingface
-            response = llm_client.chat_completion(
-                messages=messages,
-                max_tokens=500,
-                temperature=0.7
-            )
-            return response.choices[0].message.content.strip()
+        response = ai_model.generate_response(
+            messages=messages,
+            max_tokens=500,
+            temperature=temperature
+        )
+        return response
     except Exception as e:
         logger.error(f"Error generating AI response: {e}")
         return "I apologize, but I'm having trouble generating a response at the moment. Please try again later."
@@ -282,8 +381,8 @@ def rephrase_query(chat_history, query, temperature=0.2):
    - First, check if the query contains words indicating current information (e.g., "today", "now", "current", "latest"):
      - If present, do NOT add any date operators to the query
    - Otherwise, if the query mentions a specific time period (e.g., a quarter, year, or date range):
-     - Add appropriate "after:" operators to the end of the rephrased query.
-     - Use the format "after:YYYY" for date ranges.
+     - Add appropriate "after: " operators to the end of the rephrased query.
+     - Use the format "after: YYYY" for date ranges.
    - If no specific time period is mentioned and no current-time indicators are present:
      - Append "after: {CURRENT_YEAR}" to the end of the rephrased query.
    - Do not use quotes or the "+" operator when adding dates.
@@ -300,16 +399,20 @@ def rephrase_query(chat_history, query, temperature=0.2):
 
 **Scenario 2: New Topic with Specific Quarter**
 - **User Query**: "How did Bank of America perform during Q2 2024?"
-- **Rephrased Query**: "How did \"Bank of America\" perform during Q2 2024 after:2024"
+- **Rephrased Query**: "How did \"Bank of America\" perform during Q2 2024 after: 2024"
 
 **Scenario 3: Continuation with Date Range**
 - **Previous Query**: "What were Apple's sales figures for 2023?"
 - **User Query**: "How about for the first half of 2024?"
-- **Rephrased Query**: "How about \"Apple\"'s sales figures for the first half of 2024 after:2024"
+- **Rephrased Query**: "How about \"Apple\"'s sales figures for the first half of 2024 after: 2024"
 
 **Scenario 4: Current Status Query**
 - **User Query**: "What is the current market share of Toyota and Honda in the US?"
 - **Rephrased Query**: "What is the current market share of \"Toyota\" and \"Honda\" in the \"US\""
+
+**Scenario 5: Current Status Query**
+- **User Query**: "Bank of America Q2 2024 earnings?"
+- **Rephrased Query**: "\"Bank of America\" Q2 2024 earnings after: 2024""
 """
 
     # Create the user prompt with the chat history and current query
@@ -527,7 +630,7 @@ def is_content_unique(new_content, existing_contents, similarity_threshold=0.8):
     return True
 
 def assess_relevance_and_summarize(llm_client, query, document, temperature=0.2):
-    system_prompt = """You are a world-class AI assistant specializing in financial news analysis. Your task is to assess the relevance of a given document to a user's query and provide a detailed summary if it's relevant."""
+    system_prompt = """You are a world-class AI assistant specializing in news analysis. Your task is to assess the relevance of a given document to a user's query and provide a detailed summary if it's relevant."""
 
     user_prompt = f"""
 Query: {query}
@@ -542,7 +645,7 @@ Instructions:
    - Key facts and figures
    - Dates of events or announcements
    - Names of important entities mentioned
-   - Any financial metrics or changes reported
+   - Any metrics or changes reported
    - The potential impact or significance of the news
 3. If not relevant, simply state "Not relevant".
 
@@ -550,7 +653,7 @@ Your response should be in the following format:
 Relevant: [Yes/No]
 Summary: [Your detailed summary if relevant, or "Not relevant" if not]
 
-Remember to focus on financial aspects and implications in your assessment and summary. Aim to make the summary distinctive, highlighting what makes this particular news item unique compared to similar news.
+Remember to focus on key aspects and implications in your assessment and summary. Aim to make the summary distinctive, highlighting what makes this particular news item unique compared to similar news.
 """
 
     messages = [
@@ -596,7 +699,7 @@ def scrape_full_content(url, max_chars=3000, timeout=5, use_pydf2=True):
         return ""
 
 def llm_summarize(json_input, model, temperature=0.2):
-    system_prompt = """You are Sentinel, a world-class Financial analysis AI model who is expert at searching the web and answering user's queries. You are also an expert at summarizing web pages or documents and searching for content in them."""
+    system_prompt = """You are Sentinel, a world-class AI model who is expert at searching the web and answering user's queries. You are also an expert at summarizing web pages or documents and searching for content in them."""
     user_prompt = f"""
 Please provide a comprehensive summary based on the following JSON input:
 {json_input}
@@ -654,8 +757,23 @@ Instructions:
         logger.error(f"Error in LLM summarization: {e}")
         return "Error: Unable to generate a summary. Please try again."
 
-def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_range="", language="all", category="", 
-                      engines=[], safesearch=2, method="GET", llm_temperature=0.2, timeout=5, model="huggingface", use_pydf2=True):   
+def search_and_scrape(
+    query: str,
+    chat_history: str,
+    ai_model: AIModel,
+    num_results: int = 10,
+    max_chars: int = 1500,
+    time_range: str = "",
+    language: str = "en",
+    category: str = "general",
+    engines: List[str] = [],
+    safesearch: int = 2,
+    method: str = "GET",
+    llm_temperature: float = 0.2,
+    timeout: int = 5,
+    model: str = "huggingface",
+    use_pydf2: bool = True
+):
     try:
         # Step 1: Rephrase the Query
         rephrased_query = rephrase_query(chat_history, query, temperature=llm_temperature)
@@ -801,16 +919,16 @@ def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_r
 
         if not relevant_documents:
             logger.warning("No relevant and unique documents found.")
-            return "No relevant and unique financial news found for the given query."
+            return "No relevant and unique news found for the given query."
 
         # Step 5: Rerank documents based on similarity to query and prioritize entity domain
         reranked_docs = rerank_documents_with_priority(rephrased_query, relevant_documents, entity_domain, similarity_threshold=0.95, max_results=num_results)
         
         if not reranked_docs:
             logger.warning("No documents remained after reranking.")
-            return "No relevant financial news found after filtering and ranking."
+            return "No relevant news found after filtering and ranking."
         
-        logger.info(f"Reranked and filtered to top {len(reranked_docs)} unique, finance-related documents.")
+        logger.info(f"Reranked and filtered to top {len(reranked_docs)} unique, related documents.")
 
         # Step 5: Scrape full content for top documents (up to num_results)
         for doc in reranked_docs[:num_results]:
@@ -839,19 +957,39 @@ def search_and_scrape(query, chat_history, num_results=5, max_chars=3000, time_r
         logger.error(f"Unexpected error in search_and_scrape: {e}")
         return f"An unexpected error occurred during the search and scrape process: {e}"
 
-def chat_function(message: str, history: List[Tuple[str, str]], num_results: int, max_chars: int, time_range: str, language: str, category: str, engines: List[str], safesearch: int, method: str, llm_temperature: float, model: str, use_pydf2: bool):
+# Helper function to get the appropriate client for each model
+def get_client_for_model(model: str) -> Any:
+    if model == "huggingface":
+        return InferenceClient("mistralai/Mistral-Small-Instruct-2409", token=HF_TOKEN)
+    elif model == "groq":
+        return Groq(api_key=GROQ_API_KEY)
+    elif model == "mistral":
+        return Mistral(api_key=MISTRAL_API_KEY)
+    elif CUSTOM_LLM and (model in fetch_custom_models() or model == CUSTOM_LLM_DEFAULT_MODEL):
+        return None  # CustomModel doesn't need a client
+    else:
+        raise ValueError(f"Unsupported model: {model}")
+
+def chat_function(message: str, history: List[Tuple[str, str]], only_web_search: bool, num_results: int, max_chars: int, time_range: str, language: str, category: str, engines: List[str], safesearch: int, method: str, llm_temperature: float, model: str, use_pydf2: bool):
     chat_history = "\n".join([f"{role}: {msg}" for role, msg in history])
     
-    query_type = determine_query_type(message, chat_history, client)
+    # Create the appropriate AI model
+    ai_model = AIModelFactory.create_model(model, get_client_for_model(model))
+    
+    if only_web_search:
+        query_type = "web_search"
+    else:
+        query_type = determine_query_type(message, chat_history, ai_model)
     
     if query_type == "knowledge_base":
-        response = generate_ai_response(message, chat_history, client, model)
+        response = generate_ai_response(message, chat_history, ai_model, llm_temperature)
     else:  # web_search
         gr.Info("Initiating Web Search")
         yield "Request you to sit back and relax until I scrape the web for up-to-date information"
         response = search_and_scrape(
             query=message,
             chat_history=chat_history,
+            ai_model=ai_model,
             num_results=num_results,
             max_chars=max_chars,
             time_range=time_range,
@@ -870,15 +1008,16 @@ def chat_function(message: str, history: List[Tuple[str, str]], num_results: int
 
 iface = gr.ChatInterface(
     chat_function,
-    title="Web Scraper for Financial News with Sentinel AI",
+    title="Web Scraper for News with Sentinel AI",
     description="Ask Sentinel any question. It will search the web for recent information or use its knowledge base as appropriate.",
     theme=gr.Theme.from_hub("allenai/gradio-theme"),
     additional_inputs=[
-        gr.Slider(5, 20, value=10, step=1, label="Number of initial results"),
+        gr.Checkbox(label="Only do web search", value=True),  # Add this line
+        gr.Slider(5, 20, value=3, step=1, label="Number of initial results"),
         gr.Slider(500, 10000, value=1500, step=100, label="Max characters to retrieve"),
-        gr.Dropdown(["", "day", "week", "month", "year"], value="", label="Time Range"),
-        gr.Dropdown(["", "all", "en", "fr", "de", "es", "it", "nl", "pt", "pl", "ru", "zh"], value="", label="Language"),
-        gr.Dropdown(["", "general", "news", "images", "videos", "music", "files", "it", "science", "social media"], value="", label="Category"),
+        gr.Dropdown(["", "day", "week", "month", "year"], value="week", label="Time Range"),
+        gr.Dropdown(["", "all", "en", "fr", "de", "es", "it", "nl", "pt", "pl", "ru", "zh"], value="en", label="Language"),
+        gr.Dropdown(["", "general", "news", "images", "videos", "music", "files", "it", "science", "social media"], value="general", label="Category"),
         gr.Dropdown(
             ["google", "bing", "duckduckgo", "baidu", "yahoo", "qwant", "startpage"],
             multiselect=True,
@@ -886,10 +1025,10 @@ iface = gr.ChatInterface(
             label="Engines"
         ),
         gr.Slider(0, 2, value=2, step=1, label="Safe Search Level"),
-        gr.Radio(["GET", "POST"], value="POST", label="HTTP Method"),
+        gr.Radio(["GET", "POST"], value="GET", label="HTTP Method"),
         gr.Slider(0, 1, value=0.2, step=0.1, label="LLM Temperature"),
-        gr.Dropdown(["huggingface", "groq", "mistral"], value="mistral", label="LLM Model"),
-        gr.Checkbox(label="Use PyPDF2 for PDF scraping", value=False),
+        gr.Dropdown(all_models, value=default_model, label="LLM Model"),
+        gr.Checkbox(label="Use PyPDF2 for PDF scraping", value=True),
     ],
     additional_inputs_accordion=gr.Accordion("⚙️ Advanced Parameters", open=True),
     retry_btn="Retry",
@@ -904,5 +1043,31 @@ iface = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    logger.info("Starting the SearXNG Scraper for Financial News using ChatInterface with Advanced Parameters")
-    iface.launch(share=True)
+    logger.info("Starting the SearXNG Scraper for News using ChatInterface with Advanced Parameters")
+    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+