feat: demo with yugabyte

.env.sample

@@ -4,6 +4,9 @@ TABLE_SUMMARY_PATH=
 DB_URL=
 DB_SCHEMA=
 
+# DB URL with root.crt for yugabyte
+# DB_URL="postgresql://username:password@us-east-1.14747abb-66a1-48eb-a85f-d434c56d612b.aws.yugabyte.cloud:5433/yugabyte?sslmode=verify-full&sslrootcert=root.crt"
+
 GEMINI_API_KEY=
 GEMINI_MODEL=
 GEMINI_MODEL_PROVIDER=

gradio_app.py

@@ -11,6 +11,7 @@ from memory_store import MemoryStore
 import logging
 
 
+
 # ======================================= Load DB configs
 def load_db_configs():
     """Load database configurations from configs.yaml"""
@@ -121,41 +122,56 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                     scale=4
                 ),
                 theme="soft",
+                # examples=[
+                #     "Describe the database",
+                #     "List all tables in the database",
+                #     "List all tables with columns and data types",
+                #     "How many customers do you have?",
+                #     "What are the statuses my of my customers",
+                #     "Visualize with different colors and show legend",
+                #     "What are the statues of my customers and how many are in each status, show it by percentage",
+                #     "Total number of completed orders in six years by customer count show top most 10 customers",
+                #     "In january how many products has been sold ? group them by year",
+                #     "How many users and roles have been created in 2024"
+                # ],
                 examples=[
                     "Describe the database",
                     "List all tables in the database",
                     "List all tables with columns and data types",
-                    "How many customers do you have?",
-                    "What are the statuses my of my customers",
+                    "How many comments are there per ticket channel (email, chat, portal)? Also Visualize it as a pie chart",
                     "Visualize with different colors and show legend",
-                    "What are the statues of my customers and how many are in each status, show it by percentage",
-                    "Total number of completed orders in six years by customer count show top most 10 customers",
-                    "In january how many products has been sold ? group them by year",
-                    "How many users and roles have been created in 2024"
+                    "How many customers are in each industry?",
+                    "List the 5 most active agents by ticket count in 2024.",
+                    "How many tickets were reopened at least once?"
                 ],
                 save_history=True,
                 type="messages"
             )
         with gr.Column(scale=1):
             with gr.Accordion("Example Questions", open=True):
+                # gr.Markdown("""
+                # - 📊 List all tables in database
+                # - 👥 Total number of customers
+                # - 📈 Visualize it with different colors
+                # - 📋 Order statistics for last 6 years
+                # - 📆 User and role counts in 2024
+                # """)
                 gr.Markdown("""
                 - 📊 List all tables in database
                 - 👥 Total number of customers
                 - 📈 Visualize it with different colors
                 - 📋 Order statistics for last 6 years
-                - 📆 User and role counts in 2024
+                - 📆 Average ticket reopen count per year
                 """)
 
 
-
-
 # TODO: maybe we can add a mcp tool to validate the results (those converted to DataFrame) to make sure the valid type is passed to the visualization tool by ReAct agent
 
 
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
-        server_port=7860,
+        # server_port=7860,
         share=True,
         debug=True
     )

langchain_mcp_client.py

@@ -15,7 +15,6 @@ from langchain_community.chat_message_histories import ChatMessageHistory
 from memory_store import MemoryStore
 
 
-
 # set_debug(True)
 
 
@@ -103,6 +102,7 @@ def load_table_summary(path: str) -> str:
     with open(path, 'r') as file:
         return file.read()
 
