Update agent/planner.py

agent/planner.py

@@ -1,9 +1,19 @@
 import logging
+from typing import List, Dict, Set
 
 logger = logging.getLogger(__name__)
 
 class QueryPlanner:
-    def plan_query(self, query: str) -> list:
+    def __init__(self, india_finance_keywords: List[str], category_keywords: Dict[str, Set[str]]):
+        """
+        Initialize QueryPlanner with keyword lists for financial topics and categories.
+        These lists are typically provided by NewsTools or a central configuration.
+        """
+        self.india_finance_keywords = set(india_finance_keywords) # Convert to set for faster lookup
+        self.category_keywords = {k: set(v) for k, v in category_keywords.items()} # Convert values to sets
+        logger.info("QueryPlanner initialized with financial keywords and categories.")
+
+    def plan_query(self, query: str) -> List[Dict]:
         """
         Decompose a user query into a list of planned tasks (e.g., news search, stock price lookup).
         This planning logic is heuristic-based and can be expanded for more sophisticated query understanding.
@@ -15,21 +25,36 @@ class QueryPlanner:
 
             # --- News Search Task Identification ---
             # Keywords indicating a news-related query.
-            news_keywords = [
-                "trend", "news", "articles", "latest", "updates",
-                "market sentiment", "recent", "headlines", "report"
+            news_trigger_keywords = [
+                "trend", "news", "articles", "latest", "updates", "market sentiment",
+                "recent", "headlines", "report", "what's happening", "current situation",
+                "developments"
             ]
             
-            # If any news keyword is present, plan a news search task.
-            # The original query is passed to the news search tool.
-            if any(keyword in query_lower for keyword in news_keywords):
-                # We can refine this to extract specific news categories if the query gets more complex
-                tasks.append({"type": "news_search", "query": query, "category": None}) 
-                logger.debug(f"Identified news search task for query: '{query}'")
+            # Combine general news triggers with specific financial keywords for news search
+            # If any news trigger keyword or general financial keyword is present, plan a news search task.
+            is_news_query = any(keyword in query_lower for keyword in news_trigger_keywords)
+            is_general_financial_query = any(keyword in query_lower for keyword in self.india_finance_keywords)
+
+            # Determine primary category for news search if possible
+            detected_categories = []
+            for category_name, keywords_set in self.category_keywords.items():
+                if any(keyword in query_lower for keyword in keywords_set):
+                    detected_categories.append(category_name)
+            
+            # If it's a news-related query or a general financial query, plan a news search.
+            # Prioritize a specific category if detected, otherwise default to a broader search.
+            if is_news_query or is_general_financial_query:
+                # If multiple categories detected, choose the most relevant or broaden the search
+                # For simplicity, we'll pick the first detected or keep it None for broad search.
+                primary_category = detected_categories[0] if detected_categories else None
+                tasks.append({"type": "news_search", "query": query, "category": primary_category}) 
+                logger.debug(f"Identified news search task for query: '{query}' (Category: {primary_category or 'Any'})")
 
             # --- Stock Price Task Identification ---
-            # A predefined list of common stock symbols or company names to look for.
-            # This list can be expanded or dynamically loaded from a financial data source.
+            # Expanded list of common stock symbols or company names (should come from NewsTools/MCPServer for consistency)
+            # For now, let's keep a sample, but ideally, this would be dynamic or passed in.
+            # NOTE: For better dynamic stock symbol recognition, a more advanced NLP model might be needed.
             known_stock_symbols = [
                 "nifty", "reliance", "tcs", "infy", "hdfcbank", "sbin", 
                 "icicibank", "bhartiartl", "kotakbank", "lt", "asianpaint",
@@ -39,32 +64,27 @@ class QueryPlanner:
             ]
             
             found_symbols = []
-            # Iterate through known symbols to find exact matches or strong indicators in the query.
+            # Check for explicit stock symbol mentions
             for symbol in known_stock_symbols:
-                # Use word boundaries (spaces) to avoid partial matches (e.g., "fin" matching "bajajfinserv")
-                # Also handle cases where symbol is at the start or end of the query, or is the entire query.
-                if (f" {symbol.lower()} " in f" {query_lower} " or 
-                    query_lower.startswith(f"{symbol.lower()} ") or 
-                    query_lower.endswith(f" {symbol.lower()}") or 
-                    query_lower == symbol.lower()):
-                    
-                    # Add the symbol in uppercase to ensure consistency for API calls,
-                    # but only if it hasn't been added already.
-                    if symbol.upper() not in [task.get("symbol") for task in tasks if task["type"] == "stock_price"]:
+                # Using regex for word boundary check (more robust than simple space check)
+                import re
+                if re.search(r'\b' + re.escape(symbol.lower()) + r'\b', query_lower):
+                    if symbol.upper() not in found_symbols: # Ensure uniqueness
                         found_symbols.append(symbol.upper())
                         logger.debug(f"Found explicit stock symbol '{symbol.upper()}' in query.")
 
             # Fallback heuristic: If no known symbol is found but keywords like "stock" or "price" are present,
-            # try to infer a symbol from the last word of the query.
-            if not found_symbols and ("stock" in query_lower or "price" in query_lower or "quote" in query_lower):
+            # try to infer a symbol from the last significant word of the query.
+            if not found_symbols and any(kw in query_lower for kw in ["stock", "price", "quote", "share"]):
                 words = query_lower.split()
-                if words:
-                    potential_symbol = words[-1].upper()
-                    # Basic validation for inferred symbol: 2-5 characters, all alphabetic.
-                    if 2 <= len(potential_symbol) <= 5 and potential_symbol.isalpha():
-                        if potential_symbol not in [task.get("symbol") for task in tasks if task["type"] == "stock_price"]:
+                # Find the last word that is likely a symbol (alphabetic, short)
+                for i in range(len(words) -1, -1, -1):
+                    potential_symbol = words[i].upper()
+                    if 2 <= len(potential_symbol) <= 8 and potential_symbol.isalpha(): # Increased max length
+                        if potential_symbol not in found_symbols:
                             found_symbols.append(potential_symbol)
-                            logger.info(f"Inferred potential stock symbol '{potential_symbol}' from query tail.")
+                            logger.info(f"Inferred potential stock symbol '{potential_symbol}' from query.")
+                            break # Only infer one for now
 
             # Add a stock_price task for each identified symbol.
             for symbol in found_symbols:
@@ -75,3 +95,18 @@ class QueryPlanner:
         except Exception as e:
             logger.error(f"Error planning query '{query}': {e}", exc_info=True)
             return []
+
+    def generate_topic_suggestions(self) -> List[str]:
+        """
+        Generates sample financial topic suggestions for the chatbot.
+        Ported from Flask app.
+        """
+        return [
+            "What's today's Nifty 50 trend?",
+            "Explain RBI's latest repo rate decision",
+            "What's the latest on cryptocurrency in India?",
+            "How is the Indian economy doing?",
+            "What are the recent government policies affecting the market?",
+            "What are the current gold prices in India?",
+            "Tell me about recent IPOs."
+        ]