Update app.py

app.py

@@ -29,6 +29,7 @@ import os
 from mistralai import Mistral
 from dotenv import load_dotenv
 import re
+from typing import List, Tuple
 
 # Automatically get the current year
 current_year = datetime.datetime.now().year
@@ -65,6 +66,94 @@ 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:
+    system_prompt = """You are an intelligent agent tasked with determining whether a user query requires a web search or can be answered using the AI's existing knowledge base. Your task is to analyze the query and decide on the appropriate action.
+
+    Instructions:
+    1. If the query is a general conversation starter, greeting, or can be answered without real-time information, classify it as "knowledge_base".
+    2. If the query requires up-to-date information, news, or specific data that might change over time, classify it as "web_search".
+    3. Consider the chat history when making your decision.
+    4. Respond with ONLY "knowledge_base" or "web_search".
+
+    Examples:
+    - "Hi, how are you?" -> "knowledge_base"
+    - "What's the latest news in the US?" -> "web_search"
+    - "Can you explain quantum computing?" -> "knowledge_base"
+    - "What are the current stock prices for Apple?" -> "web_search"
+    """
+
+    user_prompt = f"""
+    Chat history:
+    {chat_history}
+
+    Current query: {query}
+
+    Determine if this query requires a web search or can be answered from the knowledge base.
+    """
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt}
+    ]
+
+    try:
+        response = llm_client.chat_completion(
+            messages=messages,
+            max_tokens=10,
+            temperature=0.2
+        )
+        decision = response.choices[0].message.content.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:
+    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"""
+    Chat history:
+    {chat_history}
+
+    Current query: {query}
+
+    Please provide a response to the query.
+    """
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt}
+    ]
+
+    try:
+        if model == "groq":
+            response = groq_client.chat.completions.create(
+                messages=messages,
+                model="llama-3.1-70b-instant",
+                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()
+    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."
+
+
 # Set up a session with retry mechanism
 def requests_retry_session(
     retries=0,
@@ -609,24 +698,29 @@ 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, history, num_results, max_chars, time_range, language, category, engines, safesearch, method, llm_temperature, model, use_pydf2):
+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):
     chat_history = "\n".join([f"{role}: {msg}" for role, msg in history])
     
-    response = search_and_scrape(
-        query=message,
-        chat_history=chat_history,
-        num_results=num_results,
-        max_chars=max_chars,
-        time_range=time_range,
-        language=language,
-        category=category,
-        engines=engines,
-        safesearch=safesearch,
-        method=method,
-        llm_temperature=llm_temperature,
-        model=model,
-        use_pydf2=use_pydf2
-    )
+    query_type = determine_query_type(message, chat_history, client)
+    
+    if query_type == "knowledge_base":
+        response = generate_ai_response(message, chat_history, client, model)
+    else:  # web_search
+        response = search_and_scrape(
+            query=message,
+            chat_history=chat_history,
+            num_results=num_results,
+            max_chars=max_chars,
+            time_range=time_range,
+            language=language,
+            category=category,
+            engines=engines,
+            safesearch=safesearch,
+            method=method,
+            llm_temperature=llm_temperature,
+            model=model,
+            use_pydf2=use_pydf2
+        )
     
     yield response