Update app.py

app.py

@@ -16,37 +16,335 @@ from sentence_transformers.util import pytorch_cos_sim
 from enum import Enum
 from groq import Groq
 import os
-from typing import List, Dict, Any, Set
+from typing import List, Dict, Any, Set, Optional
 from dotenv import load_dotenv
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+logger.info("Starting application initialization")
 
 # Load environment variables from .env file
 load_dotenv()
+logger.info("Environment variables loaded")
 
 # Initialize Groq client
 groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
+logger.info("Groq client initialized")
 
 class ScoringMethod(Enum):
     BM25 = "bm25"
     TFIDF = "tfidf"
     COMBINED = "combined"
 
-# First define the SafeSearch enum
 class SafeSearch(Enum):
     STRICT = 2
     MODERATE = 1
     NONE = 0
 
-# Then use it to define the options
+class QueryType(Enum):
+    KNOWLEDGE_BASE = "knowledge_base"
+    WEB_SEARCH = "web_search"
+
 SAFE_SEARCH_OPTIONS = [
     ("Strict (2)", SafeSearch.STRICT.value),
     ("Moderate (1)", SafeSearch.MODERATE.value),
     ("None (0)", SafeSearch.NONE.value)
 ]
 
+async def determine_query_type(query: str, chat_history: List[List[str]], temperature: float = 0.1) -> QueryType:
+    """
+    Determine whether a query should be answered from knowledge base or require web search.
+    Now with improved context handling.
+    """
+    logger.info(f'Determining query type for: {query}')
+    try:
+        # Format chat history into a more natural conversation format
+        formatted_history = []
+        for i, (user_msg, assistant_msg) in enumerate(chat_history[-5:], 1):  # Last 5 turns
+            formatted_history.append(f"Turn {i}:")
+            formatted_history.append(f"User: {user_msg}")
+            if assistant_msg:
+                formatted_history.append(f"Assistant: {assistant_msg}")
+        
+        chat_context = "\n".join(formatted_history)
+
+        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 April 2024, and the current date is November 2024.
+
+Rules for Classification:
+
+1. RESPOND WITH ONLY "knowledge_base" OR "web_search" - NO OTHER TEXT
+
+2. Consider conversation context:
+   - Look for references to previous turns in the conversation
+   - Check if the query is a follow-up to earlier discussion
+   - Consider if previous context requires updated information
+
+3. Classify as "web_search" if:
+   - Query explicitly asks for current/latest/recent information
+   - References events or data after April 2024
+   - Requires real-time information (prices, weather, news)
+   - Uses words like "current", "latest", "now", "today"
+   - Asks about ongoing events or situations
+   - Needs verification of recent claims
+   - Is a follow-up question about current events
+   - Previous context involves recent/ongoing events
+
+4. Classify as "knowledge_base" if:
+   - Query is about historical events or facts before April 2024
+   - Involves general knowledge, concepts, or theories
+   - Is casual conversation or greeting
+   - Asks for explanations of established topics
+   - Requires logical reasoning or analysis
+   - Is about personal opinions or hypotheticals
+   - Is a follow-up to a knowledge-base discussion
+   - Previous context is about historical or conceptual topics"""
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"Previous conversation:\n{chat_context}\n\nCurrent query: {query}\n\nClassify this query based on the rules above, considering the conversation context."}
+        ]
+
+        response = groq_client.chat.completions.create(
+            messages=messages,
+            model="llama-3.1-70b-versatile",
+            temperature=temperature,
+            max_tokens=10,
+            stream=False
+        )
+
+        result = response.choices[0].message.content.strip().lower()
+        logger.info(f'Query type determined as: {result} with context')
+        
+        return QueryType.WEB_SEARCH if result == "web_search" else QueryType.KNOWLEDGE_BASE
+
+    except Exception as e:
+        logger.error(f'Error determining query type: {e}')
+        return QueryType.WEB_SEARCH
+
+async def process_knowledge_base_query(query: str, chat_history: List[List[str]], temperature: float = 0.7) -> str:
+    """Handle queries that can be answered from the knowledge base, with context."""
+    logger.info(f'Processing knowledge base query: {query}')
+    try:
+        # Format recent conversation history
+        formatted_history = []
+        for i, (user_msg, assistant_msg) in enumerate(chat_history[-5:], 1):
+            formatted_history.append(f"Turn {i}:")
+            formatted_history.append(f"User: {user_msg}")
+            if assistant_msg:
+                formatted_history.append(f"Assistant: {assistant_msg}")
+        
+        chat_context = "\n".join(formatted_history)
+
+        system_prompt = """You are Sentinel, a highly knowledgeable AI assistant with expertise through April 2024. You provide accurate, informative responses based on your knowledge base while maintaining conversation context.
+
+Guidelines:
+1. Use the conversation history to provide contextually relevant responses
+2. Reference previous turns when appropriate
+3. Maintain consistency with previous responses
+4. Use markdown formatting for better readability
+5. Be clear about historical facts vs. analysis
+6. Note if information might be outdated
+7. Stay within knowledge cutoff date of April 2024
+8. Be direct and conversational
+9. Acknowledge and build upon previous context when relevant"""
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"Previous conversation:\n{chat_context}\n\nCurrent query: {query}\n\nProvide a comprehensive response based on your knowledge base and the conversation context."}
+        ]
+
+        response = groq_client.chat.completions.create(
+            messages=messages,
+            model="llama-3.1-70b-versatile",
+            temperature=temperature,
+            max_tokens=2000,
+            stream=False
+        )
+
+        return response.choices[0].message.content.strip()
+
+    except Exception as e:
+        logger.error(f'Error processing knowledge base query: {e}')
+        return f"I apologize, but I encountered an error while processing your query: {str(e)}"
+
+async def rephrase_query(chat_history, query, temperature=0.2) -> str:
+    """Rephrase the query based on chat history and context."""
+    logger.info(f'Rephrasing query: {query}')
+    try:
+        # Format chat history for context
+        formatted_history = []
+        for user_msg, assistant_msg in chat_history:
+            formatted_history.append({"role": "user", "content": user_msg})
+            if assistant_msg:  # Only add if there's an assistant message
+                formatted_history.append({"role": "assistant", "content": assistant_msg})
+        
+        current_year = datetime.now().year 
+        system_prompt = """You are a highly intelligent and context-aware query rephrasing assistant. Your task is to rephrase search queries while following these strict rules:
+
+1. Entity Handling:
+   - Identify main entities (organizations, brands, products, locations)
+   - Enclose ONLY the entity names in double quotes
+   - Example: "Apple" stock price, not "Apple stock price"
+
+2. Date Handling Rules (VERY IMPORTANT):
+   - For queries about current/latest/recent information:
+     * If query contains words like "latest", "current", "recent", "now", "today":
+       - Keep these words in the query
+       - ALWAYS append "after: YYYY" (current year) at the end
+     * Example: "latest news on "Apple"" becomes "latest news on "Apple" after: 2024"
+   
+   - For queries with specific time periods:
+     * Keep the original time reference
+     * Add appropriate "after: YYYY" based on the mentioned year
+     * Example: "How did "Bank of America" perform in Q2 2023" becomes
+       "How did "Bank of America" perform in Q2 2023 after: 2023"
+
+   - For queries without any time reference:
+     * ALWAYS append "after: YYYY" (current year) at the end
+     * Example: ""Toyota" market share" becomes ""Toyota" market share after: 2024"
+
+3. Output Format:
+   - First letter should be capitalized
+   - No period at the end
+   - Include all specified date operators
+   - Maintain the entire original query's meaning and context
+
+Remember: EVERY query must end with a date operator unless it explicitly references a past date/year."""
+
+        # Prepare messages for the API call
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"Current year is {current_year}. Rephrase this query: {query}"}
+        ]
+
+        # Call Groq API
+        response = groq_client.chat.completions.create(
+            messages=messages,
+            model="llama-3.1-70b-versatile",
+            temperature=temperature,
+            max_tokens=200,
+            stream=False
+        )
+
+        rephrased_query = response.choices[0].message.content.strip()
+        logger.info(f'Query rephrased to: {rephrased_query}')
+        return rephrased_query
+
+    except Exception as e:
+        logger.error(f'Error rephrasing query: {e}')
+        return query  # Return original query if rephrasing fails
+
+class ParallelScraper:
+    def __init__(self, max_workers: int = 5):
+        logger.info(f"Initializing ParallelScraper with {max_workers} workers")
+        self.executor = ThreadPoolExecutor(max_workers=max_workers)
+        self.session: Optional[aiohttp.ClientSession] = None
+
+    async def __aenter__(self):
+        logger.info("Creating aiohttp session")
+        self.session = aiohttp.ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.session:
+            logger.info("Closing aiohttp session")
+            await self.session.close()
+
+    def parse_article(self, article: Article) -> Dict[str, Any]:
+        """Parse a newspaper Article object in a separate thread"""
+        try:
+            logger.info("Parsing article")
+            article.parse()
+            return {
+                "content": article.text,
+                "publish_date": article.publish_date.isoformat() if article.publish_date else None
+            }
+        except Exception as e:
+            logger.error(f'Error parsing article: {e}')
+            return None
+
+    async def download_and_parse_html(self, url: str, max_chars: int) -> Dict[str, Any]:
+        """Download and parse HTML content asynchronously"""
+        logger.info(f'Processing HTML URL: {url}')
+        try:
+            article = Article(url)
+            await asyncio.get_event_loop().run_in_executor(self.executor, article.download)
+            result = await asyncio.get_event_loop().run_in_executor(self.executor, self.parse_article, article)
+            
+            if result:
+                result["content"] = result["content"][:max_chars]
+                logger.info(f'Successfully processed HTML from {url}')
+            return result
+        except Exception as e:
+            logger.error(f'Error processing HTML from {url}: {e}')
+            return None
+
+    async def download_and_parse_pdf(self, url: str, max_chars: int) -> Dict[str, Any]:
+        """Download and parse PDF content asynchronously"""
+        logger.info(f'Processing PDF URL: {url}')
+        try:
+            if not self.session:
+                raise RuntimeError("Session not initialized")
+                
+            async with self.session.get(url) as response:
+                pdf_bytes = await response.read()
+                
+            def process_pdf():
+                logger.info("Processing PDF content")
+                pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_bytes))
+                text = ""
+                for page in pdf_reader.pages:
+                    text += page.extract_text()
+                return text[:max_chars]
+                
+            text = await asyncio.get_event_loop().run_in_executor(self.executor, process_pdf)
+            logger.info(f'Successfully processed PDF from {url}')
+            return {"content": text, "publish_date": None}
+        except Exception as e:
+            logger.error(f'Error processing PDF from {url}: {e}')
+            return None
+
+    async def scrape_url(self, url: str, max_chars: int) -> Dict[str, Any]:
+        """Scrape content from a URL, handling both HTML and PDF formats"""
+        logger.info(f'Starting to scrape URL: {url}')
+        if url.endswith('.pdf'):
+            return await self.download_and_parse_pdf(url, max_chars)
+        else:
+            return await self.download_and_parse_html(url, max_chars)
+
+    async def scrape_urls(self, urls: list, max_chars: int) -> list:
+        """Scrape multiple URLs in parallel"""
+        logger.info(f'Starting parallel scraping of {len(urls)} URLs')
+        tasks = [self.scrape_url(url, max_chars) for url in urls]
+        return await asyncio.gather(*tasks)
+
+async def scrape_urls_parallel(results: list, max_chars: int) -> list:
+    """Scrape multiple URLs in parallel using the ParallelScraper"""
+    logger.info(f'Initializing parallel scraping for {len(results)} results')
+    async with ParallelScraper() as scraper:
+        urls = [result["url"] for result in results]
+        scraped_data = await scraper.scrape_urls(urls, max_chars)
+        
+        # Combine results with scraped data
+        valid_results = []
+        for result, article in zip(results, scraped_data):
+            if article is not None:
+                valid_results.append((result, article))
+        
+        logger.info(f'Successfully scraped {len(valid_results)} valid results')
+        return valid_results
+
 async def get_available_engines(session, base_url, headers):
     """Fetch available search engines from SearxNG instance."""
+    logger.info("Fetching available search engines")
     try:
-        # First try the search endpoint to get engines
         params = {
             "q": "test",
             "format": "json",
@@ -55,84 +353,22 @@ async def get_available_engines(session, base_url, headers):
         async with session.get(f"{base_url}/search", headers=headers, params=params) as response:
             data = await response.json()
             available_engines = set()
-            # Extract unique engine names from the response
             if "search" in data:
                 for engine_data in data["search"]:
                     if isinstance(engine_data, dict) and "engine" in engine_data:
                         available_engines.add(engine_data["engine"])
             
-            # If no engines found, try alternate endpoint
             if not available_engines:
                 async with session.get(f"{base_url}/engines", headers=headers) as response:
                     engines_data = await response.json()
                     available_engines = set(engine["name"] for engine in engines_data if engine.get("enabled", True))
             
+            logger.info(f'Found {len(available_engines)} available engines')
             return list(available_engines)
     except Exception as e:
-        logging.error(f'Error fetching search engines: {e}')
-        # Return default engines if unable to fetch
+        logger.error(f'Error fetching search engines: {e}')
         return ["google", "bing", "duckduckgo", "brave", "wikipedia"]
 
-def select_search_engines(available_engines: List[str]) -> Set[str]:
-    """Let user select search engines from available options."""
-    print("\nAvailable search engines:")
-    engines_list = sorted(available_engines)
-    for i, engine in enumerate(engines_list, 1):
-        print(f"{i}. {engine}")
-    
-    print("\nEnter the numbers of engines you want to use (comma-separated), or 'all' for all engines:")
-    selection = input("Your selection: ").strip().lower()
-    
-    if selection == 'all':
-        return set(engines_list)
-    
-    try:
-        selected_indices = [int(idx.strip()) - 1 for idx in selection.split(',')]
-        return {engines_list[idx] for idx in selected_indices if 0 <= idx < len(engines_list)}
-    except (ValueError, IndexError):
-        logging.error("Invalid selection, using all engines as fallback")
-        return set(engines_list)
-
-
-logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s: %(message)s')
-
-async def scrape_url(url, max_chars):
-    logging.info(f'Scraping URL: {url}')
-    if url.endswith(".pdf"):
-        return await scrape_pdf(url, max_chars)
-    else:
-        return await scrape_html(url, max_chars)
-
-async def scrape_html(url, max_chars):
-    try:
-        article = Article(url)
-        article.download()
-        article.parse()
-        text = article.text[:max_chars]
-        publish_date = article.publish_date
-        logging.info(f'Scraped HTML content from {url}')
-        return {"content": text, "publish_date": publish_date.isoformat() if publish_date else None}
-    except Exception as e:
-        logging.error(f'Error scraping HTML content from {url}: {e}')
-        return None
-
-async def scrape_pdf(url, max_chars):
-    try:
-        async with aiohttp.ClientSession() as session:
-            async with session.get(url) as response:
-                pdf_bytes = await response.read()
-        pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_bytes))
-        text = ""
-        for page_num in range(len(pdf_reader.pages)):
-            page = pdf_reader.pages[page_num]
-            text += page.extract_text()
-        text = text[:max_chars]
-        logging.info(f'Scraped PDF content from {url}')
-        return {"content": text, "publish_date": None}
-    except Exception as e:
-        logging.error(f'Error scraping PDF content from {url}: {e}')
-        return None
-
 def normalize_scores(scores):
     """Normalize scores to [0, 1] range using min-max normalization"""
     if not isinstance(scores, np.ndarray):
@@ -153,41 +389,46 @@ def normalize_scores(scores):
 
 async def calculate_bm25(query, documents):
     """Calculate BM25 scores for documents."""
+    logger.info("Calculating BM25 scores")
     try:
         if not documents:
             return []
             
         bm25 = BM25Okapi([doc.split() for doc in documents])
         scores = bm25.get_scores(query.split())
-        return normalize_scores(scores)
+        normalized_scores = normalize_scores(scores)
+        logger.info("BM25 scores calculated successfully")
+        return normalized_scores
         
     except Exception as e:
-        logging.error(f'Error calculating BM25 scores: {e}')
+        logger.error(f'Error calculating BM25 scores: {e}')
         return [0] * len(documents)
 
 async def calculate_tfidf(query, documents, measure="cosine"):
     """Calculate TF-IDF based similarity scores."""
+    logger.info("Calculating TF-IDF scores")
     try:
         if not documents:
             return []
             
-        model = SentenceTransformer('BAAI/bge-base-en-v1.5')
+        model = SentenceTransformer('all-MiniLM-L6-v2')
+        logger.info("Encoding query and documents")
         query_embedding = model.encode(query)
         document_embeddings = model.encode(documents)
         
-        # Normalize embeddings
         query_embedding = query_embedding / np.linalg.norm(query_embedding)
         document_embeddings = document_embeddings / np.linalg.norm(document_embeddings, axis=1)[:, np.newaxis]
 
         if measure == "cosine":
-            # Calculate cosine similarity
             scores = np.dot(document_embeddings, query_embedding)
-            return normalize_scores(scores)
+            normalized_scores = normalize_scores(scores)
+            logger.info("TF-IDF scores calculated successfully")
+            return normalized_scores
         else:
             raise ValueError("Unsupported similarity measure.")
             
     except Exception as e:
-        logging.error(f'Error calculating TF-IDF scores: {e}')
+        logger.error(f'Error calculating TF-IDF scores: {e}')
         return [0] * len(documents)
 
 def combine_scores(bm25_score, tfidf_score, weights=(0.5, 0.5)):
@@ -221,11 +462,9 @@ def get_total_score(scores, scoring_method: ScoringMethod):
         return combine_scores(bm25_score, tfidf_score)
 
 async def generate_summary(query: str, articles: List[Dict[str, Any]], temperature: float = 0.7) -> str:
-    """
-    Generate a summary of the articles using Groq's LLama 3.1 8b model.
-    """
+    """Generate a summary of the articles using Groq's LLama 3.1 8b model."""
+    logger.info(f'Generating summary for query: {query}')
     try:
