init branch

.env.sample

@@ -11,5 +11,6 @@ GEMINI_MODEL_PROVIDER=
 PANDAS_KEY=
 PANDAS_EXPORTS_PATH=
 
+OPENAI_MODEL_PROVIDER=
 OPENAI_MODEL=
 OPENAI_API_KEY=

gradio_app.py

@@ -1,3 +1,4 @@
+from asyncio.log import logger
 import yaml
 from pathlib import Path
 import gradio as gr
@@ -6,6 +7,9 @@ from langchain_mcp_client import lc_mcp_exec
 from dotenv import load_dotenv
 import os
 import base64
+from memory_store import MemoryStore
+import logging
+
 
 # ======================================= Load DB configs
 def load_db_configs():
@@ -28,34 +32,44 @@ def image_to_base64_markdown(image_path, alt_text="Customer Status"):
 
 
 # ====================================== Async-compatible wrapper
-async def run_agent(request, history):
-    response, messages = await lc_mcp_exec(request, history)
-
-    image_path = ""
-    load_dotenv()
-    PANDAS_EXPORTS_PATH = os.getenv("PANDAS_EXPORTS_PATH", "exports/charts")
-    
-    # Ensure the exports directory exists
-    os.makedirs(PANDAS_EXPORTS_PATH, exist_ok=True)
-    
-    generated_files = [f for f in os.listdir(PANDAS_EXPORTS_PATH) if f.startswith("temp_chart_") and f.endswith(".png")]
-    if generated_files:
-        image_path = os.path.join(PANDAS_EXPORTS_PATH, generated_files[0])
-        try:
-            image_markdown = image_to_base64_markdown(image_path)
-            output = f"{image_markdown}\n\n{response}"
-            # Remove the image file after reading it
-            os.remove(image_path)
-        except Exception as e:
-            print(f"Error processing image: {e}")
+async def run_agent(request, history=None):
+    try:
+        logger.info(f"Current request: {request}")
+        memory = MemoryStore.get_memory()
+        logger.info(f"Current memory messages: {memory.messages}")
+
+        # Process request using existing memory
+        response, messages = await lc_mcp_exec(request)
+        
+        # Handle image processing
+        image_path = ""
+        load_dotenv()
+        PANDAS_EXPORTS_PATH = os.getenv("PANDAS_EXPORTS_PATH", "exports/charts")
+        
+        # Ensure the exports directory exists
+        os.makedirs(PANDAS_EXPORTS_PATH, exist_ok=True)
+        
+        # Check for generated charts
+        generated_files = [f for f in os.listdir(PANDAS_EXPORTS_PATH) 
+                         if f.startswith("temp_chart_") and f.endswith(".png")]
+        
+        if generated_files:
+            image_path = os.path.join(PANDAS_EXPORTS_PATH, generated_files[0])
+            try:
+                image_markdown = image_to_base64_markdown(image_path)
+                output = f"{image_markdown}\n\n{response}"
+                os.remove(image_path)  # Clean up the image file
+            except Exception as e:
+                logger.error(f"Error processing image: {e}")
+                output = response
+        else:
             output = response
-    else:
-        output = response
 
-    print(f"Image path: {image_path}")
-    print(f"Output length: {len(output)}")
-    
-    return output
+        return output
+
+    except Exception as e:
+        logger.error(f"Error in run_agent: {str(e)}", exc_info=True)
+        return f"Error: {str(e)}"
 
 
 # ====================================== Gradio UI with history
@@ -86,9 +100,8 @@ custom_css = """
 
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     with gr.Row(elem_classes="container"):
-        # with gr.Column(scale=1):
-        #     gr.Image(value=LOGO_PATH, height=200, show_label=False)
-            
+        with gr.Column(scale=1):
+            gr.Image(value=LOGO_PATH, height=200, show_label=False)
         with gr.Column(scale=3):
             gr.Markdown(
                 """
@@ -96,7 +109,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                 <p style='text-align: center'>Ask questions about your database, analyze and visualize data.</p>
                 """
             )
-    
     with gr.Row(elem_classes="container"):
         with gr.Column(scale=3):
             chat = gr.ChatInterface(
@@ -128,7 +140,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                 save_history=True,
                 type="messages"
             )
-        
         with gr.Column(scale=1):
             with gr.Accordion("Example Questions", open=True):
                 gr.Markdown("""

langchain_mcp_client.py

@@ -11,8 +11,12 @@ from langchain.chat_models import init_chat_model
 import logging
 from dotenv import load_dotenv
 from langchain.globals import set_debug
+from langchain.memory import ChatMessageHistory
+from memory_store import MemoryStore
 
-set_debug(True)
+
+
+# set_debug(True)
 
 
 # Set up logging
