Update app.py

app.py

@@ -33,6 +33,7 @@ import requests
 from duckduckgo_search import DDGS
 import random
 import datetime
+from groq import Groq
 
 # Automatically get the current year
 current_year = datetime.datetime.now().year
@@ -55,6 +56,12 @@ client = InferenceClient(
     token=HF_TOKEN,
 )
 
+# Default API key for examples (replace with a dummy value or leave empty)
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+
+# Initialize Groq client
+groq_client =  Groq(api_key=GROQ_API_KEY)
+
 # Initialize the similarity model
 similarity_model = SentenceTransformer('all-MiniLM-L6-v2')
 
@@ -495,7 +502,7 @@ def scrape_full_content(url, scraper="bs4", max_chars=3000, timeout=5):
         logger.error(f"Error scraping full content from {url}: {e}")
         return ""
 
-def llm_summarize(json_input, llm_client, temperature=0.2):
+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."""
 
     user_prompt = f"""
@@ -522,21 +529,32 @@ Your response should be detailed, informative, accurate, and directly relevant t
     ]
 
     try:
-        response = llm_client.chat_completion(
-            messages=messages,
-            max_tokens=10000,
-            temperature=temperature,
-            frequency_penalty=1.4,
-            top_p=0.9
-        )
-        return response.choices[0].message.content.strip()
+        if model == "groq":
+            response = groq_client.chat.completions.create(
+                messages=messages,
+                model="llama-3.1-8b-instruct",
+                max_tokens=8192,
+                temperature=temperature,
+                top_p=0.9,
+                frequency_penalty=1.2,
+                stream=False
+            )
+            return response.choices[0].message.content.strip()
+        else:
+            response = client.chat_completion(
+                messages=messages,
+                max_tokens=10000,
+                temperature=temperature,
+                frequency_penalty=1.4,
+                top_p=0.9
+            )
+            return response.choices[0].message.content.strip()
     except Exception as e:
         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, scraper="bs4", max_chars=3000, time_range="", language="all", category="", 
-                      engines=[], safesearch=2, method="GET", llm_temperature=0.2, timeout=5, use_duckduckgo=False):
-    
+                      engines=[], safesearch=2, method="GET", llm_temperature=0.2, timeout=5, use_duckduckgo=False, model="huggingface"):   
     try:
         # Step 1: Rephrase the Query
         rephrased_query = rephrase_query(chat_history, query, temperature=llm_temperature)
@@ -710,10 +728,10 @@ def search_and_scrape(query, chat_history, num_results=5, scraper="bs4", max_cha
                 } for doc in reranked_docs[:num_results]
             ]
         }
-    
+
         # Step 6: LLM Summarization
-        llm_summary = llm_summarize(json.dumps(llm_input), client, temperature=llm_temperature)
-    
+        llm_summary = llm_summarize(json.dumps(llm_input), model, temperature=llm_temperature)
+        
         return llm_summary
 
     except Exception as e:
@@ -721,7 +739,7 @@ def search_and_scrape(query, chat_history, num_results=5, scraper="bs4", max_cha
         return f"An unexpected error occurred during the search and scrape process: {e}"
 
 
-def chat_function(message, history, num_results, scraper, max_chars, time_range, language, category, engines, safesearch, method, llm_temperature, use_duckduckgo):
+def chat_function(message, history, num_results, scraper, max_chars, time_range, language, category, engines, safesearch, method, llm_temperature, use_duckduckgo, model):
     chat_history = "\n".join([f"{role}: {msg}" for role, msg in history])
     
     response = search_and_scrape(
@@ -737,7 +755,8 @@ def chat_function(message, history, num_results, scraper, max_chars, time_range,
         safesearch=safesearch,
         method=method,
         llm_temperature=llm_temperature,
-        use_duckduckgo=use_duckduckgo
+        use_duckduckgo=use_duckduckgo,
+        model=model
     )
     
     yield response
@@ -764,6 +783,7 @@ iface = gr.ChatInterface(
         gr.Radio(["GET", "POST"], value="POST", label="HTTP Method"),
         gr.Slider(0, 1, value=0.2, step=0.1, label="LLM Temperature"),
         gr.Checkbox(label="Use DuckDuckGo Search", value=False),
+        gr.Dropdown(["huggingface", "groq"], value="huggingface", label="LLM Model"),
     ],
     additional_inputs_accordion=gr.Accordion("⚙️ Advanced Parameters", open=True),
     retry_btn="Retry",