Update app.py

Added the gradio MCP server app.

app.py

@@ -1,7 +1,658 @@
 import gradio as gr
+import psycopg2
+from psycopg2.extras import RealDictCursor
+import os
+from dotenv import load_dotenv
+from typing import Dict, Any, List, Optional
 
-def greet(name):
-    return "Hello " + name + "!!"
+load_dotenv()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def get_db_connection():
+    """Get database connection using environment variables."""
+    try:
+        conn = psycopg2.connect(
+            host=os.getenv("POSTGRES_HOST", "localhost"),
+            database=os.getenv("POSTGRES_DB", "erp_db"),
+            user=os.getenv("POSTGRES_USER", "postgres"),
+            password=os.getenv("POSTGRES_PASSWORD", ""),
+            port=os.getenv("POSTGRES_PORT", "5432")
+        )
+        return conn
+    except Exception as e:
+        raise Exception(f"Database connection failed: {str(e)}")
+
+def execute_query(query: str, params: str = None):
+    """
+    Execute a custom SQL query on the ERP database.
+    
+    Args:
+        query: SQL query to execute
+        params: Parameters for parameterized queries (as a comma-separated string)
+    
+    Returns:
+        Query results or error message
+    """
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        
+        # Convert string params to list if provided
+        param_list = None
+        if params and params.strip():
+            param_list = [p.strip() for p in params.split(',')]
+        
+        if param_list:
+            cursor.execute(query, param_list)
+        else:
+            cursor.execute(query)
+        
+        # Check if it's a SELECT query
+        if query.strip().upper().startswith('SELECT'):
+            results = cursor.fetchall()
+            # Convert RealDictRow to regular dict for JSON serialization
+            results = [dict(row) for row in results]
+            if not results:
+                return str({
+                    "success": True,
+                    "message": "No results found",
+                    "results": []
+                })
+            return str({
+                "success": True,
+                "message": f"Query returned {len(results)} results",
+                "count": len(results),
+                "results": results
+            })
+        else:
+            # For INSERT, UPDATE, DELETE queries
+            conn.commit()
+            return str({
+                "success": True,
+                "message": f"Query executed successfully",
+                "rows_affected": cursor.rowcount
+            })
+            
+    except Exception as e:
+        return str({
+            "success": False,
+            "error": str(e)
+        })
+    finally:
+        if 'conn' in locals():
+            conn.close()
+
+def list_erp_tables():
+    """
+    List all ERP tables in the database.
+    
+    Returns:
+        List of ERP table names
+    """
+    query = """
+    SELECT table_name 
+    FROM information_schema.tables 
+    WHERE table_schema = 'public'
+    AND table_name LIKE 'erp_%'
+    ORDER BY table_name;
+    """
+    
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        cursor.execute(query)
+        results = cursor.fetchall()
+        
+        table_names = [row['table_name'] for row in results]
+        if not table_names:
+            return str({
+                "success": False,
+                "message": "No ERP tables found",
+                "tables": []
+            })
+        return str({
+            "success": True,
+            "message": f"Found {len(table_names)} tables",
+            "count": len(table_names),
+            "tables": table_names
+        })
+        
+    except Exception as e:
+        return str({
+            "success": False,
+            "error": str(e)
+        })
+    finally:
+        if 'conn' in locals():
+            conn.close()
+
+def get_order_status(order_id: int):
+    """
+    Get the status of an order by order ID.
+    
+    Args:
+        order_id: The ID of the order to check
+    
+    Returns:
+        Order status information including shipping and destination countries
+    """
+    query = """
+    SELECT o.*, c.name as customer_name, c.email as customer_email,
+           o.shipping_country, o.destination_country
+    FROM erp_orders o
+    JOIN erp_customers c ON o.customer_id = c.customer_id
+    WHERE o.order_id = %s;
+    """
+    
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        cursor.execute(query, [order_id])
+        order = cursor.fetchone()
+        
+        if not order:
+            return str({
+                "success": False,
+                "message": f"Order with ID {order_id} not found"
+            })
+        
+        # Get order items
+        items_query = """
+        SELECT oi.order_item_id, oi.order_id, oi.product_id, oi.quantity, oi.unit_price, oi.subtotal,
+               p.product_name, p.sku
+        FROM erp_order_items oi
+        JOIN erp_products p ON oi.product_id = p.product_id
+        WHERE oi.order_id = %s;
+        """
+        cursor.execute(items_query, [order_id])
+        items = cursor.fetchall()
+        
+        # Get order history
+        history_query = """
+        SELECT * FROM erp_order_history
+        WHERE order_id = %s
+        ORDER BY timestamp DESC;
+        """
+        cursor.execute(history_query, [order_id])
+        history = cursor.fetchall()
+        
+        order_dict = dict(order)
+        items_list = [dict(item) for item in items]
+        history_list = [dict(entry) for entry in history]
+        
+        # Return JSON formatted response
+        return str({
+            "success": True,
+            "order": {
+                "order_id": order_id,
+                "customer": {
+                    "customer_id": order_dict['customer_id'],
+                    "name": order_dict['customer_name'],
+                    "email": order_dict['customer_email']
+                },
+                "status": order_dict['status'],
+                "shipping_address": order_dict.get('shipping_address', 'N/A'),
+                "shipping_country": order_dict.get('shipping_country', 'N/A'),
+                "destination_country": order_dict.get('destination_country', 'N/A'),
+                "items": items_list,
+                "history": history_list
+            }
+        })
+        
+    except Exception as e:
+        return str({
+            "success": False,
+            "error": str(e)
+        })
+    finally:
+        if 'conn' in locals():
+            conn.close()
+
+def place_new_order(
+    customer_id: int,
+    product_ids: str,
+    quantities: str,
+    shipping_address: str,
+    shipping_country: str,
+    destination_country: str,
+    previous_order_id: int = None
+):
+    """
+    Place a new order in the ERP system.
+    
+    Args:
+        customer_id: ID of the customer placing the order
+        product_ids: Comma-separated list of product IDs
+        quantities: Comma-separated list of quantities corresponding to product IDs
+        shipping_address: Shipping address for the order
+        shipping_country: Country where the order will be shipped from
+        destination_country: Country where the order will be delivered to
+        previous_order_id: ID of a previous order this is replacing (optional)
+    
+    Returns:
+        New order information
+    """
+    try:
+        # Parse product IDs and quantities
+        product_id_list = [int(pid.strip()) for pid in product_ids.split(',')]
+        quantity_list = [int(qty.strip()) for qty in quantities.split(',')]
+        
+        if len(product_id_list) != len(quantity_list):
+            return str({
+                "success": False,
+                "error": "Number of product IDs must match number of quantities"
+            })
+        
+        # Create items list
+        items = [
+            {"product_id": pid, "quantity": qty} 
+            for pid, qty in zip(product_id_list, quantity_list)
+        ]
+        
+        conn = get_db_connection()
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        
+        # Calculate total amount based on the total items and their prices
+        total_amount = 0
+        for item in items:
+            product_query = "SELECT price FROM erp_products WHERE product_id = %s;"
+            cursor.execute(product_query, [item['product_id']])
+            product = cursor.fetchone()
+            if not product:
+                return str({
+                    "success": False,
+                    "error": f"Product with ID {item['product_id']} not found"
+                })
+            total_amount += product['price'] * item['quantity']
+        
+        # Insert new order into the erp_orders table for the customer
+        order_query = """
+        INSERT INTO erp_orders (
+            customer_id, order_date, total_amount, status, 
+            shipping_address, shipping_country, destination_country, previous_order_id,
+            estimated_delivery, payment_status
+        ) VALUES (
+            %s, CURRENT_DATE, %s, 'Processing', 
+            %s, %s, %s, %s,
+            CURRENT_DATE + INTERVAL '7 days', 'Pending'
+        ) RETURNING order_id;
+        """
+        cursor.execute(order_query, [
+            customer_id, total_amount, shipping_address, shipping_country, destination_country, previous_order_id
+        ])
+        new_order_id = cursor.fetchone()['order_id']
+        
+        # Insert order items into erp_order_items table for the new order
+        for item in items:
+            # Get product price
+            cursor.execute(product_query, [item['product_id']])
+            product = cursor.fetchone()
+            unit_price = product['price']
+            subtotal = unit_price * item['quantity']
+            
+            item_query = """
+            INSERT INTO erp_order_items (
+                order_id, product_id, quantity, unit_price, subtotal
+            ) VALUES (
+                %s, %s, %s, %s, %s
+            );
+            """
+            cursor.execute(item_query, [
+                new_order_id, item['product_id'], item['quantity'], 
+                unit_price, subtotal
+            ])
+            
+            # Update product stock
+            update_stock_query = """
+            UPDATE erp_products 
+            SET stock_quantity = stock_quantity - %s 
+            WHERE product_id = %s;
+            """
+            cursor.execute(update_stock_query, [item['quantity'], item['product_id']])
+        
+        # Create order history entry into the erp_order_history for the new order
+        history_query = """
+        INSERT INTO erp_order_history (
+            order_id, timestamp, status_change, notes, updated_by
+        ) VALUES (
+            %s, CURRENT_TIMESTAMP, 'Order Created', 'New order placed', 'System'
+        );
+        """
+        cursor.execute(history_query, [new_order_id])
+        
+        # If this is a replacement order, add a note to the previous order
+        if previous_order_id:
+            prev_order_note_query = """
+            INSERT INTO erp_order_history (
+                order_id, timestamp, status_change, notes, updated_by
+            ) VALUES (
+                %s, CURRENT_TIMESTAMP, 'Replaced', 'Order replaced by order #%s', 'System'
+            );
+            """
+            cursor.execute(prev_order_note_query, [previous_order_id, new_order_id])
+        
+        # Generate invoice for the new order
+        invoice_query = """
+        INSERT INTO erp_invoices (
+            order_id, invoice_date, amount, payment_terms, due_date, is_paid, invoice_number
+        ) VALUES (
+            %s, CURRENT_DATE, %s, 'Net 30', CURRENT_DATE + INTERVAL '30 days', FALSE, 'INV-' || %s
+        ) RETURNING invoice_id;
+        """
+        cursor.execute(invoice_query, [new_order_id, total_amount, new_order_id])
+        invoice_id = cursor.fetchone()['invoice_id']
+        
+        conn.commit()
+        
+        # Get the complete new order details
+        cursor.execute("SELECT * FROM erp_orders WHERE order_id = %s;", [new_order_id])
+        order = cursor.fetchone()
+        
+        order_dict = dict(order)
+        
+        # Return JSON formatted response
+        return str({
+            "success": True,
+            "order": {
+                "order_id": new_order_id,
+                "invoice_id": invoice_id,
+                "total_amount": float(total_amount),
+                "status": order_dict['status'],
+                "estimated_delivery": str(order_dict['estimated_delivery']) if order_dict['estimated_delivery'] else None,
+                "customer_id": customer_id,
+                "shipping_address": shipping_address,
+                "shipping_country": shipping_country,
+                "destination_country": destination_country,
+                "previous_order_id": previous_order_id,
+                "items": [{"product_id": pid, "quantity": qty} for pid, qty in zip(product_id_list, quantity_list)]
+            }
+        })
+        
+    except Exception as e:
+        if 'conn' in locals():
+            conn.rollback()
+        return str({
+            "success": False,
+            "error": str(e)
+        })
+    finally:
+        if 'conn' in locals():
+            conn.close()
+
+def cancel_order(order_id: int, reason: str):
+    """
+    Cancel an existing order in the ERP system.
+    
+    Args:
+        order_id: ID of the order to cancel
+        reason: Reason for cancellation
+    
+    Returns:
+        Result of the cancellation operation
+    """
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        
+        # Check if order exists and can be cancelled
+        check_query = """
+        SELECT status, customer_id FROM erp_orders WHERE order_id = %s;
+        """
+        cursor.execute(check_query, [order_id])
+        order = cursor.fetchone()
+        
+        if not order:
+            return str({
+                "success": False,
+                "error": f"Order with ID {order_id} not found"
+            })
+        
+        if order['status'] in ['Delivered', 'Cancelled']:
+            return str({
+                "success": False,
+                "error": f"Cannot cancel order with status '{order['status']}'"
+            })
+        
+        # Get order items to restore stock
+        items_query = """
+        SELECT product_id, quantity FROM erp_order_items WHERE order_id = %s;
+        """
+        cursor.execute(items_query, [order_id])
+        items = cursor.fetchall()
+        
+        # Update order status to Cancelled
+        update_query = """
+        UPDATE erp_orders SET status = 'Cancelled', payment_status = 'Cancelled'
+        WHERE order_id = %s;
+        """
+        cursor.execute(update_query, [order_id])
+        
+        # Add entry to order history
+        history_query = """
+        INSERT INTO erp_order_history (
+            order_id, timestamp, status_change, notes, updated_by
+        ) VALUES (
+            %s, CURRENT_TIMESTAMP, 'Cancelled', %s, 'System'
+        );
+        """
+        cursor.execute(history_query, [order_id, f"Order cancelled: {reason}"])
+        
+        # Restore product stock quantities
+        for item in items:
+            restore_stock_query = """
+            UPDATE erp_products 
+            SET stock_quantity = stock_quantity + %s 
+            WHERE product_id = %s;
+            """
+            cursor.execute(restore_stock_query, [item['quantity'], item['product_id']])
+        
+        # Update invoice if exists
+        invoice_query = """
+        UPDATE erp_invoices 
+        SET is_paid = FALSE
+        WHERE order_id = %s;
+        """
+        cursor.execute(invoice_query, [order_id])
+        
+        conn.commit()
+        
+        return str({
+            "success": True,
+            "message": f"Order #{order_id} has been successfully cancelled",
+            "order_id": order_id,
+            "reason": reason,
+            "items_restored": len(items)
+        })
+        
+    except Exception as e:
+        if 'conn' in locals():
+            conn.rollback()
+        return str({
+            "success": False,
+            "error": str(e)
+        })
+    finally:
+        if 'conn' in locals():
+            conn.close()
+
+def get_invoice_details(invoice_id: int = None, order_id: int = None):
+    """
+    Get invoice details by invoice ID or order ID.
+    
+    Args:
+        invoice_id: ID of the invoice (optional)
+        order_id: ID of the order (optional)
+    
+    Returns:
+        Invoice details including customer and order information
+    """
+    if not invoice_id and not order_id:
+        return str({
+            "success": False,
+            "error": "Either invoice_id or order_id must be provided"
+        })
+    
+    try:
+        conn = get_db_connection()
+        cursor = conn.cursor(cursor_factory=RealDictCursor)
+        
+        if invoice_id:
+            query = """
+            SELECT i.*, o.order_date, o.status as order_status, 
+                   c.name as customer_name, c.email as customer_email, c.address as customer_address
+            FROM erp_invoices i
+            JOIN erp_orders o ON i.order_id = o.order_id
+            JOIN erp_customers c ON o.customer_id = c.customer_id
+            WHERE i.invoice_id = %s;
+            """
+            cursor.execute(query, [invoice_id])
+        else:
+            query = """
+            SELECT i.*, o.order_date, o.status as order_status, 
+                   c.name as customer_name, c.email as customer_email, c.address as customer_address
+            FROM erp_invoices i
+            JOIN erp_orders o ON i.order_id = o.order_id
+            JOIN erp_customers c ON o.customer_id = c.customer_id
+            WHERE i.order_id = %s;
+            """
+            cursor.execute(query, [order_id])
+        
+        invoice = cursor.fetchone()
+        
+        if not invoice:
+            return str({
+                "success": False,
+                "error": f"Invoice not found for the provided {'invoice_id' if invoice_id else 'order_id'}"
+            })
+        
+        # Get order items
+        items_query = """
+        SELECT oi.*, p.product_name, p.sku
+        FROM erp_order_items oi
+        JOIN erp_products p ON oi.product_id = p.product_id
+        WHERE oi.order_id = %s;
+        """
+        cursor.execute(items_query, [invoice['order_id']])
+        items = cursor.fetchall()
+        
+        invoice_dict = dict(invoice)
+        items_list = [dict(item) for item in items]
+        
+        # Return JSON formatted response
+        return str({
+            "success": True,
+            "invoice": {
+                "invoice_id": invoice_dict['invoice_id'],
+                "invoice_number": invoice_dict.get('invoice_number', ''),
+                "order_id": invoice_dict['order_id'],
+                "order_date": str(invoice_dict['order_date']) if invoice_dict['order_date'] else None,
+                "order_status": invoice_dict['order_status'],
+                "amount": float(invoice_dict['amount']),
+                "due_date": str(invoice_dict['due_date']) if invoice_dict['due_date'] else None,
+                "payment_status": "Paid" if invoice_dict['is_paid'] else "Unpaid",
+                "customer": {
+                    "name": invoice_dict['customer_name'],
+                    "email": invoice_dict['customer_email'],
+                    "address": invoice_dict['customer_address']
+                },
+                "items": items_list
+            }
+        })
+        
+    except Exception as e:
+        return str({
+            "success": False,
+            "error": str(e)
+        })
+    finally:
+        if 'conn' in locals():
+            conn.close()
+
+# Create Gradio interfaces for each function
+execute_query_interface = gr.Interface(
+    fn=execute_query,
+    inputs=[
+        gr.Textbox(lines=5, label="SQL Query"),
+        gr.Textbox(label="Parameters (comma-separated)", placeholder="Optional")
+    ],
+    outputs=gr.Textbox(lines=10),
+    title="Execute SQL Query",
+    description="Execute a custom SQL query on the ERP database"
+)
+
+list_tables_interface = gr.Interface(
+    fn=list_erp_tables,
+    inputs=[],
+    outputs=gr.Textbox(lines=10),
+    title="List ERP Tables",
+    description="List all ERP tables in the database"
+)
+
+order_status_interface = gr.Interface(
+    fn=get_order_status,
+    inputs=gr.Number(label="Order ID", precision=0),
+    outputs=gr.Textbox(lines=15),
+    title="Get Order Status",
+    description="Get the status of an order by order ID"
+)
+
+place_order_interface = gr.Interface(
+    fn=place_new_order,
+    inputs=[
+        gr.Number(label="Customer ID", precision=0),
+        gr.Textbox(label="Product IDs (comma-separated)", placeholder="1, 2, 3"),
+        gr.Textbox(label="Quantities (comma-separated)", placeholder="2, 1, 3"),
+        gr.Textbox(label="Shipping Address"),
+        gr.Textbox(label="Shipping Country"),
+        gr.Textbox(label="Destination Country"),
+        gr.Number(label="Previous Order ID (optional)", precision=0)
+    ],
+    outputs=gr.Textbox(lines=10),
+    title="Place New Order",
+    description="Place a new order in the ERP system"
+)
+
+cancel_order_interface = gr.Interface(
+    fn=cancel_order,
+    inputs=[
+        gr.Number(label="Order ID", precision=0),
+        gr.Textbox(label="Reason for Cancellation")
+    ],
+    outputs=gr.Textbox(lines=5),
+    title="Cancel Order",
+    description="Cancel an existing order in the ERP system"
+)
+
+invoice_details_interface = gr.Interface(
+    fn=get_invoice_details,
+    inputs=[
+        gr.Number(label="Invoice ID (optional)", precision=0),
+        gr.Number(label="Order ID (optional)", precision=0)
+    ],
+    outputs=gr.Textbox(lines=15),
+    title="Get Invoice Details",
+    description="Get invoice details by invoice ID or order ID"
+)
+
+# Create a Gradio TabItem for each interface
+demo = gr.TabbedInterface(
+    [
+        execute_query_interface,
+        list_tables_interface,
+        order_status_interface,
+        place_order_interface,
+        cancel_order_interface,
+        invoice_details_interface
+    ],
+    [
+        "Execute Query",
+        "List Tables",
+        "Order Status",
+        "Place Order",
+        "Cancel Order",
+        "Invoice Details"
+    ],
+    title="ERP System Tools"
+)
+
+# Launch the demo with MCP server enabled
+demo.launch(mcp_server=True, server_port=7090)