Merge pull request #5 from pia-team/amirkia_test

A temporary solution for conversation memory in Langchain

.env.sample

@@ -0,0 +1,2 @@
+API_KEY=
+MCP_SERVER_PATH=/home/andre/andre/postgre_mcp_server.py

.gitignore

@@ -0,0 +1,125 @@
+.idea
+.env
+.vscode
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+nosetests.xml
+coverage.xml
+*.coveragerc
+*.pylintrc
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff
+instance/
+.webassets-cache
+
+# Scrapy stuff
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath
+*.sage.py
+
+# Environments
+.env/
+.venv/
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/

chat_history.json

@@ -0,0 +1 @@
+[{"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my previous request ?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "# Result\nlist all tables\n# Explanation\nThe user's previous request was to list all the tables in the database. This was determined by reviewing the previous turn in the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "count tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\n4\n# Explanation\nThe user requested the number of tables. The previous response listed the tables, and counting them gives the answer.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my last message ?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\ncount tables\n# Explanation\nThe user's last message was \"count tables\". This was determined by reviewing the previous turn in the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my very first message in this conversation?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nlist all tables\n# Explanation\nThe user's very first message in this conversation was \"list all tables\". This was determined by reviewing the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "total number of sales in 2024", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nThe tables in the database are:\n*   dim_agreement\n*   dim_customer\n*   dim_product\n*   dim_product_order_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}]

gradio_app.py

@@ -0,0 +1,49 @@
+import yaml
+from pathlib import Path
+import gradio as gr
+import asyncio
+from postgre_mcp_client import pg_mcp_exec
+from postgre_smolagent_clinet import pg_mcp_smolagent_exec
+from langchain_mcp_client import lc_mcp_exec
+
+def load_db_configs():
+    """Load database configurations from configs.yaml"""
+    configs_path = Path("configs.yaml")
+
+    if not configs_path.exists():
+        raise FileNotFoundError("configs.yaml not found")
+
+    with open(configs_path) as f:
+        configs = yaml.safe_load(f)
+
+    return configs["db_configs"]
+
+# Async-compatible wrapper
+async def run_agent(request, history):
+    # configs = load_db_configs()
+    # final_answer, last_tool_answer, = await pg_mcp_exec(request)
+    # return final_answer, last_tool_answer
+
+    response, message = await lc_mcp_exec(request, history)
+    return response
+
+# Gradio UI
+demo = gr.Interface(
+    fn=run_agent,
+    inputs=gr.Textbox(
+        label="Natural Language Request",
+        placeholder="e.g., Show me the table of join posts and users tables."
+    ),
+    outputs=gr.Markdown(
+        label="SQL Query / Result",
+        value="# Result \n \n",
+        # max_height=1500,
+        height= 1500
+    ),
+    title="PostgreSQL Query Agent",
+    description="Ask your database in natural language and get results.",
+    flagging_mode="never"
+)
+
+if __name__ == "__main__":
+    demo.launch()

langchain_mcp_client.py

@@ -0,0 +1,167 @@
+import os.path
+import json
+from typing import Tuple, Any
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from langchain_mcp_adapters.tools import load_mcp_tools
+from langgraph.prebuilt import create_react_agent
+from langchain_core.prompts import PromptTemplate
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain.memory import ChatMessageHistory
+from langchain_community.chat_message_histories import FileChatMessageHistory
+from langchain.chat_models import init_chat_model
+from utils import parse_mcp_output, classify_intent
+import logging
+from dotenv import load_dotenv
+
+
+
+logger = logging.getLogger(__name__)
+load_dotenv()
+async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
+    """
+    Execute the full PostgreSQL MCP pipeline with persistent memory.
+    """
+    try:
+        history_file = os.path.join(os.path.dirname(__file__), "chat_history.json")
+        
+        # Initialize chat history file if it doesn't exist or is empty
+        if not os.path.exists(history_file) or os.path.getsize(history_file) == 0:
+            with open(history_file, 'w') as f:
+                json.dump({"messages": []}, f)
+
+        message_history = FileChatMessageHistory(file_path=history_file)
+        
+        try:
+            # Load existing messages or handle bootstrap scenario
+            existing_messages = message_history.messages
+        except json.JSONDecodeError:
+            # If JSON is corrupted, reinitialize the file
+            logger.warning("Chat history file corrupted, reinitializing...")
+            with open(history_file, 'w') as f:
+                json.dump({"messages": []}, f)
+            existing_messages = []
+
+        # Format existing messages properly
+        formatted_history = []
+        for msg in existing_messages:
+            if isinstance(msg, HumanMessage):
+                formatted_history.append(HumanMessage(content=msg.content))
+            elif isinstance(msg, AIMessage):
+                formatted_history.append(AIMessage(content=msg.content))
+
+        # TODO: give summary file path from config
+        table_summary = load_table_summary("table_summary.txt")
+        server_params = get_server_params()
+
+        api_key = os.getenv("API_KEY")
+        llm = init_chat_model(model="gemini-2.0-flash", model_provider="google_genai",
+                              api_key=api_key)
+
+        async with stdio_client(server_params) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+
+                tools = await load_and_enrich_tools(session)
+                intent = classify_intent(request)
+
+                # Add new user message before processing
+                message_history.add_user_message(request)
+
+                # Create agent and prepare system message
+                agent = create_react_agent(llm, tools)
+                
+                # Create base messages list with system message
+                base_message = HumanMessage(content="""You are a PostgreSQL database expert assistant. 
+                    Use the conversation history for context when available.""")
+                messages = [base_message]
+                
+                # Add history if exists
+                if formatted_history:
+                    messages.extend(formatted_history)
+                
+                # Add current request
+                messages.append(HumanMessage(content=request))
+
+                # Build prompt with conversation context
+                prompt = await build_prompt(session, intent, request, tools, table_summary, formatted_history)
+
+                # Invoke agent with proper message structure
+                agent_response = await agent.ainvoke(
+                    {
+                        "messages": prompt,
+                        "chat_history": [msg.content for msg in formatted_history]
+                    },
+                    config={"configurable": {"thread_id": "conversation_123"}}
+                )
+
+                if "messages" in agent_response:
+                    response = agent_response["messages"][-1].content
+                    # Save assistant response
+                    message_history.add_ai_message(response)
+                else:
+                    response = "No response generated"
+                    message_history.add_ai_message(response)
+
+                # Return current response and up-to-date messages
+                return response, message_history.messages
+
+    except Exception as e:
+        logger.error(f"Error in chat history handling: {str(e)}", exc_info=True)
+        # Fallback to stateless response if history fails
+        return f"Error in conversation: {str(e)}", []
+
+
+# ---------------- Helper Functions ---------------- #
+
+def load_table_summary(path: str) -> str:
+    with open(path, 'r') as file:
+        return file.read()
+
+def get_server_params() -> StdioServerParameters:
+    # TODO: give server params from config
+    load_dotenv()
+    return StdioServerParameters(
+        command="python",
+        args=[os.environ["MCP_SERVER_PATH"]],
+    )
+
+async def load_and_enrich_tools(session: ClientSession):
+    tools = await load_mcp_tools(session)
+    return tools
+
+async def build_prompt(session, intent, request, tools, summary, chat_history):
+    superset_prompt = await session.read_resource("resource://last_prompt")
+    conversation_prompt = await session.read_resource("resource://base_prompt")
+    
+    if intent == "superset_request":
+        template = superset_prompt.contents[0].text
+        return template.format(
+            new_request=request
+        )
+    else:
+        template = conversation_prompt.contents[0].text
+        tools_str = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
+        
+        # Handle history formatting with proper message access
+        history_str = ""
+        if chat_history:
+            history_sections = []
+            for msg in chat_history:
+                if isinstance(msg, HumanMessage):
+                    history_sections.append(f"Previous Human Question:\n{msg.content}\n")
+                elif isinstance(msg, AIMessage):
+                    history_sections.append(f"Previous Assistant Response:\n{msg.content}\n")
+            history_str = "\n".join(history_sections)
+
+        return template.format(
+            new_request=request,
+            tools=tools_str,
+            descriptions=summary,
+            chat_history=f"\nPrevious Conversation History:\n{history_str}" if history_str else "\nThis is a new conversation.",
+            system_instructions="""You are a PostgreSQL database expert assistant.
+                Use the conversation history when available to maintain context.
+                For new conversations, focus on understanding the initial request."""
+        )
+

main.py

@@ -5,6 +5,9 @@ import asyncio
 from postgre_mcp_client import pg_mcp_exec
 import logging
 
+from postgre_smolagent_clinet import pg_mcp_smolagent_exec
+
+
 #logger = logging.getLogger(__name__)
 # TODO add config
 def load_db_configs():
@@ -23,8 +26,11 @@ def load_db_configs():
 async def main():
     #configs = load_db_configs()
 
-    request = "Show me the table of join posts and users tables."
-    await pg_mcp_exec(request)
+    request = "send to superset"
+    # request = "Show me the table of join posts and users tables."
+
+    # await pg_mcp_exec(request)
+    # await pg_mcp_smolagent_exec(request)
 
 
 if __name__ == "__main__":

memory.json

@@ -1 +0,0 @@
-{"history": [[{"type": "user_message", "content": "=========================\n# Your Role\n=========================\n\nYou are an expert in generating SQL queries and interacting with a PostgreSQL database using FastMCP tools. These tools allow you to:\n\n- List available tables\n- Retrieve schema details\n- Execute SQL queries\n\nEach tool may also provide summaries of table contents to help you understand the data structure. You have access to **short-term memory**, which stores relevant information from earlier steps or previous queries. If the memory is not empty, you **must** use it when processing the current request. Avoid repeating the same tool with identical input unless the result is **not already present in memory**.\n\n=========================\n# Your Objective\n=========================\n\nWhen a user submits a request, you must:\n\n1. **Analyze the request** to determine the required data or action.\n2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).\n3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.\n4. **Execute the SQL query** and return the results.\n5. **Chain tools logically**, such as: List Tables \u2192 Get Schema \u2192 Write and Run Query.\n6. **Explain your reasoning and each step taken** to ensure clarity and transparency.\n\n=========================\n# Critical Rules\n=========================\n\n- Only use **SELECT** queries. Never use destructive operations (e.g., DELETE, DROP, UPDATE).\n- Always display any SQL query you generate along with the result of its execution.\n- Validate SQL syntax before execution to ensure correctness and safety.\n- Never make assumptions about the database structure \u2014 always use tools to confirm table and column names.\n- Be cautious, deliberate, and transparent in your actions.\n\n=========================\n# Short-Term Memory\n=========================\n\nYou have access to the following memory from this conversation. Use it if applicable for the current request.\n\n### Conversation Context:\n- **Previous user requests**:\n  []\n- **Tools used so far**:\n  []\n- **Last SQL queries (if any)**:\n  []\n- **Last result preview**:\n  []\n\n=========================\n# New User Request\n=========================\n\nPlease fulfill the following request based on the above context:\n\ncan you show me the result of the join of all tables?\n"}, {"type": "ai_function_call", "ai_said": "Okay, I will attempt to join all tables and show the results. Here's the plan:\n\n1.  **List tables:** I'll start by listing the tables in the database to know which tables are available.\n2.  **Get schema:** I'll retrieve the schema of each table to understand their structure and identify potential join columns.\n3.  **Construct and execute the SQL query:** Based on the table schemas, I will construct a SQL query to join the tables. Since there is no information about the relationships between tables, I will try to join them based on common columns like 'id' and table names. I will then execute the query.\n\nLet's start by listing the tables.", "tool": "list_tables", "args": "{}"}, {"type": "tool_response", "tool": "list_tables", "response": "posts\nusers"}, {"type": "ai_function_call", "ai_said": "Okay, I see two tables: \"posts\" and \"users\". I will now retrieve the schema of both tables to understand their structure and identify potential join columns.", "tool": "get_table_schema", "args": "{\"table_name\": \"users\"}"}, {"type": "tool_response", "tool": "get_table_schema", "response": "Table: posts\nColumns:\n- id (integer) NOT NULL DEFAULT nextval('posts_id_seq'::regclass)\n- user_id (integer) NULL\n- title (character varying(200)) NULL\n- content (text) NULL\n- created_at (timestamp without time zone) NULL DEFAULT CURRENT_TIMESTAMP"}, {"type": "tool_response", "tool": "get_table_schema", "response": "Table: users\nColumns:\n- id (integer) NOT NULL DEFAULT nextval('users_id_seq'::regclass)\n- username (character varying(100)) NOT NULL\n- email (character varying(100)) NOT NULL\n- created_at (timestamp without time zone) NULL DEFAULT CURRENT_TIMESTAMP"}, {"type": "ai_function_call", "ai_said": ["Okay, I need to join the `posts` and `users` tables. Based on the schema, the `posts` table has a `user_id` column that likely references the `id` column in the `users` table. I will construct a query to join these tables on the `user_id` and `id` columns.\n\nHere's the SQL query:", "```sql\nSELECT *\nFROM posts\nJOIN users ON posts.user_id = users.id;\n```", "\n\nNow, I will execute the query."], "tool": "execute_query", "args": "{\"query\": \"SELECT * FROM posts JOIN users ON posts.user_id = users.id;\"}"}, {"type": "tool_response", "tool": "execute_query", "response": "id | user_id | title | content | created_at | id | username | email | created_at\n--------------------------------------------------------------------------------\n1 | 1 | Hello World | This is my first post! | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107\n2 | 2 | Bob\u2019s Thoughts | Here are some thoughts... | 2025-04-09 11:16:04.846107 | 2 | bob | bob@example.com | 2025-04-09 11:16:04.846107\n3 | 1 | Follow-up Post | Another one from Alice. | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107"}, {"type": "ai_final_answer", "ai_said": "I have successfully joined the `posts` and `users` tables on the `user_id` and `id` columns, respectively, and displayed the results."}], [{"type": "user_message", "content": "=========================\n# Your Role\n=========================\n\nYou are an expert in generating SQL queries and interacting with a PostgreSQL database using FastMCP tools. These tools allow you to:\n\n- List available tables\n- Retrieve schema details\n- Execute SQL queries\n\nEach tool may also provide summaries of table contents to help you understand the data structure. You have access to **short-term memory**, which stores relevant information from earlier steps or previous queries. If the memory is not empty, you **must** use it when processing the current request. Avoid repeating the same tool with identical input unless the result is **not already present in memory**.\n\n=========================\n# Your Objective\n=========================\n\nWhen a user submits a request, you must:\n\n1. **Analyze the request** to determine the required data or action.\n2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).\n3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.\n4. **Execute the SQL query** and return the results.\n5. **Chain tools logically**, such as: List Tables \u2192 Get Schema \u2192 Write and Run Query.\n6. **Explain your reasoning and each step taken** to ensure clarity and transparency.\n\n=========================\n# Critical Rules\n=========================\n\n- Only use **SELECT** queries. Never use destructive operations (e.g., DELETE, DROP, UPDATE).\n- Always display any SQL query you generate along with the result of its execution.\n- Validate SQL syntax before execution to ensure correctness and safety.\n- Never make assumptions about the database structure \u2014 always use tools to confirm table and column names.\n- Be cautious, deliberate, and transparent in your actions.\n\n=========================\n# Short-Term Memory\n=========================\n\nYou have access to the following memory from this conversation. Use it if applicable for the current request.\n\n### Conversation Context:\n- **Previous user requests**:\n  ['=========================\\n# Your Role\\n=========================\\n\\nYou are an expert in generating SQL queries and interacting with a PostgreSQL database using FastMCP tools. These tools allow you to:\\n\\n- List available tables\\n- Retrieve schema details\\n- Execute SQL queries\\n\\nEach tool may also provide summaries of table contents to help you understand the data structure. You have access to **short-term memory**, which stores relevant information from earlier steps or previous queries. If the memory is not empty, you **must** use it when processing the current request. Avoid repeating the same tool with identical input unless the result is **not already present in memory**.\\n\\n=========================\\n# Your Objective\\n=========================\\n\\nWhen a user submits a request, you must:\\n\\n1. **Analyze the request** to determine the required data or action.\\n2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).\\n3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.\\n4. **Execute the SQL query** and return the results.\\n5. **Chain tools logically**, such as: List Tables \u2192 Get Schema \u2192 Write and Run Query.\\n6. **Explain your reasoning and each step taken** to ensure clarity and transparency.\\n\\n=========================\\n# Critical Rules\\n=========================\\n\\n- Only use **SELECT** queries. Never use destructive operations (e.g., DELETE, DROP, UPDATE).\\n- Always display any SQL query you generate along with the result of its execution.\\n- Validate SQL syntax before execution to ensure correctness and safety.\\n- Never make assumptions about the database structure \u2014 always use tools to confirm table and column names.\\n- Be cautious, deliberate, and transparent in your actions.\\n\\n=========================\\n# Short-Term Memory\\n=========================\\n\\nYou have access to the following memory from this conversation. Use it if applicable for the current request.\\n\\n### Conversation Context:\\n- **Previous user requests**:\\n  []\\n- **Tools used so far**:\\n  []\\n- **Last SQL queries (if any)**:\\n  []\\n- **Last result preview**:\\n  []\\n\\n=========================\\n# New User Request\\n=========================\\n\\nPlease fulfill the following request based on the above context:\\n\\ncan you show me the result of the join of all tables?\\n']\n- **Tools used so far**:\n  ['execute_query', 'get_table_schema', 'list_tables']\n- **Last SQL queries (if any)**:\n  []\n- **Last result preview**:\n  ['id | user_id | title | content | created_at | id | username | email | created_at\\n--------------------------------------------------------------------------------\\n1 | 1 | Hello World | This is my first post! | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107\\n2 | 2 | Bob\u2019s Thoughts | Here are some thoughts... | 2025-04-09 11:16:04.846107 | 2 | bob | bob@example.com | 2025-04-09 11:16:04.846107\\n3 | 1 | Follow-up Post | Another one from Alice. | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107']\n\n=========================\n# New User Request\n=========================\n\nPlease fulfill the following request based on the above context:\n\nhow many columns are there in this joined table?\n"}, {"type": "ai_function_call", "ai_said": ["Okay, I will determine the number of columns in the joined table. Since I have already performed a join operation and have the result in memory, I will use the `execute_query` tool to find the number of columns.\n\nHere's the SQL query I will use:", "```sql\nSELECT COUNT(*) FROM (SELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at FROM posts p JOIN users u ON p.user_id = u.id) AS joined_table\n```", "\n\nNow, I will execute this query using the `execute_query` tool."], "tool": "execute_query", "args": "{\"query\": \"SELECT COUNT(*) FROM (SELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at FROM posts p JOIN users u ON p.user_id = u.id) AS joined_table\"}"}, {"type": "tool_response", "tool": "execute_query", "response": "count\n-----\n3"}, {"type": "ai_final_answer", "ai_said": "The previous query joined the `posts` and `users` tables. The query to determine the number of columns in the joined table is: `SELECT COUNT(*) FROM (SELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at FROM posts p JOIN users u ON p.user_id = u.id) AS joined_table`. The result shows that there are 9 columns."}]], "tools_used": ["list_tables", "get_table_schema", "execute_query", "execute_query"], "all_queries": [], "query_results": ["id | user_id | title | content | created_at | id | username | email | created_at\n--------------------------------------------------------------------------------\n1 | 1 | Hello World | This is my first post! | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107\n2 | 2 | Bob\u2019s Thoughts | Here are some thoughts... | 2025-04-09 11:16:04.846107 | 2 | bob | bob@example.com | 2025-04-09 11:16:04.846107\n3 | 1 | Follow-up Post | Another one from Alice. | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107", "count\n-----\n3"], "user_messages": ["=========================\n# Your Role\n=========================\n\nYou are an expert in generating SQL queries and interacting with a PostgreSQL database using FastMCP tools. These tools allow you to:\n\n- List available tables\n- Retrieve schema details\n- Execute SQL queries\n\nEach tool may also provide summaries of table contents to help you understand the data structure. You have access to **short-term memory**, which stores relevant information from earlier steps or previous queries. If the memory is not empty, you **must** use it when processing the current request. Avoid repeating the same tool with identical input unless the result is **not already present in memory**.\n\n=========================\n# Your Objective\n=========================\n\nWhen a user submits a request, you must:\n\n1. **Analyze the request** to determine the required data or action.\n2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).\n3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.\n4. **Execute the SQL query** and return the results.\n5. **Chain tools logically**, such as: List Tables \u2192 Get Schema \u2192 Write and Run Query.\n6. **Explain your reasoning and each step taken** to ensure clarity and transparency.\n\n=========================\n# Critical Rules\n=========================\n\n- Only use **SELECT** queries. Never use destructive operations (e.g., DELETE, DROP, UPDATE).\n- Always display any SQL query you generate along with the result of its execution.\n- Validate SQL syntax before execution to ensure correctness and safety.\n- Never make assumptions about the database structure \u2014 always use tools to confirm table and column names.\n- Be cautious, deliberate, and transparent in your actions.\n\n=========================\n# Short-Term Memory\n=========================\n\nYou have access to the following memory from this conversation. Use it if applicable for the current request.\n\n### Conversation Context:\n- **Previous user requests**:\n  []\n- **Tools used so far**:\n  []\n- **Last SQL queries (if any)**:\n  []\n- **Last result preview**:\n  []\n\n=========================\n# New User Request\n=========================\n\nPlease fulfill the following request based on the above context:\n\ncan you show me the result of the join of all tables?\n", "=========================\n# Your Role\n=========================\n\nYou are an expert in generating SQL queries and interacting with a PostgreSQL database using FastMCP tools. These tools allow you to:\n\n- List available tables\n- Retrieve schema details\n- Execute SQL queries\n\nEach tool may also provide summaries of table contents to help you understand the data structure. You have access to **short-term memory**, which stores relevant information from earlier steps or previous queries. If the memory is not empty, you **must** use it when processing the current request. Avoid repeating the same tool with identical input unless the result is **not already present in memory**.\n\n=========================\n# Your Objective\n=========================\n\nWhen a user submits a request, you must:\n\n1. **Analyze the request** to determine the required data or action.\n2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).\n3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.\n4. **Execute the SQL query** and return the results.\n5. **Chain tools logically**, such as: List Tables \u2192 Get Schema \u2192 Write and Run Query.\n6. **Explain your reasoning and each step taken** to ensure clarity and transparency.\n\n=========================\n# Critical Rules\n=========================\n\n- Only use **SELECT** queries. Never use destructive operations (e.g., DELETE, DROP, UPDATE).\n- Always display any SQL query you generate along with the result of its execution.\n- Validate SQL syntax before execution to ensure correctness and safety.\n- Never make assumptions about the database structure \u2014 always use tools to confirm table and column names.\n- Be cautious, deliberate, and transparent in your actions.\n\n=========================\n# Short-Term Memory\n=========================\n\nYou have access to the following memory from this conversation. Use it if applicable for the current request.\n\n### Conversation Context:\n- **Previous user requests**:\n  ['=========================\\n# Your Role\\n=========================\\n\\nYou are an expert in generating SQL queries and interacting with a PostgreSQL database using FastMCP tools. These tools allow you to:\\n\\n- List available tables\\n- Retrieve schema details\\n- Execute SQL queries\\n\\nEach tool may also provide summaries of table contents to help you understand the data structure. You have access to **short-term memory**, which stores relevant information from earlier steps or previous queries. If the memory is not empty, you **must** use it when processing the current request. Avoid repeating the same tool with identical input unless the result is **not already present in memory**.\\n\\n=========================\\n# Your Objective\\n=========================\\n\\nWhen a user submits a request, you must:\\n\\n1. **Analyze the request** to determine the required data or action.\\n2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).\\n3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.\\n4. **Execute the SQL query** and return the results.\\n5. **Chain tools logically**, such as: List Tables \u2192 Get Schema \u2192 Write and Run Query.\\n6. **Explain your reasoning and each step taken** to ensure clarity and transparency.\\n\\n=========================\\n# Critical Rules\\n=========================\\n\\n- Only use **SELECT** queries. Never use destructive operations (e.g., DELETE, DROP, UPDATE).\\n- Always display any SQL query you generate along with the result of its execution.\\n- Validate SQL syntax before execution to ensure correctness and safety.\\n- Never make assumptions about the database structure \u2014 always use tools to confirm table and column names.\\n- Be cautious, deliberate, and transparent in your actions.\\n\\n=========================\\n# Short-Term Memory\\n=========================\\n\\nYou have access to the following memory from this conversation. Use it if applicable for the current request.\\n\\n### Conversation Context:\\n- **Previous user requests**:\\n  []\\n- **Tools used so far**:\\n  []\\n- **Last SQL queries (if any)**:\\n  []\\n- **Last result preview**:\\n  []\\n\\n=========================\\n# New User Request\\n=========================\\n\\nPlease fulfill the following request based on the above context:\\n\\ncan you show me the result of the join of all tables?\\n']\n- **Tools used so far**:\n  ['execute_query', 'get_table_schema', 'list_tables']\n- **Last SQL queries (if any)**:\n  []\n- **Last result preview**:\n  ['id | user_id | title | content | created_at | id | username | email | created_at\\n--------------------------------------------------------------------------------\\n1 | 1 | Hello World | This is my first post! | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107\\n2 | 2 | Bob\u2019s Thoughts | Here are some thoughts... | 2025-04-09 11:16:04.846107 | 2 | bob | bob@example.com | 2025-04-09 11:16:04.846107\\n3 | 1 | Follow-up Post | Another one from Alice. | 2025-04-09 11:16:04.846107 | 1 | alice | alice@example.com | 2025-04-09 11:16:04.846107']\n\n=========================\n# New User Request\n=========================\n\nPlease fulfill the following request based on the above context:\n\nhow many columns are there in this joined table?\n"]}

postgre_mcp_client.py

@@ -1,4 +1,6 @@
 import os.path
+from typing import Tuple, Any
+
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 from langchain_mcp_adapters.tools import load_mcp_tools
@@ -10,7 +12,7 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-async def pg_mcp_exec(request: str) -> str:
+async def pg_mcp_exec(request: str) -> tuple[Any, Any]:
     """
     Execute the full PostgreSQL MCP pipeline: load summary, connect session,
     load memory and tools, build prompt, run agent, update memory.
@@ -27,7 +29,7 @@ async def pg_mcp_exec(request: str) -> str:
     server_params = get_server_params()
 
     # TODO: give key from env
-    llm = init_chat_model(model="gemini-2.0-flash-lite", model_provider="google_genai",
+    llm = init_chat_model(model="gemini-2.0-flash", model_provider="google_genai",
                           api_key="AIzaSyAuxYmci0DVU5l5L_YcxLlxHzR5MLn70js")
 
     async with stdio_client(server_params) as (read, write):
@@ -44,12 +46,12 @@ async def pg_mcp_exec(request: str) -> str:
             agent = create_react_agent(llm, tools)
             agent_response = await agent.ainvoke({"messages": prompt})
 
-            parsed_steps, _ = parse_mcp_output(agent_response)
+            parsed_steps, final_answer, last_tool_answer, _ = parse_mcp_output(agent_response)
             memory.update_from_parsed(parsed_steps, request)
 
             await handle_memory_save_or_reset(memory, request)
 
-            return agent_response
+            return  final_answer, last_tool_answer
 
 
 # ---------------- Helper Functions ---------------- #
@@ -62,7 +64,7 @@ def get_server_params() -> StdioServerParameters:
     # TODO: give server params from config
     return StdioServerParameters(
         command="python",
-        args=[r"C:\Users\yukcus\Desktop\MCPTest\postgre_mcp_server.py"],
+        args=[r"/home/amirkia/Desktop/query_mcp_server/postgre_mcp_server.py"],
     )
 
 async def load_or_create_memory() -> ConversationMemory:

postgre_mcp_server.py

@@ -31,7 +31,7 @@ class DbContext:
 async def db_lifespan(server: FastMCP) -> AsyncIterator[DbContext]:
     """Manage database connection lifecycle"""
     # Initialize DB connection from environment variables
-    dsn = os.environ.get("DSN", "postgresql://postgres:postgres@localhost:5432/postgres")
+    dsn = os.environ.get("DSN", "postgresql://postgres:postgres@localhost:5432/DNextReporting")
     schema = os.environ.get("SCHEMA", DEFAULT_SCHEMA)
     pool = await asyncpg.create_pool(dsn)
     try:
@@ -71,10 +71,21 @@ async def base_prompt_query() -> str:
 
 
                     Each tool may also return previews or summaries of table contents to help you better understand the data structure.
+                    
+                    ---
+
+                    ==========================
+                    # Your Objective
+                    ==========================
+
+                    When a user submits a request, you must:
 
-                    You also have access to **short-term memory**, which stores relevant context from earlier queries. If memory contains the needed information, you **must use it** instead of repeating tool calls with the same input. Avoid redundant tool usage unless:
-                    - The memory is empty, or
-                    - A tool's output is outdated or missing
+                    1. **Analyze the request** to determine the required data or action.
+                    2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).
+                    3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.
+                    4. **Execute the SQL query** and return the results.
+                    5. **Chain tools logically**, such as: List Tables → Get Schema → Write and Run Query.
+                    6. **Explain your reasoning and each step taken** to ensure clarity and transparency.
 
                     ---
 
@@ -106,20 +117,14 @@ async def base_prompt_query() -> str:
                     - Use memory efficiently. Don’t rerun a tool unless necessary.
                     - If you generate a SQL query, immediately call the **execute_query** tool using that query. Do not delay or wait for user confirmation.
 
-
                     ---
-
+                    
                     ==========================
-                    # Short-Term Memory
+                    # Database Description
                     ==========================
-
-                    You have access to the following memory from this conversation. Use it if applicable for the current request.
-
-                    - Previous user requests: {user_requests}
-                    - Tools used so far: {past_tools}
-                    - Last SQL queries: {last_queries}
-                    - Last result preview: {last_results}
-
+                    
+                    {descriptions}
+                    
                     ---
 
                     ==========================
@@ -144,12 +149,30 @@ async def base_prompt_query() -> str:
                       > SELECT * FROM customers WHERE country = 'Germany';
                       > ```
 
+==========================
+# Output Format
+==========================
+
+Present your final answer using the following structure **exactly**. When necessary, bold the important parts of your answer or use `` inline code blocks.:
+
+```markdown
+# Result
+{{Take the result from the execute_query tool and format it nicely using Markdown. Use a Markdown table for tabular data (rows and columns) including headers. Use bullet points or items in markdown for answers that include lists of names or descriptions. Use plain text for single values or simple messages. Ensure data alignment and clarity.}}
+
+# Explanation
+{{Provide a concise explanation or interpretation of the results in 1-3 sentences. Explain what the data in the 'Result' section represents in the context of the user's request.}}
+
+# Query
+```sql
+{{Display the exact SQL query you generated and executed here to answer the user's request.}}
 
                     **Reminder:** 
                     **Every time you generate a SQL query, call **execute_query** right after and include the result in your final response.**
                     **If you do not execute the generated SQL query, this will be the violation of the instructions**
                     **Your final answer cannot be only a SQL query, you will have to call **execute_query** and give the result of the call with the SQL query.** 
                     ---
+                    {chat_history}
+                    ---
                     =========================
                     # New User Request
                     =========================
@@ -179,8 +202,6 @@ async def last_prompt() -> str:
 You are an expert at reading and understanding SQL queries.  
 Your task is to retrieve the **exact SQL query** that produced a previously seen result, convert the query to the **ANSI SQL query** and return **only the ANSI SQL query** — no explanation, reasoning, or commentary.
 
-You have access to a **short-term memory**, which stores relevant context from earlier interactions in the current conversation.
-
 ---
 
 ==========================
@@ -189,7 +210,7 @@ You have access to a **short-term memory**, which stores relevant context from e
 
 When a user submits a request (e.g., *"send me that table"*, *"send the last query"*, etc.), follow these steps:
 
-1. Identify which previous result the user is referring to, using your short-term memory.
+1. Identify which previous result the user is referring to.
 2. Retrieve the corresponding SQL query that produced that result.
 3. Convert the SQL query to the ANSI SQL query
 3. Return **only** that ANSI SQL query.
@@ -202,35 +223,12 @@ When a user submits a request (e.g., *"send me that table"*, *"send the last que
 
 - Do **not** ask questions or request clarification.
 - Do **not** explain anything to the user.
-- Only use the **memory** to determine which query is relevant.
 - Respond with the **exact ANSI SQL query only**, formatted cleanly.
-- Do **not** guess — only retrieve queries that actually exist in memory.
+- Do **not** guess — only retrieve queries.
 - If no query fits, respond with: "Query not found."
 
 ---
 
-==========================
-# Short-Term Memory
-==========================
-
-You have access to the following memory from this conversation:
-
-- **Previous user requests**:  
-  `{user_requests}`
-
-- **Tools used so far**:  
-  `{past_tools}`
-
-- **Recent SQL queries**:  
-  `{last_queries}`
-
-- **Result preview from last query**:  
-  `{last_results}`
-
-Use this memory to resolve any references in the user's latest request.
-
----
-
 ==========================
 # Reference Conversion Rules for PostgreSQL to ANSI SQL
 ==========================
@@ -492,6 +490,7 @@ async def execute_query(
 # Database helper functions
 async def get_all_tables(pool, schema):
     """Get all tables from the database"""
+    print(f"schema: {schema}")
     async with pool.acquire() as conn:
         result = await conn.fetch("""
             SELECT c.relname AS table_name

postgre_smolagent_clinet.py

@@ -8,6 +8,8 @@ from conversation_memory import ConversationMemory
 from utils import parse_mcp_output, classify_intent
 import logging
 from smolagents import LiteLLMModel, ToolCollection, CodeAgent
+from dotenv import load_dotenv
+
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +30,7 @@ async def pg_mcp_smolagent_exec(request: str) -> str:
     server_params = get_server_params()
 
     # TODO: give key from env
-    llm = LiteLLMModel(model_id="gemini/gemini-2.0-flash-lite", api_key="AIzaSyAuxYmci0DVU5l5L_YcxLlxHzR5MLn70js")
+    llm = LiteLLMModel(model_id="gemini/gemini-2.0-flash", api_key="AIzaSyDRoQ8yG0Rmpt3NtfHS7fdm5hgfL9su5aY")
 
     async with stdio_client(server_params) as (read, write):
         async with ClientSession(read, write) as session:
@@ -42,14 +44,15 @@ async def pg_mcp_smolagent_exec(request: str) -> str:
                 tools = await load_and_enrich_tools(session, table_summary)
                 past_data = get_memory_snapshot(memory)
 
-                prompt = await build_prompt(session, intent, request, tools, past_data)
+                prompt = await build_prompt(session, intent, request, tools, past_data, table_summary)
                 agent_response = agent.run(task=prompt, stream=False)
 
+                # TODO: add a smolagent output parser
 
-                parsed_steps, _ = parse_mcp_output(agent_response)
-                memory.update_from_parsed(parsed_steps, request)
+                #parsed_steps, _ = parse_mcp_output(agent_response)
+                #memory.update_from_parsed(parsed_steps, request)
 
-                await handle_memory_save_or_reset(memory, request)
+                #await handle_memory_save_or_reset(memory, request)
 
             return agent_response
 
@@ -62,9 +65,10 @@ def load_table_summary(path: str) -> str:
 
 def get_server_params() -> StdioServerParameters:
     # TODO: give server params from config
+    load_dotenv()
     return StdioServerParameters(
         command="python",
-        args=[r"C:\Users\yukcus\Desktop\MCPTest\postgre_mcp_server.py"],
+        args=[os.environ["MCP_SERVER_PATH"]],
     )
 
 async def load_or_create_memory() -> ConversationMemory:
@@ -75,8 +79,8 @@ async def load_or_create_memory() -> ConversationMemory:
 
 async def load_and_enrich_tools(session: ClientSession, summary: str):
     tools = await load_mcp_tools(session)
-    for tool in tools:
-        tool.description += f" {summary}"
+    #for tool in tools:
+    #    tool.description += f" {summary}"
     return tools
 
 def get_memory_snapshot(memory: ConversationMemory) -> dict:
@@ -94,7 +98,7 @@ def get_memory_snapshot(memory: ConversationMemory) -> dict:
         "past_requests": "No requests found"
     }
 
-async def build_prompt(session, intent, request, tools, past_data):
+async def build_prompt(session, intent, request, tools, past_data, summary):
     superset_prompt = await session.read_resource("resource://last_prompt")
     conversation_prompt = await session.read_resource("resource://base_prompt")
     # TODO: add uri's from config
@@ -116,7 +120,8 @@ async def build_prompt(session, intent, request, tools, past_data):
             last_queries=past_data["past_queries"],
             last_results=past_data["past_results"],
             new_request=request,
-            tools=tools_str
+            tools=tools_str,
+            descriptions=summary
         )
 
 async def handle_memory_save_or_reset(memory: ConversationMemory, request: str):

