added gemini support and fast-mcp integration

.gitignore

@@ -1,2 +1,3 @@
 demo_data/*
 .env*.png
+.env

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 import asyncio
 import os
+import json
+import plotly.io as pio
 from money_rag import MoneyRAG
 
 st.set_page_config(page_title="MoneyRAG", layout="wide")
@@ -12,7 +14,7 @@ with st.sidebar:
     
     if provider == "Google":
         models = ["gemini-3-flash-preview", "gemini-3-pro-image-preview", "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite"]
-        embeddings = ["text-embedding-004"]
+        embeddings = ["gemini-embedding-001"]
     else:
         models = ["gpt-5-mini", "gpt-5-nano", "gpt-4o-mini", "gpt-4o"]
         embeddings = ["text-embedding-3-small", "text-embedding-3-large", "text-embedding-ada-002"]
@@ -85,10 +87,34 @@ if "rag" in st.session_state:
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
+    # Helper function to cleverly render either text or a Plotly chart
+    def render_content(content):
+        # We might have mixed text and charts delimited by ===CHART=== ... ===ENDCHART===
+        if isinstance(content, str) and "===CHART===" in content:
+            parts = content.split("===CHART===")
+            # Render first text part
+            st.markdown(parts[0].strip())
+            
+            for part in parts[1:]:
+                if "===ENDCHART===" in part:
+                    chart_json, remaining_text = part.split("===ENDCHART===")
+                    try:
+                        fig = pio.from_json(chart_json.strip())
+                        st.plotly_chart(fig, use_container_width=True)
+                    except Exception as e:
+                        st.error("Failed to render chart.")
+                    
+                    if remaining_text.strip():
+                        st.markdown(remaining_text.strip())
+        else:
+            st.markdown(content)
+
+    # Render previous messages
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
-            st.markdown(message["content"])
+            render_content(message["content"])
 
+    # Handle new user input
     if prompt := st.chat_input("Ask about your spending..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
         with st.chat_message("user"):
@@ -97,7 +123,9 @@ if "rag" in st.session_state:
         with st.chat_message("assistant"):
             with st.spinner("Thinking..."):
                 response = asyncio.run(st.session_state.rag.chat(prompt))
-                st.markdown(response)
+                render_content(response)
+                
         st.session_state.messages.append({"role": "assistant", "content": response})
+
 else:
     st.info("Please authenticate in the sidebar to start.")

mcp_server.py

@@ -1,3 +1,5 @@
+import pandas as pd
+import plotly.express as px
 from fastmcp import FastMCP
 from langchain_qdrant import QdrantVectorStore
 from qdrant_client import QdrantClient
@@ -180,6 +182,50 @@ def semantic_search(query: str, top_k: int = 5) -> str:
     except Exception as e:
         return f"Error performing search: {str(e)}"
 
+
+@mcp.tool()
+def generate_interactive_chart(sql_query: str, chart_type: str, x_col: str, y_col: str, title: str) -> str:
+    """
+    Generate an interactive Plotly chart from the money_rag SQLite database.
+    Use this proactively whenever a visual representation of data would be helpful.
+    
+    CRITICAL INSTRUCTIONS:
+    1. Write a valid SQLite SELECT query. 
+    2. Aggregate data appropriately (e.g., use GROUP BY for pie/bar charts).
+    3. Pass the exact column names from your query to x_col and y_col.
+    
+    Args:
+        sql_query: The SQL SELECT query (e.g. "SELECT category, SUM(amount) as total FROM transactions GROUP BY category")
+        chart_type: Must be exactly "bar", "pie", or "line"
+        x_col: Column name from query for X-axis (or labels for pie)
+        y_col: Column name from query for Y-axis (or values for pie)
+        title: Title of the chart
+    """
+    try:
+        conn = sqlite3.connect(DB_PATH)
+        df = pd.read_sql_query(sql_query, conn)
+        conn.close()
+        if df.empty:
+            return '{"error": "No data found for this query."}'
+        if chart_type == "bar":
+            fig = px.bar(df, x=x_col, y=y_col, title=title)
+        elif chart_type == "pie":
+            fig = px.pie(df, names=x_col, values=y_col, title=title)
+        elif chart_type == "line":
+            fig = px.line(df, x=x_col, y=y_col, title=title)
+        else:
+            return f'{{"error": "Unsupported chart type: {chart_type}"}}'
+        # Write the huge JSON to a temp file instead of returning it directly to LLM context
+        chart_path = os.path.join(DATA_DIR, "latest_chart.json")
+        with open(chart_path, "w") as f:
+            f.write(fig.to_json())
+            
+        return "Chart generated successfully! It has been sent to the user's UI. Continue analyzing without outputting the JSON parameters directly."
+        
+    except Exception as e:
+        return f'{{"error": "Failed to generate chart: {str(e)}"}}'
+
+
 # A helper to clear data (useful for session reset)
 @mcp.tool()
 def clear_database() -> str:

money_rag.py

@@ -209,6 +209,9 @@ class MoneyRAG:
             "You are a financial analyst. Use the provided tools to query the database "
             "and perform semantic searches. Spending is POSITIVE (>0). "
             "Always explain your findings clearly."
+            "IMPORTANT: Whenever possible and relevant (e.g. when discussing trends, comparing categories, or showing breakdowns), "
+            "you MUST proactively use the 'generate_interactive_chart' tool to generate visual plots (bar, pie, or line charts) to accompany your analysis. "
+            "WARNING: You MUST use the actual tool call to generate the chart. DO NOT simply output a json block with chart parameters as your final text answer."
         )
         
         self.agent = create_agent(
@@ -221,6 +224,11 @@ class MoneyRAG:
     async def chat(self, query: str):
         config = {"configurable": {"thread_id": "session_1"}}
         
+        # Clear out any previous chart so we don't carry over stale plots
+        chart_path = os.path.join(self.temp_dir, "latest_chart.json")
+        if os.path.exists(chart_path):
+            os.remove(chart_path)
+        
         result = await self.agent.ainvoke(
             {"messages": [{"role": "user", "content": query}]},
             config,
@@ -235,10 +243,18 @@ class MoneyRAG:
             for block in content:
                 if isinstance(block, dict) and block.get("type") == "text":
                     text_parts.append(block.get("text", ""))
-            return "\n".join(text_parts)
-        
-        # If content is already a string (OpenAI format), return as-is
-        return content
+            final_text = "\n".join(text_parts)
+        else:
+            final_text = content
+            
+        # Check if the tool generated a chart file on disk during this turn
+        chart_path = os.path.join(self.temp_dir, "latest_chart.json")
+        if os.path.exists(chart_path):
+            with open(chart_path, "r") as f:
+                chart_json = f.read()
+            final_text += f"\n\n===CHART===\n{chart_json}\n===ENDCHART==="
+            
+        return final_text
 
     async def cleanup(self):
         """Delete temporary session files and close MCP client."""