-        # Format the articles into a structured JSON string
         json_input = json.dumps(articles, indent=2)
         
         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."""
@@ -251,6 +490,7 @@ Instructions:
 12. Make sure the answer is not short and is informative.
 13. Your response should be detailed, informative, accurate, and directly relevant to the user's query."""
 
+        logger.info("Sending request to Groq API")
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": user_prompt}
@@ -258,7 +498,7 @@ Instructions:
 
         response = groq_client.chat.completions.create(
             messages=messages,
-            model="llama-3.1-70b-versatile",  # Using LLama 3.1 8b model
+            model="llama-3.1-70b-versatile",
             max_tokens=5000,
             temperature=temperature,
             top_p=0.9,
@@ -266,21 +506,23 @@ Instructions:
             stream=False
         )
         
+        logger.info("Summary generated successfully")
         return response.choices[0].message.content.strip()
         
     except Exception as e:
-        logging.error(f'Error generating summary: {e}')
+        logger.error(f'Error generating summary: {e}')
         return f"Error generating summary: {str(e)}"
 
 class ChatBot:
     def __init__(self):
+        logger.info("Initializing ChatBot")
         self.scoring_method = ScoringMethod.COMBINED
         self.num_results = 10
         self.max_chars = 10000
         self.score_threshold = 0.8
         self.temperature = 0.1