run.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+# Replace 'myenv' with the name of your conda environment
+# conda activate myenv
+
+python gradio_app.py

streamlit_app.py

@@ -0,0 +1,38 @@
+import yaml
+from pathlib import Path
+import streamlit as st
+from postgre_mcp_client import pg_mcp_exec
+import asyncio
+import nest_asyncio
+nest_asyncio.apply()
+
+def load_db_configs():
+    """Load database configurations from configs.yaml"""
+    configs_path = Path("configs.yaml")
+
+    if not configs_path.exists():
+        st.error("configs.yaml not found")
+        return None
+
+    with open(configs_path) as f:
+        configs = yaml.safe_load(f)
+
+    return configs.get("db_configs", {})
+
+def run_agent(message):
+    response = asyncio.run(pg_mcp_exec(message))
+    # Return in message format
+    return {"role": "assistant", "content": response}
+
+# Streamlit UI
+st.title("PostgreSQL Query Agent")
+st.write("Ask your database in natural language and get results using the smolagent executor.")
+
+user_input = st.text_input("Natural Language Request", placeholder="e.g., Show me the table of join posts and users tables.")
+
+if st.button("Run Query"):
+    if user_input.strip():
+        result = run_agent(user_input)
+        st.text_area("SQL Query / Result", value=str(result), height=300)
+    else:
+        st.warning("Please enter a natural language request.")

