Update app.py

app.py

@@ -61,7 +61,7 @@ SAFE_SEARCH_OPTIONS = [
 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.
+    Now with improved context handling and response validation.
     """
     logger.info(f'Determining query type for: {query}')
     try:
@@ -78,37 +78,31 @@ async def determine_query_type(query: str, chat_history: List[List[str]], temper
         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"""
+IMPORTANT: You must ONLY respond with either "knowledge_base" or "web_search" - no other text or explanation is allowed.
+
+Classify as "web_search" if the 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
+
+Classify as "knowledge_base" if the 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."}
+            {"role": "user", "content": f"Previous conversation:\n{chat_context}\n\nQuery to classify: {query}\n\nRespond ONLY with 'knowledge_base' or 'web_search':"}
         ]
 
         response = groq_client.chat.completions.create(
@@ -120,13 +114,18 @@ Rules for Classification:
         )
 
         result = response.choices[0].message.content.strip().lower()
-        logger.info(f'Query type determined as: {result} with context')
         
+        # Validate the response
+        if result not in ["knowledge_base", "web_search"]:
+            logger.warning(f'Invalid response from LLM: {result}. Defaulting to knowledge_base')
+            return QueryType.KNOWLEDGE_BASE
+            
+        logger.info(f'Query type determined as: {result}')
         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
+        logger.error(f'Error determining query type: {e}. Defaulting to knowledge_base')
+        return QueryType.KNOWLEDGE_BASE
 
 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."""