feat: add clear cache

app.py

@@ -158,11 +158,13 @@ with gr.Blocks(css=custom_css, theme=xtheme) as demo:
         with gr.Column(scale=1):
             with gr.Accordion("Example Questions", open=True):
                 gr.Markdown("""
+                - 📋 Describe the database
                 - 📊 List all tables in database
                 - 👥 Total number of customers
                 - 📈 Top 10 customers by ticket count
                 - 📋 Ticket count by status and visualize
-                - 📆 Average ticket reopen count per year and month
+                - 📆 Average ticket reopen count per month
+                - 🧹 Clear memory : `/clear-cache`
                 """)
 
 

langchain_mcp_client.py

@@ -36,13 +36,14 @@ async def lc_mcp_exec(request: str, history=None) -> Tuple[str, list]:
         # Load table summary and server parameters
         table_summary = load_table_summary(os.environ["TABLE_SUMMARY_PATH"])
         server_params = get_server_params()
-
-        if os.environ["OPENAI_API_KEY"]:
+        
+        OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+        if OPENAI_API_KEY:
             # Initialize the LLM for OpenAI
             llm = init_chat_model(
                 model_provider=os.environ["OPENAI_MODEL_PROVIDER"],
                 model=os.environ["OPENAI_MODEL"],
-                api_key=os.environ["OPENAI_API_KEY"]
+                api_key=OPENAI_API_KEY
             )
         else:
             # Initialize the LLM for Gemini
@@ -61,6 +62,11 @@ async def lc_mcp_exec(request: str, history=None) -> Tuple[str, list]:
                 tools = await load_and_enrich_tools(session)
                 agent = create_react_agent(llm, tools)
 
+                # clear the memory
+                if request == "/clear-cache":
+                    message_history.clear()
+                    return "Memory cleared", []
+                
                 # Add new user message to memory
                 message_history.add_user_message(request)
 

postgre_mcp_server.py

@@ -117,7 +117,7 @@ When a user submits a request, you must:
 - Always show the SQL query you generate along with the execution result.
 - Validate SQL syntax before execution.
 - Never assume table or column names. Use tools to confirm structure.
-- Use memory efficiently. Don’t rerun a tool unless necessary.
+- Use memory efficiently. Don't rerun a tool unless necessary.
 - If you generate a SQL query, immediately call the **execute_query** tool.
 - If the user requests a visualization, call the **visualize_results** tool with:
   - A visualization prompt starting with "plot" (e.g., "plot a bar chart showing sales by region").
@@ -487,12 +487,12 @@ async def generate_analytical_query(table_name: str) -> list[PromptMessage]:
 
 @mcp.tool(
     description="Identifies both explicit and implied foreign key relationships for a given table using schema analysis and naming patterns.")
-def find_relationships(table_name: str, schema: str = 'public') -> str:
+def find_relationships(table_name: str, db_schema: str = 'public') -> str:
     """Find both explicit and implied relationships for a table.
 
     Args:
         table_name: The name of the table to analyze relationships for
-        schema: The schema name (defaults to 'public')
+        db_schema: The schema name (defaults to 'public')
     """
     try:
         # First get explicit foreign key relationships
@@ -511,7 +511,7 @@ def find_relationships(table_name: str, schema: str = 'public') -> str:
             ON ccu.constraint_name = tc.constraint_name
             AND ccu.table_schema = tc.table_schema
         WHERE tc.constraint_type = 'FOREIGN KEY'
-            AND tc.table_schema = {schema}
+            AND tc.table_schema = {db_schema}
             AND tc.table_name = {table_name}
         """
 
@@ -521,7 +521,7 @@ def find_relationships(table_name: str, schema: str = 'public') -> str:
             -- Get all ID-like columns from our table
             SELECT column_name, data_type
             FROM information_schema.columns
-            WHERE table_schema = {schema}
+            WHERE table_schema = {db_schema}
             AND table_name = {table_name}
             AND (
                 column_name LIKE '%%id' 
@@ -557,7 +557,7 @@ def find_relationships(table_name: str, schema: str = 'public') -> str:
                 ON c.table_schema = t.table_schema 
                 AND c.table_name = t.table_name
                 AND (c.column_name = 'id' OR c.column_name = sc.column_name)
-            WHERE t.table_schema = {schema}
+            WHERE t.table_schema = {db_schema}
                 AND t.table_name != {table_name}  -- Exclude self-references
         )
         SELECT 
@@ -619,10 +619,10 @@ async def visualize_results(json_data: dict, vis_prompt: str) -> str:
     """
     try:
         # Debug prints to see what's being received
-        
-        OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
-        pllm = OpenAI(api_token=OPENAI_API_KEY)
-        pai.config.set({"llm": pllm})
+        OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+        if OPENAI_API_KEY:
+            pllm = OpenAI(api_token=OPENAI_API_KEY)
+            pai.config.set({"llm": pllm})
 
         # Convert JSON to DataFrame
         df = pd.DataFrame(json_data["data"], columns=json_data["columns"])