table_summary.txt

@@ -1,3 +1,342 @@
-The users table stores information about the individuals who use the application. Each user is assigned a unique, auto-incrementing id that serves as the primary key. The username field holds the user's chosen display name and cannot be null, while the email field stores the user’s unique email address, also required and constrained to be unique to avoid duplicates. To track when a user was added to the system, the created_at column records the timestamp of their creation, with a default value set to the current time.
+## Table: `dim_agreement`
 
-The posts table represents content created by users, such as blog posts or messages. Like the users table, each entry has a unique, auto-incrementing id as the primary key. The user_id field links each post to its author by referencing the id field in the users table, establishing a one-to-many relationship between users and posts. The title column holds a brief summary or headline of the post, while the content field contains the full text. A created_at timestamp is also included to record when each post was created, with a default value of the current time.
+### Schema: `public`
+
+### All Columns and their Descriptions:
+- `agreement_id` (text): Unique identifier for the agreement. This is likely the primary key.
+- `customer_id` (text):  Identifier for the customer associated with the agreement.  This is likely a foreign key referencing a customer table.
+- `agreement_type` (text):  The type of agreement (e.g., Service Level Agreement, Sales Agreement).
+- `agreement_name` (text):  The name or title of the agreement.
+- `agreement_status` (text): The current status of the agreement (e.g., Active, Inactive, Expired).
+- `description` (text):  A textual description of the agreement.
+- `document_number` (text): The document number associated with the agreement.
+- `statement_of_intent` (text):  Indicates whether a statement of intent is associated with the agreement.
+- `version` (text): The version number of the agreement.
+- `initial_date` (timestamp without time zone): The date the agreement was initially created.
+- `agreement_period_start_date` (timestamp without time zone): The start date of the agreement period.
+- `agreement_period_end_date` (timestamp without time zone): The end date of the agreement period.
+- `completion_date_start_date` (timestamp without time zone): The start date of the completion date.
+- `completion_date_end_date` (text): The end date of the completion date.
+- `auto_renewal` (text): Indicates whether the agreement has auto-renewal enabled.
+- `commitment_unit` (text): The unit for the commitment amount (e.g., USD, hours).
+- `rating_type` (text): The type of rating associated with the agreement.
+- `commitment_amount` (text): The committed amount associated with the agreement.
+- `remaining_payment_cycle` (text): The remaining payment cycle.
+- `termination_date` (text): The date the agreement was terminated.
+
+### Relationships with Other Tables (Foreign Keys):
+- `customer_id` (FK) references a `customer` table (likely `dim_customer` or similar) on the `customer_id` column. This links the agreement to a specific customer.
+
+### Cardinality of Relationships:
+- Relationship with `dim_customer` (or similar) is one-to-many. One customer can have multiple agreements.
+
+### Common Use Cases/Example Queries:
+- Describe all active agreements for a specific customer.
+- Analyze agreements by type or status.
+- Track agreement start and end dates.
+- Example SQL Snippet: `SELECT * FROM dim_agreement WHERE customer_id = 'customer_id_value' AND agreement_status = 'Active';`
+
+### Data Constraints and Business Rules:
+- `agreement_id` is likely unique and a primary key.
+- `agreement_status` might have a controlled vocabulary (e.g., Active, Inactive).
+- Date fields should be validated for data integrity.
+- Business rules may dictate agreement approval processes or renewal cycles.
+
+### Data Update Frequency/Volatility:
+- Data is likely updated when agreements are created, modified, or terminated.
+- Update frequency would depend on the business processes.
+
+### Potential Pitfalls/Things to Avoid When Querying:
+- Ensure correct date ranges when filtering by agreement period or completion dates.
+- Be aware of potential data quality issues in text fields.
+
+### Important Notes/Considerations for Querying:
+- Join with the `dim_customer` table to retrieve customer information.
+- Consider indexing frequently queried columns (e.g., `customer_id`, `agreement_status`, `agreement_type`).
+
+---
+
+## Table: `dim_customer`
+
+### Schema: `public`
+
+### All Columns and their Descriptions:
+-   `customer_id` (text):  Likely the primary identifier for a customer.  This is probably the primary key, although not explicitly stated.
+-   `customer_name` (text): The name of the customer.
+-   `customer_start_date` (timestamp without time zone): The date the customer record became active.
+-   `customer_end_date` (timestamp without time zone): The date the customer record became inactive.
+-   `customer_status` (text): The status of the customer (e.g., Active, Inactive).
+-   `is_customer_active` (integer):  Indicates whether the customer is currently active (1 or 0).
+-   `customer_status_reason` (text): Reason for the customer's status.
+-   `market_segment` (text): The market segment the customer belongs to.
+-   `customer_segment` (text): The customer segment the customer belongs to.
+-   `is_individual_customer` (integer): Indicates whether the customer is an individual (1 or 0).
+-   `organization_name` (text): The name of the organization, if the customer is an organization.
+-   `trading_name` (text): The trading name of the customer.
+-   `organization_type_id` (text): The ID of the organization type.
+-   `is_head_office` (boolean): Indicates if the customer is a head office.
+-   `is_legal_entity` (boolean): Indicates if the customer is a legal entity.
+-   `organization_start_date` (timestamp without time zone): The date the organization record became active.
+-   `organization_end_date` (timestamp without time zone): The date the organization record became inactive.
+-   `birth_date` (timestamp without time zone): The birth date of the customer, if an individual.
+-   `death_date` (timestamp without time zone): The death date of the customer, if an individual.
+-   `country_of_birth` (text): The country of birth of the customer, if an individual.
+-   `party_status` (text): The status of the party.
+-   `title` (text): The title of the customer, if an individual.
+-   `generation` (text): The generation of the customer, if an individual.
+-   `preferred_given_name` (text): The preferred given name of the customer, if an individual.
+-   `given_name` (text): The given name of the customer, if an individual.
+-   `middle_name` (text): The middle name of the customer, if an individual.
+-   `family_name` (text): The family name of the customer, if an individual.
+-   `family_name_prefix` (text): The family name prefix of the customer, if an individual.
+-   `formatted_name` (text): The formatted name of the customer.
+-   `full_name` (text): The full name of the customer.
+-   `legal_name` (text): The legal name of the customer.
+-   `gender` (text): The gender of the customer, if an individual.
+-   `location` (text): The location of the customer.
+-   `marital_status` (text): The marital status of the customer, if an individual.
+-   `nationality` (text): The nationality of the customer, if an individual.
+-   `place_of_birth` (text): The place of birth of the customer, if an individual.
+-   `juridical_info` (text): Juridical information of the customer.
+-   `latitude` (text): The latitude of the customer's location.
+-   `longitude` (text): The longitude of the customer's location.
+-   `number_of_employees` (text): The number of employees for the customer, if an organization.
+-   `tax_exempt` (text): Indicates if the customer is tax-exempt.
+-   `dunning_level` (text): The dunning level for the customer.
+-   `contact_verified` (text): Indicates if the contact information is verified.
+-   `customer_number` (text): The customer number.
+-   `communication_method` (text): The communication method for the customer.
+-   `sales_agent_id` (text): The ID of the sales agent.  Likely a foreign key to a `sales_agent` table.
+-   `account_manager_id` (text): The ID of the account manager. Likely a foreign key to an `account_manager` table.
+-   `sales_partner_id` (text): The ID of the sales partner. Likely a foreign key to a `sales_partner` table.
+-   `party_identifier` (text): The party identifier.
+-   `customer_revision` (bigint): The customer revision number.
+-   `party_revision` (bigint): The party revision number.
+
+### Relationships with Other Tables (Foreign Keys):
+-   Based on the column names, it is *highly likely* that the table has foreign keys to tables such as: `sales_agent`, `account_manager`, and `sales_partner`. Without the ability to see the foreign key constraints, this is an educated guess.
+
+### Cardinality of Relationships:
+-   Relationship with `sales_agent`, `account_manager`, and `sales_partner` is likely one-to-many. One sales agent/account manager/sales partner can be associated with multiple customers.
+
+### Common Use Cases/Example Queries:
+-   Describe common use cases:
+    *   Reporting on customer demographics.
+    *   Analyzing customer sales by segment.
+    *   Tracking customer status and activity.
+    *   Identifying customers by sales agent or account manager.
+-   Example SQL Snippet:
+    ```sql
+    SELECT customer_name, market_segment, customer_status
+    FROM dim_customer
+    WHERE sales_agent_id = 'XYZ123';
+    ```
+
+### Data Constraints and Business Rules:
+-   `customer_id` is likely the primary key and should be unique and not null.
+-   `customer_start_date` should be before or equal to `customer_end_date`.
+-   `is_customer_active` should reflect the status of the customer based on `customer_start_date`, `customer_end_date`, and `customer_status`.
+-   Data quality checks should be in place to ensure the accuracy of demographic data (e.g., valid gender, valid countries).
+
+### Data Update Frequency/Volatility:
+-   The data is likely updated frequently as customer information changes.
+-   The volatility of the data will vary depending on the specific attributes (e.g., contact information will be more volatile than birthdate).
+
+### Potential Pitfalls/Things to Avoid When Querying:
+-   Joining to other tables without properly understanding the relationships, especially the role of `sales_agent_id`, `account_manager_id`, and `sales_partner_id`.
+-   Incorrectly filtering on date ranges.  Consider using `customer_start_date` and `customer_end_date` to filter for active customers.
+-   Not considering the impact of customer status when analyzing customer data.
+
+### Important Notes/Considerations for Querying:
+-   Always check `is_customer_active` or `customer_status` when querying to ensure you are getting the correct set of customers.
+-   Be mindful of the date ranges when analyzing customer data over time.
+-   Use the appropriate join conditions when joining with other tables.
+-   Consider the impact of null values in fields such as `birth_date`, `death_date`, etc.
+
+
+---
+
+
+## Table: `dim_product`
+
+### Schema: `public`
+
+### All Columns and their Descriptions:
+
+*   **product\_id (text):** Likely the primary identifier for the product. It's a text field. It is likely the primary key, though I don't have FK information.
+*   **customer\_id (text):** The identifier for the customer associated with the product. Foreign key to the `dim_customer` table (inferred, as I lack foreign key info).
+*   **agreement\_id (text):**  An identifier for the agreement related to the product.  Foreign key to the `dim_agreement` table (inferred).
+*   **agreement\_name (text):**  The name of the agreement related to the product.
+*   **product\_offering\_id (text):**  An identifier for the product offering. Potentially a foreign key to a `dim_product_offering` table (inferred).
+*   **product\_offering\_name (text):** The name of the product offering.
+*   **product\_name (text):** The name of the product.
+*   **place\_id (text):** An identifier related to the product's location or place. Foreign key to `dim_place` (inferred).
+*   **product\_class (text):** Categorization of the product (e.g., service, hardware).
+*   **product\_status (text):** The current status of the product (e.g., active, inactive).
+*   **is\_product\_active (integer):** Indicates if the product is active (1) or not (0).
+*   **is\_product\_suspended (integer):** Indicates if the product is suspended (1) or not (0).
+*   **is\_product\_in\_active (integer):** Indicates if the product is inactive (1) or not (0).
+*   **is\_bundle (boolean):** Indicates if the product is part of a bundle (true) or not (false).
+*   **order\_date (timestamp without time zone):** The date the product was ordered.
+*   **start\_date (timestamp without time zone):** The date the product service started.
+*   **price\_type (text):** The type of pricing associated with the product (e.g., recurring, one-time).
+*   **price\_type\_value (text):** The value associated with the price type.
+*   **recurring\_charge\_period (text):**  The period for recurring charges (e.g., monthly, annually).
+*   **tax\_rate (double precision):** The tax rate applied to the product.
+*   **duty\_free\_amount (double precision):** The duty-free amount for the product.
+*   **tax\_included\_amount (double precision):** The amount of tax included in the product price.
+*   **add\_product\_order\_item\_id (text):** Identifier for adding a product to an order item.
+*   **delete\_product\_order\_item\_id (text):** Identifier for deleting a product from an order item.
+*   **sales\_agent\_id (text):** The identifier of the sales agent. Foreign key to `dim_sales_agent` (inferred).
+*   **sales\_partner\_id (text):** The identifier of the sales partner. Foreign key to `dim_sales_partner` (inferred).
+*   **commitment\_duration\_units (text):** The units for the commitment duration (e.g., months, years).
+*   **commitment\_duration (double precision):** The duration of the commitment.
+*   **commitment\_term\_name (text):** The name of the commitment term.
+*   **commitment\_term\_type (text):** The type of commitment term.
+*   **usage\_duration\_units (text):** The units for the usage duration.
+*   **usage\_duration (double precision):** The duration of the product usage.
+*   **usage\_term\_name (text):** The name of the usage term.
+*   **usage\_term\_type (text):** The type of usage term.
+*   **guarantee\_amount (text):** The amount of guarantee associated with the product.
+*   **device\_type (text):** The type of device associated with the product.
+*   **gl\_code (text):** General Ledger code associated with the product.
+*   **infrastructure (text):** The infrastructure associated with the product.
+*   **ip\_address (text):** The IP address associated with the product.
+*   **mac\_address (text):** The MAC address associated with the product.
+*   **oss\_code (text):** OSS code associated with the product.
+*   **smart\_card\_serialnumber (text):** The serial number of the smart card associated with the product.
+*   **sla (text):** Service Level Agreement associated with the product.
+*   **spec\_type (text):** Specification type of the product.
+*   **specsub\_type (text):** Specification subtype of the product.
+*   **resource\_model (text):** The resource model associated with the product.
+*   **rating\_type (text):** The type of rating associated with the product.
+*   **postpaid\_type (text):** The postpaid type associated with the product.
+*   **brand\_name (text):** The brand name of the product.
+*   **tv\_infrastructure (text):** The TV infrastructure associated with the product.
+*   **revision (bigint):** Revision number of the product.
+*   **href (text):** Hypertext reference (URL) associated with the product.
+
+### Relationships with Other Tables (Foreign Keys):
+
+*   **customer\_id:** Foreign key referencing `dim_customer`.
+*   **agreement\_id:** Foreign key referencing `dim_agreement`.
+*   **product\_offering\_id:** Foreign key referencing `dim_product_offering`.
+*   **place\_id:** Foreign key referencing `dim_place`.
+*   **sales\_agent\_id:** Foreign key referencing `dim_sales_agent`.
+*   **sales\_partner\_id:** Foreign key referencing `dim_sales_partner`.
+
+### Cardinality of Relationships:
+
+*   Relationship with `dim_customer` is likely one-to-many (one customer can have many products).
+*   Relationship with `dim_agreement` is likely one-to-many (one agreement can have many products).
+*   Relationship with `dim_product_offering` is likely one-to-many (one product offering can be associated with many products).
+*   Relationship with `dim_place` is likely one-to-many (one place can have many products).
+*   Relationship with `dim_sales_agent` is likely one-to-many (one sales agent can be associated with many products).
+*   Relationship with `dim_sales_partner` is likely one-to-many (one sales partner can be associated with many products).
+
+### Common Use Cases/Example Queries:
+
+*   **Tracking Product Details:** Retrieving detailed information about a specific product.
+    *   Example SQL Snippet: `SELECT * FROM dim_product WHERE product_id = 'your_product_id';`
+*   **Customer Product Overview:**  Listing all products associated with a specific customer.
+    *   Example SQL Snippet: `SELECT * FROM dim_product WHERE customer_id = 'your_customer_id';`
+*   **Aggregating product counts:** Calculating the number of products per product offering.
+    *   Example SQL Snippet: `SELECT product_offering_name, count(*) FROM dim_product GROUP BY product_offering_name;`
+
+### Data Constraints and Business Rules:
+
+*   `product_id` is likely unique and not null.
+*   `is_product_active`, `is_product_suspended`, and `is_product_in_active` should be mutually exclusive (only one can be true at a time).
+*   Date fields (order\_date, start\_date) should be consistent with business timelines.
+
+### Data Update Frequency/Volatility:
+
+*   Data is likely updated frequently, with changes reflecting product status, customer associations, and order information.
+
+### Potential Pitfalls/Things to Avoid When Querying:
+
+*   Ensure proper joins when querying across related tables (e.g., `dim_customer`, `dim_agreement`).
+*   Be aware of product lifecycle (status) when filtering data.  Consider `is_product_active` and other status flags.
+*   Performance can suffer if large datasets are queried without appropriate indexes.
+
+### Important Notes/Considerations for Querying:
+
+*   Always join to other dimension tables (e.g., `dim_customer`, `dim_agreement`) to get a complete view of the product's context.
+*   Use appropriate date ranges when analyzing product activity over time.
+*   Consider the business rules related to product states when building queries.
+
+---
+
+
+## Table: `dim_product_order_item`
+
+### Schema: `public`
+
+### All Columns and their Descriptions:
+
+-   `product_order_item_id` (text): Unique identifier for a product order item. Likely the primary key.
+-   `product_order_id` (text): Likely a foreign key referencing a table containing product order information (e.g., `dim_product_order`).
+-   `customer_id` (text): Likely a foreign key referencing the `dim_customer` table.
+-   `agreement_id` (text): Likely a foreign key referencing the `dim_agreement` table.
+-   `agreement_name` (text): Name of the agreement.
+-   `order_item_agreement_id` (text): Identifier for the order item agreement.
+-   `order_item_agreement_name` (text): Name of the order item agreement.
+-   `billing_account_id` (text): Identifier for the billing account.
+-   `sales_channel` (text): The sales channel through which the order was placed.
+-   `order_date` (timestamp without time zone): The date the order was placed.
+-   `order_status` (text): The status of the order (e.g., 'Pending', 'Shipped', 'Delivered').
+-   `order_item_status` (text): The status of the order item.
+-   `order_description` (text): Description of the order.
+-   `order_revision` (bigint): The revision number of the order.
+-   `order_class` (text): The class of the order.
+-   `order_href` (text): Hyperlink related to the order.
+-   `order_item_type` (text): The type of the order item (e.g., 'Product', 'Service').
+-   `quantity` (bigint): The quantity of the product ordered.
+-   `order_item_price_type` (text): The price type of the order item (e.g., 'Recurring', 'One-time').
+-   `order_item_recurring_charge_period` (text): The recurring charge period if applicable.
+-   `order_item_tax_rate` (double precision): The tax rate applied to the order item.
+-   `order_item_duty_free_amount` (double precision): The duty-free amount of the order item.
+-   `order_item_tax_included_amount` (double precision): The tax-included amount of the order item.
+-   `order_item_action` (text): The action performed on the order item (e.g., 'Add', 'Update', 'Delete').
+
+### Relationships with Other Tables (Foreign Keys):
+
+-   `customer_id` (FK) references `dim_customer`: Links the order item to a specific customer.
+-   `agreement_id` (FK) references `dim_agreement`: Links the order item to a specific agreement.
+-   `product_order_id`:  Likely references a table containing product order information (e.g., `dim_product_order`).
+
+### Cardinality of Relationships:
+
+-   Relationship with `dim_customer` is one-to-many. One customer can have multiple order items.
+-   Relationship with `dim_agreement` is one-to-many. One agreement can have multiple order items.
+-   Relationship with the product order table is one-to-many. One product order can have multiple order items.
+
+### Common Use Cases/Example Queries:
+
+-   Track sales and revenue by customer, product, and agreement.
+-   Analyze order item status and trends.
+-   Calculate total revenue and taxes.
+-   Example SQL Snippet: `SELECT sum(quantity * order_item_tax_included_amount) FROM dim_product_order_item WHERE customer_id = 'customer123';`
+
+### Data Constraints and Business Rules:
+
+-   `product_order_item_id` should be unique.
+-   `order_date` should be a valid date.
+-   `quantity` should be a non-negative number.
+-   `order_item_tax_rate`, `order_item_duty_free_amount`, and `order_item_tax_included_amount` should be valid numeric values.
+
+### Data Update Frequency/Volatility:
+
+-   Data is likely updated frequently, reflecting new orders, order updates, and cancellations.
+
+### Potential Pitfalls/Things to Avoid When Querying:
+
+-   Ensure correct join conditions when joining with other tables.
+-   Be mindful of the `order_item_status` to filter for relevant order items.
+-   Consider time-based filtering using `order_date` for trend analysis.
+
+### Important Notes/Considerations for Querying:
+
+-   This table is central to understanding product order details.
+-   Join with `dim_customer` and `dim_agreement` for customer and agreement details.
+-   The table containing product order information (e.g. `dim_product_order`) will be crucial to join.

