Update generator_engine.py

generator_engine.py

@@ -1,73 +1,82 @@
 from __future__ import annotations
 
-from typing import List, Optional
 
-from models import RetrievedChunk
+def normalize_category(category: str | None) -> str:
+    c = (category or "").strip().lower()
 
+    if c in {"quantitative", "quant", "q", "math"}:
+        return "Quantitative"
+    if c in {"datainsight", "data_insight", "data insight", "di", "data"}:
+        return "DataInsight"
+    if c in {"verbal", "v"}:
+        return "Verbal"
+    if c in {"general", "", "unknown", "none", "null"}:
+        return "General"
 
-class GeneratorEngine:
+    return category or "General"
 
-    def __init__(self):
-        pass
 
-    def _build_prompt(
-        self,
-        user_text: str,
-        question_text: Optional[str],
-        topic: str,
-        intent: str,
-        retrieval_context: List[RetrievedChunk],
-    ) -> str:
+def classify_question(question_text: str, category: str | None = None) -> dict:
+    q = (question_text or "").lower()
+    normalized = normalize_category(category)
 
-        context_lines = []
-        for chunk in retrieval_context:
-            context_lines.append(f"- {chunk.text}")
+    if normalized == "Quantitative":
+        if ("percent" in q or "%" in q) and any(
+            k in q for k in ["then", "after", "followed by", "successive", "increase", "decrease", "discount"]
+        ):
+            return {"category": normalized, "topic": "percent", "type": "successive_percent"}
 
-        context_block = "\n".join(context_lines)
+        if "percent" in q or "%" in q:
+            return {"category": normalized, "topic": "percent", "type": "percent_change"}
 
-        prompt = f"""
-You are a GMAT quantitative tutor.
+        if "ratio" in q or ":" in q:
+            return {"category": normalized, "topic": "ratio", "type": "ratio_total"}
 
-Topic: {topic}
-Intent: {intent}
+        if "probability" in q or "chosen at random" in q:
+            return {"category": normalized, "topic": "probability", "type": "simple_probability"}
 
-Question:
-{question_text or user_text}
+        if "divisible" in q or "remainder" in q or "mod" in q:
+            return {"category": normalized, "topic": "number_theory", "type": "remainder_or_divisibility"}
 
-Relevant teaching notes:
-{context_block}
+        if "|" in q:
+            return {"category": normalized, "topic": "algebra", "type": "absolute_value"}
 
-Explain how to approach the problem without immediately revealing the answer.
-"""
+        if any(k in q for k in ["circle", "radius", "circumference", "triangle", "perimeter", "area"]):
+            return {"category": normalized, "topic": "geometry", "type": "geometry"}
 
-        return prompt.strip()
+        if any(k in q for k in ["average", "mean", "median"]):
+            return {"category": normalized, "topic": "statistics", "type": "average"}
 
-    def generate(
-        self,
-        user_text: str,
-        question_text: Optional[str] = None,
-        topic: str = "",
-        intent: str = "answer",
-        retrieval_context: Optional[List[RetrievedChunk]] = None,
-        chat_history=None,
-        **kwargs,
-    ) -> str:
+        if "sequence" in q:
+            return {"category": normalized, "topic": "sequence", "type": "sequence"}
 
-        retrieval_context = retrieval_context or []
+        if "=" in q:
+            return {"category": normalized, "topic": "algebra", "type": "equation"}
 
-        prompt = self._build_prompt(
-            user_text,
-            question_text,
-            topic,
-            intent,
-            retrieval_context,
-        )
+        return {"category": normalized, "topic": "quant", "type": "general"}
 
-        # simple fallback explanation system
-        if intent == "hint":
-            return "Focus on identifying the relationship between the quantities before calculating."
+    if normalized == "DataInsight":
+        if "percent" in q or "%" in q:
+            return {"category": normalized, "topic": "percent", "type": "percent_change"}
+        if any(k in q for k in ["mean", "median", "distribution"]):
+            return {"category": normalized, "topic": "statistics", "type": "distribution"}
+        if any(k in q for k in ["correlation", "scatter", "trend", "table", "chart"]):
+            return {"category": normalized, "topic": "data", "type": "correlation_or_graph"}
+        return {"category": normalized, "topic": "data", "type": "general"}
 
-        if intent in {"method", "walkthrough", "instruction"}:
-            return "Start by translating the wording of the problem into an equation or numerical relationship."
+    if normalized == "Verbal":
+        if "meaning" in q or "definition" in q:
+            return {"category": normalized, "topic": "vocabulary", "type": "definition"}
+        if "grammatically" in q or "sentence correction" in q:
+            return {"category": normalized, "topic": "grammar", "type": "sentence_correction"}
+        if "argument" in q or "author" in q:
+            return {"category": normalized, "topic": "reasoning", "type": "argument_analysis"}
+        return {"category": normalized, "topic": "verbal", "type": "general"}
 
-        return "Let’s work through the structure of the problem before jumping to the calculation."
+    if any(k in q for k in ["percent", "%", "ratio", "remainder", "divisible", "probability", "circle", "triangle", "="]):
+        return classify_question(question_text, "Quantitative")
+
+    if any(k in q for k in ["table", "chart", "scatter", "trend", "distribution"]):
+        return classify_question(question_text, "DataInsight")
+
+    return {"category": "General", "topic": "unknown", "type": "unknown"}