@@ -23,85 +27,74 @@ load_dotenv()
 
 async def lc_mcp_exec(request: str, history=None) -> Tuple[str, list]:
     """
-    Execute the PostgreSQL MCP pipeline with persistent memory in a JSON file.
+    Execute the PostgreSQL MCP pipeline with in-memory chat history.
     Returns the response and the updated message history.
     """
     try:
-        # Define the path for the chat history JSON file
-        history_file = os.path.join(os.path.dirname(__file__), "chat_history.json")
-
-        # Initialize FileChatMessageHistory for persistent storage
-        message_history = FileChatMessageHistory(file_path=history_file)
+        # Get the singleton memory store instance
+        message_history = MemoryStore.get_memory()
 
         # Load table summary and server parameters
         table_summary = load_table_summary(os.environ["TABLE_SUMMARY_PATH"])
         server_params = get_server_params()
 
         # Initialize the LLM
-        llm = init_chat_model(model_provider=os.getenv("GEMINI_MODEL_PROVIDER"),model=os.getenv("GEMINI_MODEL"),api_key=os.getenv("GEMINI_API_KEY"))
-        #llm =  init_chat_model(model_provider="openai", model=os.getenv("OPENAI_MODEL"), api_key=os.getenv("OPENAI_API_KEY"))
+        # llm = init_chat_model(
+        #     model_provider=os.getenv("OPENAI_MODEL_PROVIDER"), 
+        #     model=os.getenv("OPENAI_MODEL"), 
+        #     api_key=os.getenv("OPENAI_API_KEY")
+        # )
+
+        llm = init_chat_model(
+            model_provider=os.getenv("GEMINI_MODEL_PROVIDER"), 
+            model=os.getenv("GEMINI_MODEL"), 
+            api_key=os.getenv("GEMINI_API_KEY")
+        )
 
         # Initialize the MCP client
         async with stdio_client(server_params) as (read, write):
             async with ClientSession(read, write) as session:
                 await session.initialize()
-
-                # Load tools
                 tools = await load_and_enrich_tools(session)
-
-                # Create the ReAct agent with tools
                 agent = create_react_agent(llm, tools)
 
-                # Add the new user message to history as HumanMessage
+                # Add new user message to memory
                 message_history.add_user_message(request)
 
-                # Prepare the system prompt for the agent
+                # Get system prompt and create system message
                 system_prompt = await build_prompt(session, tools, table_summary)
+                system_message = SystemMessage(content=system_prompt)
 
-                # Init the system message
-                system_message = SystemMessage(
-                    content=system_prompt
-                )
-
-                # Add the system message to the history
+                # Combine system message with chat history
                 input_messages = [system_message] + message_history.messages
 
-                # Invoke the agent with the message list
+                # Invoke agent
                 agent_response = await agent.ainvoke(
                     {"messages": input_messages},
                     config={"configurable": {"thread_id": "conversation_123"}}
                 )
 
-                # Extract the latest response and save all new messages
+                # Process agent response
                 response_content = "No response generated"
                 if "messages" in agent_response and agent_response["messages"]:
-                    # Identify new messages (those not in input_messages)
                     new_messages = agent_response["messages"][len(input_messages):]
-
-                    # Save all new messages to history
+                    
+                    # Save new messages to memory
                     for msg in new_messages:
-                        if isinstance(msg, AIMessage):
-                            # Save AIMessage, including tool_calls if present
-                            message_history.add_message(msg)
-                        elif isinstance(msg, ToolMessage):
-                            # Save ToolMessage with content and tool_call_id
+                        if isinstance(msg, (AIMessage, ToolMessage)):
                             message_history.add_message(msg)
                         else:
-                            # Log unexpected message types for debugging
                             logger.debug(f"Skipping unexpected message type: {type(msg)}")
 
-                    # Use the last message’s content as the response
                     response_content = agent_response["messages"][-1].content
-
                 else:
                     message_history.add_ai_message(response_content)
 
-                # Return the response and the updated history
                 return response_content, message_history.messages
-            
+
     except Exception as e:
         logger.error(f"Error in execution: {str(e)}", exc_info=True)
-        return f"Error: {str(e)}", message_history.messages if 'message_history' in locals() else []
+        return f"Error: {str(e)}", []
 
 # ---------------- Helper Functions ---------------- #
 

memory_store.py

@@ -0,0 +1,28 @@
+from langchain.memory import ChatMessageHistory
+from typing import Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+class MemoryStore:
+    _instance: Optional['MemoryStore'] = None
+    _memory: Optional[ChatMessageHistory] = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(MemoryStore, cls).__new__(cls)
+            cls._memory = ChatMessageHistory()
+            logger.info("New MemoryStore instance created")
+        return cls._instance
+
+    @classmethod
+    def get_memory(cls) -> ChatMessageHistory:
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._memory
+
+    @classmethod
+    def clear_memory(cls):
+        if cls._memory is not None:
+            cls._memory = ChatMessageHistory()
+            logger.info("Memory cleared")