utils.py

@@ -8,7 +8,8 @@ def parse_mcp_output(output_dict):
     result = []
     messages = output_dict.get("messages", [])
     query_store = []
-
+    last_tool_answer = ""
+    last_answer = ""
     for msg in messages:
         role_name = msg.__class__.__name__  # Example: HumanMessage, AIMessage, ToolMessage
         content = getattr(msg, "content", "")
@@ -33,9 +34,11 @@ def parse_mcp_output(output_dict):
                 # Check for presence of "query" key
                 if "query" in arguments_dict:
                     #print("query detected!!!")
-                    print(f"=============== AI Reasoning Step ===============")
-                    print(content[0])
-                    print()
+
+                    if content:
+                        print(f"=============== AI Reasoning Step ===============")
+                        print(content[0])
+                        print()
                     print("=============== AI used the following tools ===============")
                     print(tool_name)
                     print()
@@ -73,6 +76,7 @@ def parse_mcp_output(output_dict):
                     })
 
             else:
+                final_answer = content
                 #print(f"ai final answer:{content}")
                 logger.info(f"ai final answer:{content}")
                 print("=============== AI's final answer ===============")
@@ -87,6 +91,7 @@ def parse_mcp_output(output_dict):
             tool_name = getattr(msg, "name", None)
             print("=============== The tool returned the following response ===============")
             print(content)
+            last_tool_answer = content
             logger.info(f"tool response:{content}")
             result.append({
                 "type": "tool_response",
@@ -94,7 +99,7 @@ def parse_mcp_output(output_dict):
                 "response": content
             })
 
-    return result, query_store
+    return result, final_answer, last_tool_answer, query_store