From f9fe348db062b2cd5865171db4335c70e5d7fa0e Mon Sep 17 00:00:00 2001 From: Chih-Yu Yeh Date: Wed, 29 May 2024 14:29:16 +0800 Subject: [PATCH] refine ai pipelines (#310) * update followup prompt * use deployed verion's mdl * update demo data version * upload original modified mdl --- .../sample_dataset/ecommerce2_duckdb_mdl.json | 715 ++++++++++++++++++ .../sample_dataset/ecommerce_duckdb_mdl.json | 412 +++------- wren-ai-service/demo/utils.py | 44 +- .../ask/followup_generation_pipeline.py | 82 +- 4 files changed, 931 insertions(+), 322 deletions(-) create mode 100644 wren-ai-service/demo/sample_dataset/ecommerce2_duckdb_mdl.json diff --git a/wren-ai-service/demo/sample_dataset/ecommerce2_duckdb_mdl.json b/wren-ai-service/demo/sample_dataset/ecommerce2_duckdb_mdl.json new file mode 100644 index 000000000..41cf4a4cf --- /dev/null +++ b/wren-ai-service/demo/sample_dataset/ecommerce2_duckdb_mdl.json @@ -0,0 +1,715 @@ +{ + "catalog": "wrenai", + "schema": "ecommerce", + "models": [ + { + "name": "customers", + "refSql": "select * from main.customers", + "columns": [ + { + "name": "City", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "The Customer City, where the customer company is located. Also called 'customer segment'." + } + }, + { + "name": "Id", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A unique identifier for each customer in the data model." + } + }, + { + "name": "State", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A field indicating the state where the customer is located." + } + }, + { + "name": "orders", + "type": "orders", + "relationship": "CustomersOrders", + "isCalculated": false, + "notNull": false, + "properties": {} + }, + { + "name": "LatestRecord", + "type": "DATE", + "isCalculated": true, + "expression": "max(orders.PurchaseTimestamp)", + "notNull": false, + "properties": {} + }, + { + "name": "FirstRecord", + "type": "DATE", + "isCalculated": true, + "expression": "min(orders.PurchaseTimestamp)", + "notNull": false, + "properties": {} + }, + { + "name": "VIP", + "type": "BOOLEAN", + "isCalculated": true, + "expression": "sum(orders.Size) > 2", + "notNull": false, + "properties": {} + }, + { + "name": "OrderCount", + "type": "BIGINT", + "isCalculated": true, + "expression": "count(orders.OrderId)", + "notNull": false, + "properties": {} + }, + { + "name": "Debit", + "type": "DOUBLE", + "isCalculated": true, + "expression": "sum(orders.OrderBalance)", + "notNull": false, + "properties": {} + }, + { + "name": "ReviewRate", + "type": "DOUBLE", + "isCalculated": true, + "expression": "count(orders.IsReviewed = TRUE) / count(DISTINCT orders.OrderId)", + "notNull": false, + "properties": {} + } + ], + "primaryKey": "Id", + "cached": false, + "refreshTime": "30.00m", + "properties": { + "schema": "main", + "catalog": "memory", + "description": "A table of customers who have made purchases, including their city" + } + }, + { + "name": "order_items", + "refSql": "select * from main.order_items", + "columns": [ + { + "name": "FreightValue", + "type": "DOUBLE", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A numerical value representing the cost of shipping for an item in an order." + } + }, + { + "name": "ItemNumber", + "type": "BIGINT", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "The sequential number of the order item in this order. Each order item in an order has its unique ItemNumber." + } + }, + { + "name": "OrderId", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A VARCHAR value indicating the order that this order_item belongs to. The column is used to map the order_item to Orders model in the OrdersOrder_items relationship." + } + }, + { + "name": "Price", + "type": "DOUBLE", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A numerical value representing the price of an item in an order." + } + }, + { + "name": "ProductId", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A VARCHAR value representing the product of this order_item. The column is used to map the order_item to Products model using ProductsOrder_items relationship." + } + }, + { + "name": "ShippingLimitDate", + "type": "DATE", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A date value indicating the limit by which an item should be shipped according to the order. It helps track the deadline for shipping items in the 'order_items' model." + } + }, + { + "name": "orders", + "type": "orders", + "relationship": "OrdersOrder_items", + "isCalculated": false, + "notNull": false, + "properties": {} + }, + { + "name": "products", + "type": "products", + "relationship": "ProductsOrder_items", + "isCalculated": false, + "notNull": false, + "properties": {} + } + ], + "primaryKey": "", + "cached": false, + "refreshTime": "30.00m", + "properties": { + "schema": "main", + "catalog": "memory", + "description": "The model is used to store information about items in orders, including details like prices, product IDs, shipping limits, and relationships with orders and products tables." + } + }, + { + "name": "orders", + "refSql": "select * from main.orders", + "columns": [ + { + "name": "ApprovedTimestamp", + "type": "TIMESTAMP", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column that represents the timestamp when the order was approved." + } + }, + { + "name": "CustomerId", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A unique identifier representing the customer who purchased this order." + } + }, + { + "name": "DeliveredCarrierDate", + "type": "DATE", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column that represents the date when the order was delivered by the carrier." + } + }, + { + "name": "DeliveredCustomerDate", + "type": "DATE", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column that represents the date when the order was delivered to the customer." + } + }, + { + "name": "EstimatedDeliveryDate", + "type": "DATE", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column that represents the estimated delivery date of the order." + } + }, + { + "name": "OrderId", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column that represents a unique identifier of this order." + } + }, + { + "name": "PurchaseTimestamp", + "type": "TIMESTAMP", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column that represents the timestamp when the order was purchased." + } + }, + { + "name": "Status", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column representing the status of the order." + } + }, + { + "name": "customers", + "type": "customers", + "relationship": "CustomersOrders", + "isCalculated": false, + "notNull": false, + "properties": {} + }, + { + "name": "order_items", + "type": "order_items", + "relationship": "OrdersOrder_items", + "isCalculated": false, + "notNull": false, + "properties": {} + }, + { + "name": "reviews", + "type": "reviews", + "relationship": "OrdersReviews", + "isCalculated": false, + "notNull": false, + "properties": {} + }, + { + "name": "payments", + "type": "payments", + "relationship": "PaymentsOrders", + "isCalculated": false, + "notNull": false, + "properties": {} + }, + { + "name": "Rating", + "type": "DOUBLE", + "isCalculated": true, + "expression": "avg(reviews.Score)", + "notNull": false, + "properties": { + "description": "A calculated column representing averge score of the reviews associated with the order.", + "display_name": "Rating" + } + }, + { + "name": "ReviewCount", + "type": "BIGINT", + "isCalculated": true, + "expression": "count(reviews.Id)", + "notNull": false, + "properties": { + "description": "A calculated column representing the number of the reviews associated with the order.", + "display_name": "Review Count" + } + }, + { + "name": "IsReviewed", + "type": "BOOLEAN", + "isCalculated": true, + "expression": "count(reviews.Id) > 0", + "notNull": false, + "properties": {} + }, + { + "name": "Size", + "type": "BIGINT", + "isCalculated": true, + "expression": "count(order_items.ItemNumber)", + "notNull": false, + "properties": { + "description": "A calculated column representing the number of the items associated with the order.", + "display_name": "Package Size" + } + }, + { + "name": "Large", + "type": "BOOLEAN", + "isCalculated": true, + "expression": "count(order_items.ItemNumber) > 1", + "notNull": false, + "properties": {} + }, + { + "name": "TotalValue", + "type": "DOUBLE", + "isCalculated": true, + "expression": "sum(order_items.Price)", + "notNull": false, + "properties": {} + }, + { + "name": "PaidValue", + "type": "DOUBLE", + "isCalculated": true, + "expression": "sum(payments.Value)", + "notNull": false, + "properties": {} + }, + { + "name": "OrderBalance", + "type": "DOUBLE", + "isCalculated": true, + "expression": "sum(order_items.Price) - sum(payments.Value)", + "notNull": false, + "properties": {} + } + ], + "primaryKey": "OrderId", + "cached": false, + "refreshTime": "30.00m", + "properties": { + "schema": "main", + "catalog": "memory", + "description": "A model representing the orders data." + } + }, + { + "name": "payments", + "refSql": "select * from main.payments", + "columns": [ + { + "name": "Installments", + "type": "BIGINT", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column representing the number of installments in the payments data model." + } + }, + { + "name": "OrderId", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column representing the order id associated with this payment. The column is used to map the payment to the order using PaymentsOrders relationship." + } + }, + { + "name": "Sequential", + "type": "BIGINT", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column representing the sequential number of the payment in its corresponding order. Each payment in the order has its unique sequential number." + } + }, + { + "name": "Type", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column representing the type of payment in the payments data model." + } + }, + { + "name": "Value", + "type": "DOUBLE", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A column representing the value of the payment in the payments data model." + } + }, + { + "name": "orders", + "type": "orders", + "relationship": "PaymentsOrders", + "isCalculated": false, + "notNull": false, + "properties": {} + } + ], + "primaryKey": "", + "cached": false, + "refreshTime": "30.00m", + "properties": { + "schema": "main", + "catalog": "memory", + "description": "A model representing the payment records, including installments, order IDs, sequential numbers, payment types, values, and relationships with orders." + } + }, + { + "name": "products", + "refSql": "select * from main.products", + "columns": [ + { + "name": "Category", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A category that classifies the products in the data model." + } + }, + { + "name": "Id", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A unique identifier assigned to each product in the data model." + } + }, + { + "name": "Name", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A name of the product in the data model." + } + }, + { + "name": "order_items", + "type": "order_items", + "relationship": "ProductsOrder_items", + "isCalculated": false, + "notNull": false, + "properties": {} + }, + { + "name": "Revenue", + "type": "DOUBLE", + "isCalculated": true, + "notNull": false, + "expression": "sum(order_items.Price)", + "properties": { + "description": "The calculated field representing the product revenue by summing up the product price in each order.", + "display_name": "Revenue" + } + } + ], + "primaryKey": "Id", + "cached": false, + "refreshTime": "30.00m", + "properties": { + "schema": "main", + "catalog": "memory", + "description": "A data model containing information about products such as category, ID, and name, with a relationship to order items." + } + }, + { + "name": "reviews", + "refSql": "select * from main.reviews", + "columns": [ + { + "name": "AnswerTimestamp", + "type": "TIMESTAMP", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "The date when the answer was provided." + } + }, + { + "name": "CreationTimestamp", + "type": "TIMESTAMP", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "The date when the review was created." + } + }, + { + "name": "Id", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "A unique identifier assigned to each review entry." + } + }, + { + "name": "OrderId", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "The order id of the order which the review belongs to." + } + }, + { + "name": "Score", + "type": "BIGINT", + "isCalculated": false, + "notNull": false, + "properties": { + "description": "The score associated with each review entry." + } + }, + { + "name": "orders", + "type": "orders", + "relationship": "OrdersReviews", + "isCalculated": false, + "notNull": false, + "properties": {} + } + ], + "primaryKey": "", + "cached": false, + "refreshTime": "30.00m", + "properties": { + "schema": "main", + "catalog": "memory", + "description": "A model containing information about review of orders." + } + } + ], + "relationships": [ + { + "name": "CustomersOrders", + "models": [ + "customers", + "orders" + ], + "joinType": "ONE_TO_MANY", + "condition": "customers.Id = orders.CustomerId", + "manySideSortKeys": [], + "properties": {} + }, + { + "name": "OrdersOrder_items", + "models": [ + "orders", + "order_items" + ], + "joinType": "ONE_TO_MANY", + "condition": "orders.OrderId = order_items.OrderId", + "manySideSortKeys": [], + "properties": {} + }, + { + "name": "ProductsOrder_items", + "models": [ + "products", + "order_items" + ], + "joinType": "ONE_TO_MANY", + "condition": "products.Id = order_items.ProductId", + "manySideSortKeys": [], + "properties": {} + }, + { + "name": "OrdersReviews", + "models": [ + "orders", + "reviews" + ], + "joinType": "ONE_TO_MANY", + "condition": "orders.OrderId = reviews.OrderId", + "manySideSortKeys": [], + "properties": {} + }, + { + "name": "PaymentsOrders", + "models": [ + "payments", + "orders" + ], + "joinType": "MANY_TO_ONE", + "condition": "payments.OrderId = orders.OrderId", + "manySideSortKeys": [], + "properties": {} + } + ], + "enumDefinitions": [], + "metrics": [ + { + "name": "Revenue", + "baseObject": "orders", + "dimension": [ + { + "name": "PurchaseTimestamp", + "type": "TIMESTAMP" + }, + { + "name": "CustomerId", + "type": "VARCHAR" + }, + { + "name": "Status", + "type": "VARCHAR" + }, + { + "name": "Size", + "type": "BIGINT" + }, + { + "name": "IsReviewed", + "type": "BOOLEAN" + } + ], + "measure": [ + { + "name": "Price", + "type": "DOUBLE", + "expression": "sum(order_items.Price)" + }, + { + "name": "MoM_Price", + "type": "DOUBLE", + "expression": "sum(order_items.Price) - COALESCE(LAG(sum(order_items.Price)) OVER (ORDER BY date_trunc('MONTH', PurchaseTimestamp)), 0) / COALESCE(LAG(sum(order_items.Price)) OVER (ORDER BY date_trunc('MONTH', PurchaseTimestamp)), 1)" + }, + { + "name": "NumberOfOrders", + "type": "BIGINT", + "expression": "count(OrderId)" + }, + { + "name": "Peak", + "type": "DOUBLE", + "expression": "max(TotalValue)" + } + ], + "timeGrain": [ + { + "name": "PurchaseTimestamp", + "refColumn": "PurchaseTimestamp", + "dateParts": [ + "YEAR", + "MONTH", + "DAY" + ] + } + ] + } + ], + "cumulativeMetrics": [], + "views": [ + { + "name": "total_items_view", + "statement": "WITH customer_order_items AS (SELECT c.State, oi.ItemNumber FROM customers c JOIN orders o ON c.Id = o.CustomerId JOIN order_items oi ON o.OrderId = oi.OrderId) SELECT State, COUNT(ItemNumber) AS Total_Items_Purchased FROM customer_order_items GROUP BY State", + "properties": { + "question": "How many items have been purchased in orders placed by customers in each state?", + "summary": "The query counts the total number of items purchased by customers in each state.", + "viewId": "20b69c05-e9a3-4c4d-a421-06ba7e7a385a" + } + } + ], + "macros": [ + { + "name": "MoM", + "definition": "(calculation: Expression, timestamp: Expression) => ({{ calculation }} - COALESCE(LAG({{ calculation }}) OVER (ORDER BY {{ timestamp }}), 0)) / COALESCE(LAG({{ calculation }}) OVER (ORDER BY {{ timestamp }}), 1)" + } + ], + "dateSpine": { + "unit": "DAY", + "start": "1970-01-01", + "end": "2077-12-31", + "properties": {} + } +} diff --git a/wren-ai-service/demo/sample_dataset/ecommerce_duckdb_mdl.json b/wren-ai-service/demo/sample_dataset/ecommerce_duckdb_mdl.json index 41cf4a4cf..21bda2971 100644 --- a/wren-ai-service/demo/sample_dataset/ecommerce_duckdb_mdl.json +++ b/wren-ai-service/demo/sample_dataset/ecommerce_duckdb_mdl.json @@ -1,6 +1,6 @@ { "catalog": "wrenai", - "schema": "ecommerce", + "schema": "public", "models": [ { "name": "customers", @@ -12,7 +12,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "The Customer City, where the customer company is located. Also called 'customer segment'." + "description": "The Customer City, where the customer company is located. Also called \"customer segment\".", + "displayName": "City" } }, { @@ -21,7 +22,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A unique identifier for each customer in the data model." + "description": "A unique identifier for each customer in the data model.", + "displayName": "Id" } }, { @@ -30,64 +32,17 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A field indicating the state where the customer is located." + "description": "A field indicating the state where the customer is located.", + "displayName": "State" } }, { "name": "orders", "type": "orders", - "relationship": "CustomersOrders", + "relationship": "CustomersIdOrdersCustomerid", "isCalculated": false, "notNull": false, "properties": {} - }, - { - "name": "LatestRecord", - "type": "DATE", - "isCalculated": true, - "expression": "max(orders.PurchaseTimestamp)", - "notNull": false, - "properties": {} - }, - { - "name": "FirstRecord", - "type": "DATE", - "isCalculated": true, - "expression": "min(orders.PurchaseTimestamp)", - "notNull": false, - "properties": {} - }, - { - "name": "VIP", - "type": "BOOLEAN", - "isCalculated": true, - "expression": "sum(orders.Size) > 2", - "notNull": false, - "properties": {} - }, - { - "name": "OrderCount", - "type": "BIGINT", - "isCalculated": true, - "expression": "count(orders.OrderId)", - "notNull": false, - "properties": {} - }, - { - "name": "Debit", - "type": "DOUBLE", - "isCalculated": true, - "expression": "sum(orders.OrderBalance)", - "notNull": false, - "properties": {} - }, - { - "name": "ReviewRate", - "type": "DOUBLE", - "isCalculated": true, - "expression": "count(orders.IsReviewed = TRUE) / count(DISTINCT orders.OrderId)", - "notNull": false, - "properties": {} } ], "primaryKey": "Id", @@ -96,7 +51,8 @@ "properties": { "schema": "main", "catalog": "memory", - "description": "A table of customers who have made purchases, including their city" + "description": "A table of customers who have made purchases, including their city", + "displayName": "customers" } }, { @@ -109,7 +65,17 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A numerical value representing the cost of shipping for an item in an order." + "description": "A numerical value representing the cost of shipping for an item in an order.", + "displayName": "FreightValue" + } + }, + { + "name": "Id", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "displayName": "Id" } }, { @@ -118,7 +84,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "The sequential number of the order item in this order. Each order item in an order has its unique ItemNumber." + "description": "The sequential number of the order item in this order. Each order item in an order has its unique ItemNumber.", + "displayName": "ItemNumber" } }, { @@ -127,7 +94,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A VARCHAR value indicating the order that this order_item belongs to. The column is used to map the order_item to Orders model in the OrdersOrder_items relationship." + "description": "A VARCHAR value indicating the order that this order_item belongs to. The column is used to map the order_item to Orders model in the OrdersOrder_items relationship.", + "displayName": "OrderId" } }, { @@ -136,7 +104,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A numerical value representing the price of an item in an order." + "description": "A numerical value representing the price of an item in an order.", + "displayName": "Price" } }, { @@ -145,7 +114,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A VARCHAR value representing the product of this order_item. The column is used to map the order_item to Products model using ProductsOrder_items relationship." + "description": "A VARCHAR value representing the product of this order_item. The column is used to map the order_item to Products model using ProductsOrder_items relationship.", + "displayName": "ProductId" } }, { @@ -154,13 +124,14 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A date value indicating the limit by which an item should be shipped according to the order. It helps track the deadline for shipping items in the 'order_items' model." + "description": "A date value indicating the limit by which an item should be shipped according to the order. It helps track the deadline for shipping items in the \"order_items\" model.", + "displayName": "ShippingLimitDate" } }, { "name": "orders", "type": "orders", - "relationship": "OrdersOrder_items", + "relationship": "OrdersOrderidOrder_itemsOrderid", "isCalculated": false, "notNull": false, "properties": {} @@ -168,19 +139,20 @@ { "name": "products", "type": "products", - "relationship": "ProductsOrder_items", + "relationship": "ProductsIdOrder_itemsProductid", "isCalculated": false, "notNull": false, "properties": {} } ], - "primaryKey": "", + "primaryKey": "Id", "cached": false, "refreshTime": "30.00m", "properties": { "schema": "main", "catalog": "memory", - "description": "The model is used to store information about items in orders, including details like prices, product IDs, shipping limits, and relationships with orders and products tables." + "description": "The model is used to store information about items in orders, including details like prices, product IDs, shipping limits, and relationships with orders and products tables.", + "displayName": "order_items" } }, { @@ -193,7 +165,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column that represents the timestamp when the order was approved." + "description": "A column that represents the timestamp when the order was approved.", + "displayName": "ApprovedTimestamp" } }, { @@ -202,7 +175,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A unique identifier representing the customer who purchased this order." + "description": "A unique identifier representing the customer who purchased this order.", + "displayName": "CustomerId" } }, { @@ -211,7 +185,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column that represents the date when the order was delivered by the carrier." + "description": "A column that represents the date when the order was delivered by the carrier.", + "displayName": "DeliveredCarrierDate" } }, { @@ -220,7 +195,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column that represents the date when the order was delivered to the customer." + "description": "A column that represents the date when the order was delivered to the customer.", + "displayName": "DeliveredCustomerDate" } }, { @@ -229,7 +205,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column that represents the estimated delivery date of the order." + "description": "A column that represents the estimated delivery date of the order.", + "displayName": "EstimatedDeliveryDate" } }, { @@ -238,7 +215,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column that represents a unique identifier of this order." + "description": "A column that represents a unique identifier of this order.", + "displayName": "OrderId" } }, { @@ -247,7 +225,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column that represents the timestamp when the order was purchased." + "description": "A column that represents the timestamp when the order was purchased.", + "displayName": "PurchaseTimestamp" } }, { @@ -256,13 +235,14 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column representing the status of the order." + "description": "A column representing the status of the order.", + "displayName": "Status" } }, { "name": "customers", "type": "customers", - "relationship": "CustomersOrders", + "relationship": "CustomersIdOrdersCustomerid", "isCalculated": false, "notNull": false, "properties": {} @@ -270,7 +250,7 @@ { "name": "order_items", "type": "order_items", - "relationship": "OrdersOrder_items", + "relationship": "OrdersOrderidOrder_itemsOrderid", "isCalculated": false, "notNull": false, "properties": {} @@ -278,7 +258,7 @@ { "name": "reviews", "type": "reviews", - "relationship": "OrdersReviews", + "relationship": "OrdersOrderidReviewsOrderid", "isCalculated": false, "notNull": false, "properties": {} @@ -286,83 +266,10 @@ { "name": "payments", "type": "payments", - "relationship": "PaymentsOrders", + "relationship": "PaymentsOrderidOrdersOrderid", "isCalculated": false, "notNull": false, "properties": {} - }, - { - "name": "Rating", - "type": "DOUBLE", - "isCalculated": true, - "expression": "avg(reviews.Score)", - "notNull": false, - "properties": { - "description": "A calculated column representing averge score of the reviews associated with the order.", - "display_name": "Rating" - } - }, - { - "name": "ReviewCount", - "type": "BIGINT", - "isCalculated": true, - "expression": "count(reviews.Id)", - "notNull": false, - "properties": { - "description": "A calculated column representing the number of the reviews associated with the order.", - "display_name": "Review Count" - } - }, - { - "name": "IsReviewed", - "type": "BOOLEAN", - "isCalculated": true, - "expression": "count(reviews.Id) > 0", - "notNull": false, - "properties": {} - }, - { - "name": "Size", - "type": "BIGINT", - "isCalculated": true, - "expression": "count(order_items.ItemNumber)", - "notNull": false, - "properties": { - "description": "A calculated column representing the number of the items associated with the order.", - "display_name": "Package Size" - } - }, - { - "name": "Large", - "type": "BOOLEAN", - "isCalculated": true, - "expression": "count(order_items.ItemNumber) > 1", - "notNull": false, - "properties": {} - }, - { - "name": "TotalValue", - "type": "DOUBLE", - "isCalculated": true, - "expression": "sum(order_items.Price)", - "notNull": false, - "properties": {} - }, - { - "name": "PaidValue", - "type": "DOUBLE", - "isCalculated": true, - "expression": "sum(payments.Value)", - "notNull": false, - "properties": {} - }, - { - "name": "OrderBalance", - "type": "DOUBLE", - "isCalculated": true, - "expression": "sum(order_items.Price) - sum(payments.Value)", - "notNull": false, - "properties": {} } ], "primaryKey": "OrderId", @@ -371,20 +278,31 @@ "properties": { "schema": "main", "catalog": "memory", - "description": "A model representing the orders data." + "description": "A model representing the orders data.", + "displayName": "orders" } }, { "name": "payments", "refSql": "select * from main.payments", "columns": [ + { + "name": "Id", + "type": "VARCHAR", + "isCalculated": false, + "notNull": false, + "properties": { + "displayName": "Id" + } + }, { "name": "Installments", "type": "BIGINT", "isCalculated": false, "notNull": false, "properties": { - "description": "A column representing the number of installments in the payments data model." + "description": "A column representing the number of installments in the payments data model.", + "displayName": "Installments" } }, { @@ -393,7 +311,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column representing the order id associated with this payment. The column is used to map the payment to the order using PaymentsOrders relationship." + "description": "A column representing the order id associated with this payment. The column is used to map the payment to the order using PaymentsOrders relationship.", + "displayName": "OrderId" } }, { @@ -402,7 +321,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column representing the sequential number of the payment in its corresponding order. Each payment in the order has its unique sequential number." + "description": "A column representing the sequential number of the payment in its corresponding order. Each payment in the order has its unique sequential number.", + "displayName": "Sequential" } }, { @@ -411,7 +331,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column representing the type of payment in the payments data model." + "description": "A column representing the type of payment in the payments data model.", + "displayName": "Type" } }, { @@ -420,25 +341,27 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A column representing the value of the payment in the payments data model." + "description": "A column representing the value of the payment in the payments data model.", + "displayName": "Value" } }, { "name": "orders", "type": "orders", - "relationship": "PaymentsOrders", + "relationship": "PaymentsOrderidOrdersOrderid", "isCalculated": false, "notNull": false, "properties": {} } ], - "primaryKey": "", + "primaryKey": "Id", "cached": false, "refreshTime": "30.00m", "properties": { "schema": "main", "catalog": "memory", - "description": "A model representing the payment records, including installments, order IDs, sequential numbers, payment types, values, and relationships with orders." + "description": "A model representing the payment records, including installments, order IDs, sequential numbers, payment types, values, and relationships with orders.", + "displayName": "payments" } }, { @@ -451,7 +374,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A category that classifies the products in the data model." + "description": "A category that classifies the products in the data model.", + "displayName": "Category" } }, { @@ -460,7 +384,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A unique identifier assigned to each product in the data model." + "description": "A unique identifier assigned to each product in the data model.", + "displayName": "Id" } }, { @@ -469,27 +394,17 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A name of the product in the data model." + "description": "A name of the product in the data model.", + "displayName": "Name" } }, { "name": "order_items", "type": "order_items", - "relationship": "ProductsOrder_items", + "relationship": "ProductsIdOrder_itemsProductid", "isCalculated": false, "notNull": false, "properties": {} - }, - { - "name": "Revenue", - "type": "DOUBLE", - "isCalculated": true, - "notNull": false, - "expression": "sum(order_items.Price)", - "properties": { - "description": "The calculated field representing the product revenue by summing up the product price in each order.", - "display_name": "Revenue" - } } ], "primaryKey": "Id", @@ -498,7 +413,8 @@ "properties": { "schema": "main", "catalog": "memory", - "description": "A data model containing information about products such as category, ID, and name, with a relationship to order items." + "description": "A data model containing information about products such as category, ID, and name, with a relationship to order items.", + "displayName": "products" } }, { @@ -511,7 +427,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "The date when the answer was provided." + "description": "The date when the answer was provided.", + "displayName": "AnswerTimestamp" } }, { @@ -520,7 +437,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "The date when the review was created." + "description": "The date when the review was created.", + "displayName": "CreationTimestamp" } }, { @@ -529,7 +447,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "A unique identifier assigned to each review entry." + "description": "A unique identifier assigned to each review entry.", + "displayName": "Id" } }, { @@ -538,7 +457,8 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "The order id of the order which the review belongs to." + "description": "The order id of the order which the review belongs to.", + "displayName": "OrderId" } }, { @@ -547,165 +467,77 @@ "isCalculated": false, "notNull": false, "properties": { - "description": "The score associated with each review entry." + "description": "The score associated with each review entry.", + "displayName": "Score" } }, { "name": "orders", "type": "orders", - "relationship": "OrdersReviews", + "relationship": "OrdersOrderidReviewsOrderid", "isCalculated": false, "notNull": false, "properties": {} } ], - "primaryKey": "", + "primaryKey": "Id", "cached": false, "refreshTime": "30.00m", "properties": { "schema": "main", "catalog": "memory", - "description": "A model containing information about review of orders." + "description": "A model containing information about review of orders.", + "displayName": "reviews" } } ], "relationships": [ { - "name": "CustomersOrders", - "models": [ - "customers", - "orders" - ], + "name": "CustomersIdOrdersCustomerid", + "models": ["customers", "orders"], "joinType": "ONE_TO_MANY", - "condition": "customers.Id = orders.CustomerId", + "condition": "\"customers\".Id = \"orders\".CustomerId", "manySideSortKeys": [], "properties": {} }, { - "name": "OrdersOrder_items", - "models": [ - "orders", - "order_items" - ], + "name": "OrdersOrderidOrder_itemsOrderid", + "models": ["orders", "order_items"], "joinType": "ONE_TO_MANY", - "condition": "orders.OrderId = order_items.OrderId", + "condition": "\"orders\".OrderId = \"order_items\".OrderId", "manySideSortKeys": [], "properties": {} }, { - "name": "ProductsOrder_items", - "models": [ - "products", - "order_items" - ], + "name": "ProductsIdOrder_itemsProductid", + "models": ["products", "order_items"], "joinType": "ONE_TO_MANY", - "condition": "products.Id = order_items.ProductId", + "condition": "\"products\".Id = \"order_items\".ProductId", "manySideSortKeys": [], "properties": {} }, { - "name": "OrdersReviews", - "models": [ - "orders", - "reviews" - ], + "name": "OrdersOrderidReviewsOrderid", + "models": ["orders", "reviews"], "joinType": "ONE_TO_MANY", - "condition": "orders.OrderId = reviews.OrderId", + "condition": "\"orders\".OrderId = \"reviews\".OrderId", "manySideSortKeys": [], "properties": {} }, { - "name": "PaymentsOrders", - "models": [ - "payments", - "orders" - ], + "name": "PaymentsOrderidOrdersOrderid", + "models": ["payments", "orders"], "joinType": "MANY_TO_ONE", - "condition": "payments.OrderId = orders.OrderId", + "condition": "\"payments\".OrderId = \"orders\".OrderId", "manySideSortKeys": [], "properties": {} } ], "enumDefinitions": [], - "metrics": [ - { - "name": "Revenue", - "baseObject": "orders", - "dimension": [ - { - "name": "PurchaseTimestamp", - "type": "TIMESTAMP" - }, - { - "name": "CustomerId", - "type": "VARCHAR" - }, - { - "name": "Status", - "type": "VARCHAR" - }, - { - "name": "Size", - "type": "BIGINT" - }, - { - "name": "IsReviewed", - "type": "BOOLEAN" - } - ], - "measure": [ - { - "name": "Price", - "type": "DOUBLE", - "expression": "sum(order_items.Price)" - }, - { - "name": "MoM_Price", - "type": "DOUBLE", - "expression": "sum(order_items.Price) - COALESCE(LAG(sum(order_items.Price)) OVER (ORDER BY date_trunc('MONTH', PurchaseTimestamp)), 0) / COALESCE(LAG(sum(order_items.Price)) OVER (ORDER BY date_trunc('MONTH', PurchaseTimestamp)), 1)" - }, - { - "name": "NumberOfOrders", - "type": "BIGINT", - "expression": "count(OrderId)" - }, - { - "name": "Peak", - "type": "DOUBLE", - "expression": "max(TotalValue)" - } - ], - "timeGrain": [ - { - "name": "PurchaseTimestamp", - "refColumn": "PurchaseTimestamp", - "dateParts": [ - "YEAR", - "MONTH", - "DAY" - ] - } - ] - } - ], + "metrics": [], "cumulativeMetrics": [], - "views": [ - { - "name": "total_items_view", - "statement": "WITH customer_order_items AS (SELECT c.State, oi.ItemNumber FROM customers c JOIN orders o ON c.Id = o.CustomerId JOIN order_items oi ON o.OrderId = oi.OrderId) SELECT State, COUNT(ItemNumber) AS Total_Items_Purchased FROM customer_order_items GROUP BY State", - "properties": { - "question": "How many items have been purchased in orders placed by customers in each state?", - "summary": "The query counts the total number of items purchased by customers in each state.", - "viewId": "20b69c05-e9a3-4c4d-a421-06ba7e7a385a" - } - } - ], - "macros": [ - { - "name": "MoM", - "definition": "(calculation: Expression, timestamp: Expression) => ({{ calculation }} - COALESCE(LAG({{ calculation }}) OVER (ORDER BY {{ timestamp }}), 0)) / COALESCE(LAG({{ calculation }}) OVER (ORDER BY {{ timestamp }}), 1)" - } - ], + "views": [], + "macros": [], "dateSpine": { "unit": "DAY", "start": "1970-01-01", diff --git a/wren-ai-service/demo/utils.py b/wren-ai-service/demo/utils.py index be576dba9..aa702e1f8 100644 --- a/wren-ai-service/demo/utils.py +++ b/wren-ai-service/demo/utils.py @@ -366,30 +366,32 @@ def generate_mdl_metadata(mdl_model_json: dict): def prepare_duckdb(dataset_name: str): assert dataset_name in ["music", "nba", "ecommerce"] + DATASET_VERSION = "v0.3.0" + init_sqls = { - "music": """ -CREATE TABLE album AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/Music/Album.csv',header=true); -CREATE TABLE artist AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/Music/Artist.csv',header=true); -CREATE TABLE customer AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/Music/Customer.csv',header=true); -CREATE TABLE genre AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/Music/Genre.csv',header=true); -CREATE TABLE invoice AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/Music/Invoice.csv',header=true); -CREATE TABLE invoiceLine AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/Music/InvoiceLine.csv',header=true); -CREATE TABLE track AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/Music/Track.csv',header=true); + "music": f""" +CREATE TABLE album AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/Music/Album.csv',header=true); +CREATE TABLE artist AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/Music/Artist.csv',header=true); +CREATE TABLE customer AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/Music/Customer.csv',header=true); +CREATE TABLE genre AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/Music/Genre.csv',header=true); +CREATE TABLE invoice AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/Music/Invoice.csv',header=true); +CREATE TABLE invoiceLine AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/Music/InvoiceLine.csv',header=true); +CREATE TABLE track AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/Music/Track.csv',header=true); """, - "nba": """ -CREATE TABLE game AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/NBA/game.csv',header=true); -CREATE TABLE line_score AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/NBA/line_score.csv',header=true); -CREATE TABLE player_games AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/NBA/player_game.csv',header=true); -CREATE TABLE player AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/NBA/player.csv',header=true); -CREATE TABLE team AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/NBA/team.csv',header=true); + "nba": f""" +CREATE TABLE game AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/NBA/game.csv',header=true); +CREATE TABLE line_score AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/NBA/line_score.csv',header=true); +CREATE TABLE player_games AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/NBA/player_game.csv',header=true); +CREATE TABLE player AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/NBA/player.csv',header=true); +CREATE TABLE team AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/NBA/team.csv',header=true); """, - "ecommerce": """ -CREATE TABLE customers AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/E-Commerce/customers.csv',header=true); -CREATE TABLE order_items AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/E-Commerce/order_items.csv',header=true); -CREATE TABLE orders AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/E-Commerce/orders.csv',header=true); -CREATE TABLE payments AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/E-Commerce/payments.csv',header=true); -CREATE TABLE products AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/E-Commerce/products.csv',header=true); -CREATE TABLE reviews AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/E-Commerce/reviews.csv',header=true); + "ecommerce": f""" +CREATE TABLE customers AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/E-Commerce/customers.csv',header=true); +CREATE TABLE order_items AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/E-Commerce/order_items.csv',header=true); +CREATE TABLE orders AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/E-Commerce/orders.csv',header=true); +CREATE TABLE payments AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/E-Commerce/payments.csv',header=true); +CREATE TABLE products AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/E-Commerce/products.csv',header=true); +CREATE TABLE reviews AS FROM read_csv('https://wrenai-public.s3.amazonaws.com/demo/{DATASET_VERSION}/E-Commerce/reviews.csv',header=true); """, } diff --git a/wren-ai-service/src/pipelines/ask/followup_generation_pipeline.py b/wren-ai-service/src/pipelines/ask/followup_generation_pipeline.py index 403900640..a92e943fd 100644 --- a/wren-ai-service/src/pipelines/ask/followup_generation_pipeline.py +++ b/wren-ai-service/src/pipelines/ask/followup_generation_pipeline.py @@ -20,25 +20,84 @@ logger = logging.getLogger("wren-ai-service") text_to_sql_with_followup_user_prompt_template = """ ### TASK ### -Given the following user query and the history of the last query along with the generated SQL result, -generate appropriate SQL queries that match the user's current request. -Generate at most 3 SQL queries in order to interpret the user query in various plausible ways. +Given the following user's follow-up question and previous SQL query and summary, +generate at most 3 SQL queries in order to interpret the user's question in various plausible ways. ### DATABASE SCHEMA ### {% for document in documents %} {{ document.content }} {% endfor %} -Generated SQL Queries amd Summaries: +### EXAMPLES ### + +Example 1 +[INPUT] +Previous SQL Summary: A query to find the number of employees in each department. +Previous SQL Query: SELECT department, COUNT(*) as employee_count FROM employees GROUP BY department; +User's Question: How do I modify this to only show departments with more than 10 employees? + +[OUTPUT] { "results": [ { - "sql": "SELECT users.* FROM users JOIN purchases ON users.id = purchases.user_id WHERE users.sign_up_date >= '2023-01-01';", - "summary": "Users joined in 2023 with purchases." + "sql": "SELECT department, COUNT() as employee_count FROM employees GROUP BY department HAVING COUNT() > 10", + "summary": "Modified to show only departments with more than 10 employees." }, { - "sql": "SELECT DISTINCT users.* FROM users INNER JOIN purchases ON users.id = purchases.user_id WHERE users.sign_up_date >= '2023-01-01';", - "summary": "Unique users with purchases since 2023." + "sql": "SELECT department FROM employees GROUP BY department HAVING COUNT() > 10", + "summary": "Shows only the names of departments with more than 10 employees." + }, + { + "sql": "SELECT department, COUNT() as employee_count FROM employees WHERE department IN (SELECT department FROM employees GROUP BY department HAVING COUNT(*) > 10)", + "summary": "Lists departments and their employee count, including only those with more than 10 employees." + } + ] +} + +Example 2 +[INPUT] +Previous SQL Summary: A query to retrieve the total sales per product. +Previous SQL Query: SELECT product_id, SUM(sales) as total_sales FROM sales GROUP BY product_id; +User's Question: Can you adjust this to include the product name as well? + +[OUTPUT] +{ + "results": [ + { + "sql": "SELECT products.name, SUM(sales.sales) as total_sales FROM sales JOIN products ON sales.product_id = products.id GROUP BY products.name", + "summary": "Includes product name with total sales." + }, + { + "sql": "SELECT p.name, s.total_sales FROM (SELECT product_id, SUM(sales) as total_sales FROM sales GROUP BY product_id) s JOIN products p ON s.product_id = p.id", + "summary": "Joins product table to include names in the total sales summary." + }, + { + "sql": "SELECT p.name, IFNULL(SUM(s.sales), 0) as total_sales FROM products p LEFT JOIN sales s ON p.id = s.product_id GROUP BY p.name", + "summary": "Includes all products, even those with no sales, showing total sales with product names." + } + ] +} + +Example 3 +[INPUT] +Previous SQL Summary: Query to find the highest salary in each department. +Previous SQL Query: SELECT department_id, MAX(salary) as highest_salary FROM employees GROUP BY department_id; +User's Question: What if I want to see the employee names with the highest salary in each department? + +[OUTPUT] +{ + "results": [ + { + "sql": "SELECT department_id, employee_name, salary FROM employees WHERE (department_id, salary) IN (SELECT department_id, MAX(salary) FROM employees GROUP BY department_id)", + "summary": "Shows the names of employees who earn the highest salary in their respective departments." + }, + { + "sql": "SELECT e.department_id, e.employee_name, e.salary FROM employees e INNER JOIN (SELECT department_id, MAX(salary) as max_salary FROM employees GROUP BY department_id) d ON e.department_id = d.department_id AND e.salary = d.max_salary", + "summary": "Lists employees with the highest salary in each department." + }, + { + "sql": "WITH MaxSalaries AS (SELECT department_id, MAX(salary) as max_salary FROM employees GROUP BY department_id) SELECT e.department_id, e.employee_name, e.salary FROM employees e JOIN MaxSalaries m ON e.department_id = m.department_id AND e.salary = m.max_salary", + "summary": "Utilizes a CTE to display each department's highest earners along with their names and salaries." } ] } @@ -49,7 +108,8 @@ The final answer must be the JSON format like following: { "results": [ {"sql": , "summary": }, - {"sql": , "summary": } + {"sql": , "summary": }, + {"sql": , "summary": } ] } @@ -57,8 +117,8 @@ The final answer must be the JSON format like following: ### QUESTION ### Previous SQL Summary: {{ history.summary }} -Previous Generated SQL Query: {{ history.sql }} -Current User Query: {{ query }} +Previous SQL Query: {{ history.sql }} +User's Follow-up Question: {{ query }} Let's think step by step. """