Merge remote-tracking branch 'origin/selcan_test' into selcan_test
Browse files# Conflicts:
# gradio_app.py
# memory.json
# postgre_mcp_client.py
# postgre_smolagent_clinet.py
# table_summary.txt
- gradio_app.py +11 -6
- memory.json +0 -1
- postgre_mcp_client.py +1 -1
- postgre_mcp_server.py +17 -0
- postgre_smolagent_clinet.py +2 -2
- table_summary.txt +7 -7
gradio_app.py
CHANGED
|
@@ -4,8 +4,6 @@ import gradio as gr
|
|
| 4 |
import asyncio
|
| 5 |
from postgre_mcp_client import pg_mcp_exec
|
| 6 |
from postgre_smolagent_clinet import pg_mcp_smolagent_exec
|
| 7 |
-
from postgre_smolagent_clinet import pg_mcp_smolagent_exec
|
| 8 |
-
|
| 9 |
|
| 10 |
def load_db_configs():
|
| 11 |
"""Load database configurations from configs.yaml"""
|
|
@@ -24,8 +22,9 @@ async def run_agent(request):
|
|
| 24 |
# configs = load_db_configs()
|
| 25 |
# final_answer, last_tool_answer, = await pg_mcp_exec(request)
|
| 26 |
# return final_answer, last_tool_answer
|
| 27 |
-
|
| 28 |
-
|
|
|
|
| 29 |
|
| 30 |
# Gradio UI
|
| 31 |
demo = gr.Interface(
|
|
@@ -34,9 +33,15 @@ demo = gr.Interface(
|
|
| 34 |
label="Natural Language Request",
|
| 35 |
placeholder="e.g., Show me the table of join posts and users tables."
|
| 36 |
),
|
| 37 |
-
outputs=gr.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 38 |
title="PostgreSQL Query Agent",
|
| 39 |
-
description="Ask your database in natural language and get results."
|
|
|
|
| 40 |
)
|
| 41 |
|
| 42 |
if __name__ == "__main__":
|
|
|
|
| 4 |
import asyncio
|
| 5 |
from postgre_mcp_client import pg_mcp_exec
|
| 6 |
from postgre_smolagent_clinet import pg_mcp_smolagent_exec
|
|
|
|
|
|
|
| 7 |
|
| 8 |
def load_db_configs():
|
| 9 |
"""Load database configurations from configs.yaml"""
|
|
|
|
| 22 |
# configs = load_db_configs()
|
| 23 |
# final_answer, last_tool_answer, = await pg_mcp_exec(request)
|
| 24 |
# return final_answer, last_tool_answer
|
| 25 |
+
|
| 26 |
+
result = await pg_mcp_smolagent_exec(request)
|
| 27 |
+
return result
|
| 28 |
|
| 29 |
# Gradio UI
|
| 30 |
demo = gr.Interface(
|
|
|
|
| 33 |
label="Natural Language Request",
|
| 34 |
placeholder="e.g., Show me the table of join posts and users tables."
|
| 35 |
),
|
| 36 |
+
outputs=gr.Markdown(
|
| 37 |
+
label="SQL Query / Result",
|
| 38 |
+
value="# Result \n \n",
|
| 39 |
+
# max_height=1500,
|
| 40 |
+
height= 1500
|
| 41 |
+
),
|
| 42 |
title="PostgreSQL Query Agent",
|
| 43 |
+
description="Ask your database in natural language and get results.",
|
| 44 |
+
flagging_mode="never"
|
| 45 |
)
|
| 46 |
|
| 47 |
if __name__ == "__main__":
|
memory.json
DELETED
|
@@ -1 +0,0 @@
|
|
| 1 |
-
{"tools_used": ["list_tables", "get_table_schema", "execute_query", "execute_query", "execute_query", "execute_query", "execute_query", "list_tables", "list_tables", "list_tables", "list_tables"], "all_queries": ["SELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at\nFROM posts p\nJOIN users u ON p.user_id = u.id;", "```sql\nSELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at\nFROM posts p\nJOIN users u ON p.user_id = u.id;\n```", "SELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at\nFROM posts p\nJOIN users u ON p.user_id = u.id;", "```sql\nSELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at\nFROM posts p\nJOIN users u ON p.user_id = u.id;\n```", "```sql\nSELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at\nFROM posts p\nJOIN users u ON p.user_id = u.id;\n```", "SELECT p.id, p.user_id, p.title, p.content, p.created_at, u.id, u.username, u.email, u.created_at\nFROM posts p\nJOIN users u ON p.user_id = u.id;"], "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", "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 14:59:14.801730 | 1 | alice | alice@example.com | 2025-04-09 14:59:14.801730\n2 | 2 | Bob\u2019s Thoughts | Here are some thoughts... | 2025-04-09 14:59:14.801730 | 2 | bob | bob@example.com | 2025-04-09 14:59:14.801730\n3 | 1 | Follow-up Post | Another one from Alice. | 2025-04-09 14:59:14.801730 | 1 | alice | alice@example.com | 2025-04-09 14:59:14.801730", "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 14:59:14.801730 | 1 | alice | alice@example.com | 2025-04-09 14:59:14.801730\n2 | 2 | Bob\u2019s Thoughts | Here are some thoughts... | 2025-04-09 14:59:14.801730 | 2 | bob | bob@example.com | 2025-04-09 14:59:14.801730\n3 | 1 | Follow-up Post | Another one from Alice. | 2025-04-09 14:59:14.801730 | 1 | alice | alice@example.com | 2025-04-09 14:59:14.801730", "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 14:59:14.801730 | 1 | alice | alice@example.com | 2025-04-09 14:59:14.801730\n2 | 2 | Bob\u2019s Thoughts | Here are some thoughts... | 2025-04-09 14:59:14.801730 | 2 | bob | bob@example.com | 2025-04-09 14:59:14.801730\n3 | 1 | Follow-up Post | Another one from Alice. | 2025-04-09 14:59:14.801730 | 1 | alice | alice@example.com | 2025-04-09 14:59:14.801730"], "request": ["Show me the table of join posts and users tables.", "Show me the table of join posts and users tables.", "send to superset", "Show me the table of join posts and users tables.", "Show me the table of join posts and users tables.", "send to superset", "hi", "list all tables", "list all tables", "list all tables", "list all tables"], "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."}]], "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
CHANGED
|
@@ -29,7 +29,7 @@ async def pg_mcp_exec(request: str) -> tuple[Any, Any]:
|
|
| 29 |
server_params = get_server_params()
|
| 30 |
|
| 31 |
# TODO: give key from env
|
| 32 |
-
llm = init_chat_model(model="gemini-
|
| 33 |
api_key="AIzaSyAuxYmci0DVU5l5L_YcxLlxHzR5MLn70js")
|
| 34 |
|
| 35 |
async with stdio_client(server_params) as (read, write):
|
|
|
|
| 29 |
server_params = get_server_params()
|
| 30 |
|
| 31 |
# TODO: give key from env
|
| 32 |
+
llm = init_chat_model(model="gemini-2.0-flash", model_provider="google_genai",
|
| 33 |
api_key="AIzaSyAuxYmci0DVU5l5L_YcxLlxHzR5MLn70js")
|
| 34 |
|
| 35 |
async with stdio_client(server_params) as (read, write):
|
postgre_mcp_server.py
CHANGED
|
@@ -167,6 +167,22 @@ async def base_prompt_query() -> str:
|
|
| 167 |
> SELECT * FROM customers WHERE country = 'Germany';
|
| 168 |
> ```
|
| 169 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 170 |
|
| 171 |
**Reminder:**
|
| 172 |
**Every time you generate a SQL query, call **execute_query** right after and include the result in your final response.**
|
|
@@ -515,6 +531,7 @@ async def execute_query(
|
|
| 515 |
# Database helper functions
|
| 516 |
async def get_all_tables(pool, schema):
|
| 517 |
"""Get all tables from the database"""
|
|
|
|
| 518 |
async with pool.acquire() as conn:
|
| 519 |
result = await conn.fetch("""
|
| 520 |
SELECT c.relname AS table_name
|
|
|
|
| 167 |
> SELECT * FROM customers WHERE country = 'Germany';
|
| 168 |
> ```
|
| 169 |
|
| 170 |
+
==========================
|
| 171 |
+
# Output Format
|
| 172 |
+
==========================
|
| 173 |
+
|
| 174 |
+
Present your final answer using the following structure **exactly**. When necessary, bold the important parts of your answer or use `` inline code blocks.:
|
| 175 |
+
|
| 176 |
+
```markdown
|
| 177 |
+
# Result
|
| 178 |
+
{{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.}}
|
| 179 |
+
|
| 180 |
+
# Explanation
|
| 181 |
+
{{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.}}
|
| 182 |
+
|
| 183 |
+
# Query
|
| 184 |
+
```sql
|
| 185 |
+
{{Display the exact SQL query you generated and executed here to answer the user's request.}}
|
| 186 |
|
| 187 |
**Reminder:**
|
| 188 |
**Every time you generate a SQL query, call **execute_query** right after and include the result in your final response.**
|
|
|
|
| 531 |
# Database helper functions
|
| 532 |
async def get_all_tables(pool, schema):
|
| 533 |
"""Get all tables from the database"""
|
| 534 |
+
print(f"schema: {schema}")
|
| 535 |
async with pool.acquire() as conn:
|
| 536 |
result = await conn.fetch("""
|
| 537 |
SELECT c.relname AS table_name
|
postgre_smolagent_clinet.py
CHANGED
|
@@ -28,7 +28,7 @@ async def pg_mcp_smolagent_exec(request: str) -> str:
|
|
| 28 |
server_params = get_server_params()
|
| 29 |
|
| 30 |
# TODO: give key from env
|
| 31 |
-
llm = LiteLLMModel(model_id="gemini/gemini-2.0-flash
|
| 32 |
|
| 33 |
async with stdio_client(server_params) as (read, write):
|
| 34 |
async with ClientSession(read, write) as session:
|
|
@@ -65,7 +65,7 @@ def get_server_params() -> StdioServerParameters:
|
|
| 65 |
# TODO: give server params from config
|
| 66 |
return StdioServerParameters(
|
| 67 |
command="python",
|
| 68 |
-
args=[r"
|
| 69 |
)
|
| 70 |
|
| 71 |
async def load_or_create_memory() -> ConversationMemory:
|
|
|
|
| 28 |
server_params = get_server_params()
|
| 29 |
|
| 30 |
# TODO: give key from env
|
| 31 |
+
llm = LiteLLMModel(model_id="gemini/gemini-2.0-flash", api_key="AIzaSyDRoQ8yG0Rmpt3NtfHS7fdm5hgfL9su5aY")
|
| 32 |
|
| 33 |
async with stdio_client(server_params) as (read, write):
|
| 34 |
async with ClientSession(read, write) as session:
|
|
|
|
| 65 |
# TODO: give server params from config
|
| 66 |
return StdioServerParameters(
|
| 67 |
command="python",
|
| 68 |
+
args=[r"/home/amirkia/Desktop/query_mcp_server/postgre_mcp_server.py"],
|
| 69 |
)
|
| 70 |
|
| 71 |
async def load_or_create_memory() -> ConversationMemory:
|
table_summary.txt
CHANGED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
## Table: `dim_agreement`
|
| 2 |
|
| 3 |
-
### Schema: `
|
| 4 |
|
| 5 |
### All Columns and their Descriptions:
|
| 6 |
- `agreement_id` (text): Unique identifier for the agreement. This is likely the primary key.
|
|
@@ -58,7 +58,7 @@
|
|
| 58 |
|
| 59 |
## Table: `dim_customer`
|
| 60 |
|
| 61 |
-
### Schema: `
|
| 62 |
|
| 63 |
### All Columns and their Descriptions:
|
| 64 |
- `customer_id` (text): Likely the primary identifier for a customer. This is probably the primary key, although not explicitly stated.
|
|
@@ -159,7 +159,7 @@
|
|
| 159 |
|
| 160 |
## Table: `dim_product`
|
| 161 |
|
| 162 |
-
### Schema: `
|
| 163 |
|
| 164 |
### All Columns and their Descriptions:
|
| 165 |
|
|
@@ -237,11 +237,11 @@
|
|
| 237 |
### Common Use Cases/Example Queries:
|
| 238 |
|
| 239 |
* **Tracking Product Details:** Retrieving detailed information about a specific product.
|
| 240 |
-
* Example SQL Snippet: `SELECT * FROM
|
| 241 |
* **Customer Product Overview:** Listing all products associated with a specific customer.
|
| 242 |
-
* Example SQL Snippet: `SELECT * FROM
|
| 243 |
* **Aggregating product counts:** Calculating the number of products per product offering.
|
| 244 |
-
* Example SQL Snippet: `SELECT product_offering_name, count(*) FROM
|
| 245 |
|
| 246 |
### Data Constraints and Business Rules:
|
| 247 |
|
|
@@ -270,7 +270,7 @@
|
|
| 270 |
|
| 271 |
## Table: `dim_product_order_item`
|
| 272 |
|
| 273 |
-
### Schema: `
|
| 274 |
|
| 275 |
### All Columns and their Descriptions:
|
| 276 |
|
|
|
|
| 1 |
## Table: `dim_agreement`
|
| 2 |
|
| 3 |
+
### Schema: `public`
|
| 4 |
|
| 5 |
### All Columns and their Descriptions:
|
| 6 |
- `agreement_id` (text): Unique identifier for the agreement. This is likely the primary key.
|
|
|
|
| 58 |
|
| 59 |
## Table: `dim_customer`
|
| 60 |
|
| 61 |
+
### Schema: `public`
|
| 62 |
|
| 63 |
### All Columns and their Descriptions:
|
| 64 |
- `customer_id` (text): Likely the primary identifier for a customer. This is probably the primary key, although not explicitly stated.
|
|
|
|
| 159 |
|
| 160 |
## Table: `dim_product`
|
| 161 |
|
| 162 |
+
### Schema: `public`
|
| 163 |
|
| 164 |
### All Columns and their Descriptions:
|
| 165 |
|
|
|
|
| 237 |
### Common Use Cases/Example Queries:
|
| 238 |
|
| 239 |
* **Tracking Product Details:** Retrieving detailed information about a specific product.
|
| 240 |
+
* Example SQL Snippet: `SELECT * FROM dim_product WHERE product_id = 'your_product_id';`
|
| 241 |
* **Customer Product Overview:** Listing all products associated with a specific customer.
|
| 242 |
+
* Example SQL Snippet: `SELECT * FROM dim_product WHERE customer_id = 'your_customer_id';`
|
| 243 |
* **Aggregating product counts:** Calculating the number of products per product offering.
|
| 244 |
+
* Example SQL Snippet: `SELECT product_offering_name, count(*) FROM dim_product GROUP BY product_offering_name;`
|
| 245 |
|
| 246 |
### Data Constraints and Business Rules:
|
| 247 |
|
|
|
|
| 270 |
|
| 271 |
## Table: `dim_product_order_item`
|
| 272 |
|
| 273 |
+
### Schema: `public`
|
| 274 |
|
| 275 |
### All Columns and their Descriptions:
|
| 276 |
|