-        self.history = []
-        self.base_url = "https://shreyas094-searxng-local.hf.space\search"
+        self.conversation_history = []
+        self.base_url = "https://shreyas094-searxng-local.hf.space"
         self.headers = {
             "X-Searx-API-Key": "f9f07f93b37b8483aadb5ba717f556f3a4ac507b281b4ca01e6c6288aa3e3ae5"
         }
@@ -298,35 +540,52 @@ class ChatBot:
             "ja": "Japanese",
             "ko": "Korean"
         }
+        logger.info("ChatBot initialized successfully")
+
+    def format_chat_history(self, history: List[List[str]]) -> str:
+        """Format chat history into a readable string with clear turn markers."""
+        formatted_history = []
+        for i, (user_msg, assistant_msg) in enumerate(history, 1):
+            formatted_history.append(f"Turn {i}:")
+            formatted_history.append(f"User: {user_msg}")
+            if assistant_msg:
+                formatted_history.append(f"Assistant: {assistant_msg}")
+        return "\n".join(formatted_history)
 
     async def get_search_results(self, 
                                query: str,
+                               history: List[List[str]],
                                num_results: int,
                                max_chars: int,
                                score_threshold: float,
                                temperature: float,
-                               scoring_method_str: str,
+                               scoring_method: str,
                                selected_engines: List[str],
                                safe_search: str,
                                language: str) -> str:
