Enhance portfolio extraction in financial agent: Update extraction logic to include total investment amount and improve handling of natural language inputs. Modify JSON response format and refine regex patterns for better accuracy in parsing holdings.

agents/financial_agent.py

@@ -169,20 +169,25 @@ Return only valid JSON, nothing else."""
 
         elif tool_name == "portfolio_analyzer":
             # Use OpenAI to extract portfolio information from natural language
-            extraction_prompt = f"""Extract portfolio holdings from this message: "{message}"
+            extraction_prompt = f"""Extract portfolio holdings and total investment from this message: "{message}"
 
-Convert the portfolio information to JSON format with holdings array.
+Convert the portfolio information to JSON format with holdings array and total investment amount.
 Each holding should have symbol and either shares or percentage.
 
 Return format:
-{{"holdings": [{{"symbol": "AAPL", "shares": 100}}, {{"symbol": "GOOGL", "percentage": 30}}]}}
+{{"holdings": [{{"symbol": "AAPL", "shares": 100}}, {{"symbol": "GOOGL", "percentage": 30}}], "total_investment": 100000}}
 
 Examples:
-- "My portfolio: AAPL 100 shares, GOOGL 50 shares" -> {{"holdings": [{{"symbol": "AAPL", "shares": 100}}, {{"symbol": "GOOGL", "shares": 50}}]}}
-- "I have 40% AAPL, 30% MSFT, 30% TSLA" -> {{"holdings": [{{"symbol": "AAPL", "percentage": 40}}, {{"symbol": "MSFT", "percentage": 30}}, {{"symbol": "TSLA", "percentage": 30}}]}}
-- "Portfolio with Apple 200 shares and Microsoft 25%" -> {{"holdings": [{{"symbol": "AAPL", "shares": 200}}, {{"symbol": "MSFT", "percentage": 25}}]}}
+- "My portfolio: AAPL 100 shares, GOOGL 50 shares" -> {{"holdings": [{{"symbol": "AAPL", "shares": 100}}, {{"symbol": "GOOGL", "shares": 50}}], "total_investment": 0}}
+- "I have 40% AAPL, 30% MSFT, 30% TSLA. I have invested total of 100K USD" -> {{"holdings": [{{"symbol": "AAPL", "percentage": 40}}, {{"symbol": "MSFT", "percentage": 30}}, {{"symbol": "TSLA", "percentage": 30}}], "total_investment": 100000}}
+- "Portfolio with Apple 200 shares and Microsoft 25%, total investment $50,000" -> {{"holdings": [{{"symbol": "AAPL", "shares": 200}}, {{"symbol": "MSFT", "percentage": 25}}], "total_investment": 50000}}
 
-If no clear portfolio data is found, return: {{"holdings": []}}
+Important:
+- Extract total investment amount if mentioned (convert K=1000, M=1000000)
+- If total investment not mentioned, set to 0
+- Convert company names to stock symbols (Apple->AAPL, Microsoft->MSFT, Tesla->TSLA, etc.)
+
+If no clear portfolio data is found, return: {{"holdings": [], "total_investment": 0}}
 
 Return only valid JSON, nothing else."""
 

agents/tools.py

@@ -551,84 +551,84 @@ class FinancialTools:
             try:
                 import re
 
-                # Smart extraction using multiple approaches
+                # Try to parse as JSON first (from OpenAI extraction)
                 total_investment = 0
                 holdings_info = []