+
 def get_server_params() -> StdioServerParameters:
     return StdioServerParameters(
         command="python",

postgre_mcp_server.py

@@ -15,7 +15,10 @@ import logging
 DEFAULT_QUERY_LIMIT = 100
 
 # logging info
-logging.basicConfig(level=logging.INFO)
+# logging.basicConfig(level=logging.INFO)
+
+# get logger
+logger = logging.getLogger(__name__)
 
 # Define our own PromptMessage class if the MCP one isn't available
 @dataclass
@@ -148,10 +151,10 @@ You can use the following FastMCP tools to create **read-only** queries (e.g., `
 # Output Format
 ==========================
 
-Present your final answer using the following structure **exactly** in markdown language. When necessary, bold the important parts of your answer or use `` for inline code blocks:
+Present your final answer using the following structure in markdown language:
 
 # 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.}}
+{{Take the result from the execute_query tool and format it nicely using Markdown. Use a beautiful Markdown table for tabular data (rows and columns) including headers and show such simple results using a table. 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 (and visualization, if applicable) in 1-3 sentences. Explain what the data and visualization (if any) represent in the context of the user's request.}}

table_summary_yugabyte.txt

@@ -0,0 +1,153 @@
+## Database: helpdesk
+### Description:
+This database is designed for managing customer support tickets and related comments in a helpdesk system. The database consists of three main tables containing customer information, support tickets, and ticket comments.
+ 
+### All Tables:
+- customer
+- ticket
+- ticket_comment
+ 
+### Default Schema:
+- public
+ 
+---
+ 
+## Table: `customer`
+ 
+### Schema: `public`
+ 
+### All Columns and their Descriptions:
+- `customer_id` (integer): Unique identifier for the customer. Primary key.
+- `name` (text): Name of the customer company.
+- `industry` (text): Industry sector of the customer (Finance, Telecommunications, IT Services, Retail).
+- `contact_email` (text): Contact email address.
+- `contact_phone` (text): Contact phone number.
+- `region` (text): Customer's region (Dammam, Riyadh, Jeddah, Mecca).
+- `created_at` (timestamp): Date when the customer record was created.
+- `status` (text): Current status of the customer (Active, Inactive, Prospect, Suspended, Blacklisted, Pending Verification, Closed).
+ 
+### Relationships with Other Tables (Foreign Keys):
+- `customer_id` is referenced in the `ticket` table.
+ 
+### Cardinality of Relationships:
+- One customer can have multiple support tickets (one-to-many).
+ 
+### Common Use Cases/Example Queries:
+- List active customers in a specific region
+- Analyze customer distribution by industry
+- Generate statistics based on customer status
+- Example SQL Snippet: `SELECT * FROM customer WHERE status = 'Active' AND region = 'Riyadh';`
+ 
+### Data Constraints and Business Rules:
+- `customer_id` must be unique and is the primary key
+- `status` values are restricted to a controlled vocabulary
+- `created_at` must be a valid date
+ 
+### Data Update Frequency/Volatility:
+- Customer information typically changes with new customer registrations or updates to existing customer information
+- Status updates occur regularly
+ 
+### Important Notes/Considerations for Querying:
+- Be careful when filtering by status
+- Pay attention to correct region names in region-based analyses
+ 
+---
+ 
+## Table: `ticket`
+ 
+### Schema: `public`
+ 
+### All Columns and their Descriptions:
+- `ticket_id` (text): Unique identifier for the support ticket. Primary key.
+- `status` (text): Ticket status (Open, Resolved-Cancelled, Resolved-Completed).
+- `sub_status` (text): Ticket sub-status (Pending for approval, Cancelled, Completed).
+- `created_at` (timestamp): Date and time when the ticket was created.
+- `service_number` (text): Service number.
+- `customer_id` (integer): Customer ID. Foreign key to customer table.
+- `sr_classification` (text): Service request classification (Financial).
+- `type` (text): Ticket type (Complaint, Request, Inquiry).
+- `area` (text): Main area (Billing).
+- `sub_area` (text): Sub-area (Auto Adjustments, Mass Adjustment, Bulk Balance Transfer, Refunds).
+- `owner` (text): Ticket owner (Agent A, Agent B, Agent C, etc.).
+- `ticket_group` (text): Ticket group (Technical Support, Billing Team, Customer Service, NOC, Field Operations).
+- `description` (text): Ticket description.
+- `priority` (text): Priority level (Low, Medium, High).
+- `re_open_count` (integer): Number of times the ticket was reopened.
+- `follow_up_count` (integer): Number of follow-ups.
+- `collaboration_task_count` (integer): Number of collaboration tasks.
+- `customer_level_value` (text): Customer level (Blue, Platinum).
+- `customer_segment` (text): Customer segment (CS, HP, G2, G4).
+- `cst_area_code` (text): Customer area code.
+- `cst_sub_are_code` (text): Customer sub-area code.
+- `circuit_name` (text): Circuit name.
+- `collaboration_task` (text): Collaboration task.
+- `repeated_ticket_based_on_area_sub_area` (integer): Number of repeated tickets based on area/sub-area.
+- `service_based_repeated_ticket` (integer): Number of service-based repeated tickets.
+- `contact_number` (text): Contact number.
+- `product` (text): Product name (Fiber Enterprise, IP VPN, Hosted PBX, etc.).
+ 
+### Relationships with Other Tables (Foreign Keys):
+- `customer_id` (FK) references `customer` table
+- `ticket_id` is referenced in the `ticket_comment` table
+ 
+### Cardinality of Relationships:
+- One customer can have multiple tickets (many-to-one with customer)
+- One ticket can have multiple comments (one-to-many with ticket_comment)
+ 
+### Common Use Cases/Example Queries:
+- List of open tickets
+- Ticket distribution by priority
+- Product-based ticket analysis
+- Example SQL Snippet: `SELECT * FROM ticket WHERE status = 'Open' AND priority = 'High';`
+ 
+### Data Constraints and Business Rules:
+- `ticket_id` must be unique and is the primary key
+- `status` and `sub_status` are restricted to specific values
+- `priority` can only be Low, Medium, or High
+ 
+### Data Update Frequency/Volatility:
+- Tickets are continuously created and updated
+- Status and sub_status are frequently updated
+ 
+### Important Notes/Considerations for Querying:
+- Status and sub_status should be evaluated together
+- Consider timezone when querying date-based data
+ 
+---
+ 
+## Table: `ticket_comment`
+ 
+### Schema: `public`
+ 
+### All Columns and their Descriptions:
+- `comment_id` (integer): Unique identifier for the comment. Primary key.
+- `ticket_id` (text): Related ticket ID. Foreign key to ticket table.
+- `created_at` (timestamp): Date and time when the comment was created.
+- `user_type` (text): Type of user making the comment (customer, agent).
+- `channel` (text): Comment channel (email, chat, portal).
+- `comment_text` (text): Comment text.
+ 
+### Relationships with Other Tables (Foreign Keys):
+- `ticket_id` (FK) references `ticket` table
+ 
+### Cardinality of Relationships:
+- One ticket can have multiple comments (many-to-one with ticket)
+ 
+### Common Use Cases/Example Queries:
+- View all comments for a specific ticket
+- Analysis of comments by channel
+- Distribution of comments by user type
+- Example SQL Snippet: `SELECT * FROM ticket_comment WHERE ticket_id = 'W-137014' ORDER BY created_at DESC;`
+ 
+### Data Constraints and Business Rules:
+- `comment_id` must be unique and is the primary key
+- `user_type` can only be 'customer' or 'agent'
+- `channel` can only be 'email', 'chat', or 'portal'
+ 
+### Data Update Frequency/Volatility:
+- Comments are continuously added
+- Existing comments are rarely updated or deleted
+ 
+### Important Notes/Considerations for Querying:
+- Important to order comments by date
+- Check consistency in ticket-comment relationships