+        logger.info(f'Processing search request for query: {query}')
         try:
+            # First, rephrase the query using chat history
+            rephrased_query = await rephrase_query(history, query, temperature=0.2)
+            logger.info(f'Original query: {query}')
+            logger.info(f'Rephrased query: {rephrased_query}')
+
             scoring_method_map = {
                 "BM25": ScoringMethod.BM25,
                 "TF-IDF": ScoringMethod.TFIDF,
                 "Combined": ScoringMethod.COMBINED
             }
-            self.scoring_method = scoring_method_map[scoring_method_str]
+            self.scoring_method = scoring_method_map[scoring_method]
             
             safe_search_map = dict(SAFE_SEARCH_OPTIONS)
             safe_search_value = safe_search_map.get(safe_search, SafeSearch.MODERATE.value)
 
+            logger.info(f'Search parameters - Engines: {selected_engines}, Results: {num_results}, Method: {scoring_method}')
+            
+            # Use the rephrased query for the search
             async with aiohttp.ClientSession() as session:
-                logging.info(f'Using engines: {", ".join(selected_engines)}')
-                logging.info(f'Parameters: Results={num_results}, Chars={max_chars}, Threshold={score_threshold}, '
-                           f'Temp={temperature}, Method={scoring_method_str}, SafeSearch={safe_search_value}, Language={language}')
-                
                 params = {
-                    "q": query,
+                    "q": rephrased_query,  # Use rephrased query here
                     "format": "json",
                     "engines": ",".join(selected_engines),
                     "limit": num_results,
@@ -336,24 +595,30 @@ class ChatBot:
                 if language != "all":
                     params["language"] = language
                 
+                logger.info("Sending search request to SearxNG")
                 try:
                     async with session.get(f"{self.base_url}/search", headers=self.headers, params=params) as response:
                         data = await response.json()
                 except Exception as e:
+                    logger.error(f'SearxNG connection error: {e}')
                     return f"Error: Could not connect to search service. Please check if SearxNG is running at {self.base_url}. Error: {str(e)}"
 
                 if "results" not in data or not data["results"]:
+                    logger.info("No search results found")
                     return "No results found."
 
                 results = data["results"][:num_results]
+                logger.info(f'Processing {len(results)} search results')
                 valid_results = await scrape_urls_parallel(results, max_chars)
                 
                 if not valid_results:
+                    logger.info("No valid articles found after scraping")
                     return "No valid articles found after scraping."
 
                 results, scraped_data = zip(*valid_results)
                 contents = [article["content"] for article in scraped_data]
                 
+                logger.info("Calculating document scores")
                 scores = await get_document_scores(query, contents, self.scoring_method)
 
                 scored_articles = []
@@ -381,15 +646,15 @@ class ChatBot:
                         unique_articles.append(article)
 
                 # Generate summary using Groq API
-                summary = await generate_summary(query, unique_articles, self.temperature)
+                summary = await generate_summary(query, unique_articles, temperature)
 
-                # Update the response format to include new parameters
+                # Update the response format to use scoring_method instead of scoring_method_str
                 response = f"**Search Parameters:**\n"
                 response += f"- Results: {num_results}\n"
                 response += f"- Max Characters: {max_chars}\n"
                 response += f"- Score Threshold: {score_threshold}\n"
                 response += f"- Temperature: {temperature}\n"
-                response += f"- Scoring Method: {scoring_method_str}\n"
+                response += f"- Scoring Method: {scoring_method}\n"  # Updated this line
                 response += f"- Search Engines: {', '.join(selected_engines)}\n"
                 response += f"- Safe Search: Level {safe_search_value}\n"
                 response += f"- Language: {self.available_languages.get(language, language)}\n\n"
@@ -404,9 +669,65 @@ class ChatBot:
                 return response
 
         except Exception as e:
-            logging.error(f'Error in search_and_summarize: {e}')
+            logger.error(f'Error in search_and_summarize: {e}')
             return f"Error occurred: {str(e)}"
 
+    async def get_response(self,
+                          query: str,
+                          history: List[List[str]],
+                          num_results: int,
+                          max_chars: int,
+                          score_threshold: float,
+                          temperature: float,
+                          scoring_method: str,
+                          selected_engines: List[str],
+                          safe_search: str,
+                          language: str,
+                          force_web_search: bool = False) -> str:
+        """Determine query type and route to appropriate handler with context."""
+        logger.info(f'Processing query: {query}')
+        try:
+            # Update conversation history
+            formatted_history = self.format_chat_history(history)
+            logger.info(f'Current conversation context:\n{formatted_history}')
+
+            # If force_web_search is True, skip query type determination
+            if force_web_search:
+                logger.info('Force web search mode enabled - bypassing query type determination')
+                query_type = QueryType.WEB_SEARCH
+            else:
+                # Determine query type with context
+                query_type = await determine_query_type(query, history, temperature)
+            
+            if query_type == QueryType.KNOWLEDGE_BASE and not force_web_search:
+                logger.info('Using knowledge base to answer query')
+                response = await process_knowledge_base_query(
+                    query=query,
+                    chat_history=history,
+                    temperature=temperature
+                )
+            else:
+                logger.info('Using web search to answer query')
+                response = await self.get_search_results(
+                    query=query,
+                    history=history,
+                    num_results=num_results,
+                    max_chars=max_chars,
+                    score_threshold=score_threshold,
+                    temperature=temperature,
+                    scoring_method=scoring_method,
+                    selected_engines=selected_engines,
+                    safe_search=safe_search,
+                    language=language
+                )
+
+            logger.info(f'Generated response type: {query_type}')
+            return response
+
+        except Exception as e:
+            logger.error(f'Error in get_response: {e}')
+            return f"I apologize, but I encountered an error: {str(e)}"
+
     def chat(self, 
              message: str, 
              history: List[List[str]], 
@@ -417,15 +738,18 @@ class ChatBot:
              scoring_method: str,
              engines: List[str],
              safe_search: str,
-             language: str) -> str:
-        """
-        Process chat messages and return responses with custom parameters.
-        """
-        # Extract language code from the selection (e.g., "en - English" -> "en")
+             language: str,
+             force_web_search: bool) -> str:
+        """Process chat messages with context and return responses."""
+        # Extract language code and process response
         language_code = language.split(" - ")[0]
         
-        response = asyncio.run(self.get_search_results(
+        # Update conversation history from the Gradio history
+        self.conversation_history = history
+        
+        response = asyncio.run(self.get_response(
             message,
+            self.conversation_history,
             num_results,
             max_chars,
             score_threshold,
@@ -433,7 +757,8 @@ class ChatBot:
             scoring_method,
             engines,
             safe_search,
-            language_code
+            language_code,
+            force_web_search
         ))
         return response
 
@@ -442,17 +767,7 @@ def create_gradio_interface() -> gr.Interface:
     
     # Define language options
     language_choices = [
-        "all",     # All languages
-        "en",      # English
-        "es",      # Spanish
-        "fr",      # French
-        "de",      # German
-        "it",      # Italian
-        "pt",      # Portuguese
-        "ru",      # Russian
-        "zh",      # Chinese
-        "ja",      # Japanese
-        "ko"       # Korean
+        "all", "en", "es", "fr", "de", "it", "pt", "ru", "zh", "ja", "ko"
     ]
     
     # Create mapping for language display names
@@ -526,11 +841,21 @@ def create_gradio_interface() -> gr.Interface:
                 value="all - All Languages",
                 label="Language",
                 info="Select the preferred language for search results"
+            ),
+            gr.Radio(
+                choices=["Auto (Knowledge Base + Web)", "Web Search Only"],
+                value="Auto (Knowledge Base + Web)",
+                label="Search Mode",
+                info="Choose whether to use both knowledge base and web search, or force web search only"
             )
         ],
         additional_inputs_accordion=gr.Accordion("⚙️ Advanced Parameters", open=True),
+        retry_btn="Retry",
+        undo_btn="Undo",
+        clear_btn="Clear",
         chatbot=gr.Chatbot(
             show_copy_button=True,
+            likeable=True,
             layout="bubble",
             height=500,
         )
@@ -558,6 +883,9 @@ def create_parameter_description():
     - **Language**: Preferred language for search results
         - All languages: No language restriction
         - Specific languages: Filter results to selected language
+    - **Search Mode**: Control how queries are processed
+        - Auto: Automatically choose between knowledge base and web search
+        - Web Search Only: Always use web search regardless of query type
     """
 
 if __name__ == "__main__":