-
-                # First, try to extract investment amount using improved patterns
-                def extract_investment_amount(text):
-                    patterns = [
-                        r"(?:invested|investment|total|have)\s*(?:of)?\s*(?:\$)?(\d+(?:[,\d]*)?(?:\.\d+)?)\s*([KMB]?)\s*(?:USD|dollars?|\$)?",
-                        r"(\d+(?:[,\d]*)?(?:\.\d+)?)\s*([KMB]?)\s*(?:USD|dollars?)",
-                        r"\$(\d+(?:[,\d]*)?(?:\.\d+)?)\s*([KMB]?)",
-                    ]
-                    
-                    for pattern in patterns:
-                        match = re.search(pattern, text, re.IGNORECASE)
-                        if match:
-                            amount_str = match.group(1).replace(",", "")
-                            suffix = match.group(2).upper() if len(match.groups()) > 1 else ""
-                            
-                            multiplier = {"K": 1000, "M": 1000000, "B": 1000000000}.get(suffix, 1)
-                            return float(amount_str) * multiplier
-                    return 0
-
-                total_investment = extract_investment_amount(input_str)
-
-                # Extract holdings - percentages vs shares
-                def extract_holdings(text):
-                    holdings = []
-                    
-                    # First try percentage patterns (with % symbol)
-                    percentage_patterns = [
-                        r"([A-Z]{2,5})\s*[:\s]*(\d+(?:\.\d+)?)%",
-                        r"([A-Z]{2,5}):\s*(\d+(?:\.\d+)?)%",
-                        r"([A-Z]{2,5})\s+(\d+(?:\.\d+)?)%",
-                    ]
-                    
-                    for pattern in percentage_patterns:
-                        matches = re.findall(pattern, text, re.IGNORECASE)
-                        if matches:
-                            for symbol, percentage in matches:
-                                holdings.append({
-                                    "symbol": symbol.upper(),
-                                    "percentage": float(percentage)
-                                })
-                            return holdings
-                    
-                    # If no percentages found, try shares patterns (without % symbol)
-                    shares_patterns = [
-                        r"([A-Z]{2,5})\s*[:\s]*(\d+(?:\.\d+)?)\s*(?!%)",
-                        r"([A-Z]{2,5}):\s*(\d+(?:\.\d+)?)\s*(?!%)",
-                        r"([A-Z]{2,5})\s+(\d+(?:\.\d+)?)\s*(?!%)",
-                    ]
-                    
-                    for pattern in shares_patterns:
-                        matches = re.findall(pattern, text, re.IGNORECASE)
-                        if matches:
-                            for symbol, shares in matches:
-                                holdings.append({
-                                    "symbol": symbol.upper(),
-                                    "shares": float(shares)
-                                })
-                            return holdings
-                    
-                    # If no percentage matches, try JSON format
-                    if not holdings:
-                        json_match = re.search(r"\{.*\}|\[.*\]", text, re.DOTALL)
-                        if json_match:
-                            try:
-                                data = json.loads(json_match.group(0))
-                                if isinstance(data, list):
-                                    holdings = data
-                                elif isinstance(data, dict) and "holdings" in data:
-                                    holdings = data["holdings"]
-                            except:
-                                pass
-                    
-                    return holdings
-
-                holdings_info = extract_holdings(input_str)
+                
+                try:
+                    # First try to parse as JSON
+                    data = json.loads(input_str)
+                    if isinstance(data, dict):
+                        holdings_info = data.get("holdings", [])
+                        total_investment = data.get("total_investment", 0)
+                except:
+                    # If JSON parsing fails, extract from natural language
+                    pass
+
+                # If no JSON data found, extract from natural language using regex
+                if not holdings_info:
+                    # Extract investment amount using improved patterns
+                    def extract_investment_amount(text):
+                        patterns = [
+                            r"(?:invested|investment|total|have)\s*(?:of)?\s*(?:\$)?(\d+(?:[,\d]*)?(?:\.\d+)?)\s*([KMB]?)\s*(?:USD|dollars?|\$)?",
+                            r"(\d+(?:[,\d]*)?(?:\.\d+)?)\s*([KMB]?)\s*(?:USD|dollars?)",
+                            r"\$(\d+(?:[,\d]*)?(?:\.\d+)?)\s*([KMB]?)",
+                        ]
+                        
+                        for pattern in patterns:
+                            match = re.search(pattern, text, re.IGNORECASE)
+                            if match:
+                                amount_str = match.group(1).replace(",", "")
+                                suffix = match.group(2).upper() if len(match.groups()) > 1 else ""
+                                
+                                multiplier = {"K": 1000, "M": 1000000, "B": 1000000000}.get(suffix, 1)
+                                return float(amount_str) * multiplier
+                        return 0
+
+                    if total_investment == 0:
+                        total_investment = extract_investment_amount(input_str)
+
+                    # Extract holdings using regex
+                    def extract_holdings(text):
+                        holdings = []
+                        
+                        # First try percentage patterns (with % symbol)
+                        percentage_patterns = [
+                            r"([A-Z]{2,5})\s*[:\s]*(\d+(?:\.\d+)?)%",
+                            r"([A-Z]{2,5}):\s*(\d+(?:\.\d+)?)%",
+                            r"([A-Z]{2,5})\s+(\d+(?:\.\d+)?)%",
+                        ]
+                        
+                        for pattern in percentage_patterns:
+                            matches = re.findall(pattern, text, re.IGNORECASE)
+                            if matches:
+                                for symbol, percentage in matches:
+                                    holdings.append({
+                                        "symbol": symbol.upper(),
+                                        "percentage": float(percentage)
+                                    })
+                                return holdings
+                        
+                        # If no percentages found, try shares patterns (without % symbol)
+                        shares_patterns = [
+                            r"([A-Z]{2,5})\s*[:\s]*(\d+(?:\.\d+)?)\s*(?!%)",
+                            r"([A-Z]{2,5}):\s*(\d+(?:\.\d+)?)\s*(?!%)",
+                            r"([A-Z]{2,5})\s+(\d+(?:\.\d+)?)\s*(?!%)",
+                        ]
+                        
+                        for pattern in shares_patterns:
+                            matches = re.findall(pattern, text, re.IGNORECASE)
+                            if matches:
+                                for symbol, shares in matches:
+                                    holdings.append({
+                                        "symbol": symbol.upper(),
+                                        "shares": float(shares)
+                                    })
+                                return holdings
+                        
+                        return holdings
+
+                    holdings_info = extract_holdings(input_str)
                 
                 # If no valid holdings found, return early to avoid using this tool
                 if not holdings_info: