[paopa]

This PR introduces significant improvements to the Question Recommendation pipeline:

1. Category-Based Question Generation:

   • Added support for specific categories in question generation
   • Implemented even distribution of questions across categories
   • Updated system prompt to better handle category-based generation

2. Enhanced Question Validation:

   • Added SQL validation for generated questions
   • Added SQL storage with validated questions

3. Regeneration Support:

   • Added regenerate flag to handle cases where initial generation doesn't meet requirements
   • Implemented partial updates for question categories
   • Added tracking of unfilled categories

4. Response Structure:

   • Modified response structure to organize questions by category
   • Added SQL query storage with each validated question
   • Removed explanation field from question model

Detailed Changes

1. Updated Question Model

class Question(BaseModel):
    question: str
    category: str
    # explanation field removed for simplicity

2. Enhanced Validation Process

@observe(name="Validate Question")
async def _validate_question(self, candidate: dict, request_id: str, ...):
    # Added SQL validation and storage
    # Implemented partial updates for categories

3. System Prompt Updates

• Refined guidelines for category-based question generation
• Added specific instructions for handling provided categories
• Updated output format requirements

Testing

Testing with docker

image: https://github.com/Canner/WrenAI/pkgs/container/wren-ai-service/302503606?tag=commit-d34c6d7

Testing from endpoint

{
  "mdl": "{\"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\":\"NumberOfOrders\",\"type\":\"BIGINT\",\"expression\":\"count(OrderId)\"},{\"name\":\"Peak\",\"type\":\"DOUBLE\",\"expression\":\"max(TotalValue)\"}],\"timeGrain\":[{\"name\":\"PurchaseTimestamp\",\"refColumn\":\"PurchaseTimestamp\",\"dateParts\":[\"YEAR\",\"MONTH\",\"DAY\"]}]},{\"name\":\"orders\",\"refSql\":\"select * from main.orders\",\"columns\":[{\"name\":\"CustomerId\",\"type\":\"VARCHAR\",\"isCalculated\":false,\"notNull\":false,\"properties\":{\"description\":\"The unique identifier of the customer who placed the order.\"}},{\"name\":\"OrderId\",\"type\":\"VARCHAR\",\"isCalculated\":false,\"notNull\":false,\"properties\":{\"description\":\"A unique identifier for each order in the dataset.\"}},{\"name\":\"PurchaseTimestamp\",\"type\":\"TIMESTAMP\",\"isCalculated\":false,\"notNull\":false,\"properties\":{\"description\":\"The date and time when the order was placed.\"}},{\"name\":\"TotalValue\",\"type\":\"DOUBLE\",\"isCalculated\":false,\"notNull\":false,\"properties\":{\"description\":\"The total monetary value of the order.\"}},{\"name\":\"order_items\",\"type\":\"order_items\",\"relationship\":\"OrdersOrderItems\",\"isCalculated\":false,\"notNull\":false,\"properties\":{}},{\"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\":{}}}",
  "previous_questions": [],
  "max_questions": 3,
  "max_categories": 3,
  "regenerate": true,
  "configuration": {
    "language": "English",
    "timezone": {
      "name": "Asia/Taipei"
    }
  }
}

Testing from python code

if __name__ == "__main__":
    import json

    from langfuse.decorators import langfuse_context

    from src.pipelines.generation import question_recommendation
    from src.pipelines.generation.sql_generation import SQLGeneration
    from src.pipelines.retrieval.retrieval import Retrieval
    from src.providers import generate_components
    from src.utils import load_env_vars

    load_env_vars()

    pipe_components = generate_components()
    service = QuestionRecommendation(
        pipelines={
            "question_recommendation": question_recommendation.QuestionRecommendation(
                **pipe_components["question_recommendation"]
            ),
            "retrieval": Retrieval(**pipe_components["retrieval"]),
            "sql_generation": SQLGeneration(**pipe_components["sql_generation"]),
        },
    )

    with open("sample/ecommerce_duckdb_mdl.json", "r") as file:
        mdl = json.load(file)

    request = QuestionRecommendation.Input(
        id="1",
        mdl=orjson.dumps(mdl),
        max_questions=3,
        max_categories=3,
        regenerate=True,
    )

    result = asyncio.run(service.recommend(request))
    print(result)

    langfuse_context.flush()

This PR significantly improves the question recommendation feature by adding category support, SQL validation, and regeneration capabilities.

[cyyeh]

LGTM