Bug fix: Added logic when domain classify confidence is low

docs/architecture.md

@@ -93,21 +93,33 @@ IDs) that are not redacted by Presidio.
 
 | Attribute       | Value                                              |
 |-----------------|----------------------------------------------------|
-| Architecture    | `distilbert-base-uncased` + linear classification head |
-| Task            | Multi-class text classification                    |
-| Classes (6)     | `ecommerce`, `telecom`, `banking`, `cibil`, `insurance`, `general` |
-| Training data   | CFPB Consumer Complaint Database (3M+ rows) — one-time download from Kaggle. Save as `data/raw/complaints.csv`. |
-| Script          | `python -m src.classifier.train --cfpb_csv data/raw/complaints.csv --output_dir models/domain_classifier` |
-| Input           | Redacted complaint text (string, max 512 tokens)   |
-| Output          | `{domain: str, confidence: float, all_probs: dict}` |
-| Fine-tune time  | ~30 min CPU / ~5 min GPU (T4)                      |
-| Library         | HuggingFace `transformers` + `datasets`            |
+| Architecture         | `distilbert-base-uncased` + linear classification head |
+| Task                 | Multi-class text classification                    |
+| Classes (6)          | `ecommerce`, `telecom`, `banking`, `cibil`, `insurance`, `general` |
+| Training data        | CFPB Consumer Complaint Database (3M+ rows) — one-time download from Kaggle. Save as `data/raw/complaints.csv`. |
+| Script               | `python -m src.classifier.train --cfpb_csv data/raw/complaints.csv --output_dir models/domain_classifier` |
+| Input                | Redacted complaint text (string, max 512 tokens)   |
+| Output               | `DomainResult(domain: str, confidence: float, all_probs: dict, low_confidence: bool)` |
+| Confidence threshold | `0.50` — results below this set `low_confidence=True` |
+| Low-confidence path  | Agent asks user one clarifying domain question; does not proceed until user confirms |
+| Keyword fallback     | Used when no checkpoint exists; always returns `confidence=0.0`, `low_confidence=True` |
+| Fine-tune time       | ~30 min CPU / ~5 min GPU (T4)                      |
+| Library              | HuggingFace `transformers` + `datasets`            |
 
 **Why DistilBERT?**  
 DistilBERT is 40% smaller and 60% faster than BERT-base with only 3% accuracy
 loss. For a project with limited compute, it is the ideal starting point.
 The CFPB dataset maps naturally to our 6 classes after label remapping.
 
+**Low-confidence handling:**  
+`general` is the intentional catch-all class — the model never returns an error, only a
+domain + probability. However, a low probability on all classes (e.g., the complaint text
+is too short or ambiguous) means the winning domain is unreliable. When `confidence < 0.50`
+the `low_confidence` flag is set and the CMA agent pauses to ask the user one clarifying
+question ("Is this about e-commerce, telecom, banking, credit score, insurance, or other?")
+before continuing. The user's answer overrides the model's suggestion and is stored with
+`domain_source = "user_confirmed"` so later tools know the domain is authoritative.
+
 ---
 
 #### 2.1.2 EvidenceNER
@@ -193,7 +205,13 @@ You are G.U.I.D.E., an expert consumer complaint assistant.
 PII has already been redacted locally — work with placeholders as-is.
 
 Rules:
-1. Always classify the domain first using the classify_domain tool.
+1. Always classify the domain first using classify_domain().
+   • If low_confidence=false (≥ 0.50): store domain and proceed.
+   • If low_confidence=true (< 0.50 or keyword fallback): ask the user ONE
+     clarifying question ("Is this about e-commerce, telecom, banking, credit
+     score, insurance, or other?") before continuing. Store domain_source=
+     "user_confirmed" when the domain comes from the user.
+   • If classify_domain() errors: same clarifying question as above.
 2. Ask ONE targeted follow-up question at a time if information is missing.
 3. If documents are uploaded, always run process_document before drafting.
 4. HITL gate: Before calling draft_complaint, present extracted details

src/agent/prompts.py

@@ -30,15 +30,33 @@ step. Never attempt to echo, reconstruct, or guess the original PII behind any p
 
 ## Operating Rules
 
-**Rule 1 — Classify domain first.**
+**Rule 1 — Classify domain first; verify when uncertain.**
 At the start of every new complaint thread, call classify_domain() with the complaint text as the \
 very first action — before asking any clarifying questions, before extract_entities(), before \
-anything else. Store the returned domain in session memory:
+anything else.
+
+Inspect the result's `low_confidence` field:
+
+*Case A — low_confidence is false (confidence ≥ 0.50):*
+The model is confident. Store the domain and proceed:
     store_memory(key="domain", value=<domain string>)
-Every downstream decision (which questions to ask, which authority to recommend, which draft \
-template to use) depends on the domain. If classify_domain() returns an error, fall back to \
-asking the user once: "What type of service is involved? (e-commerce, telecom, banking, credit \
-score, insurance, or other)". Store the user's answer with the same key.
+    store_memory(key="domain_source", value="model")
+
+*Case B — low_confidence is true (confidence < 0.50, or keyword fallback was used):*
+The model is uncertain. Do NOT silently use the suggested domain. Instead, ask the user \
+exactly ONE clarifying question as your next response — do not call any other tool first:
+    "To route your complaint correctly, could you tell me which type of service this is about?
+     Choose one: e-commerce / online shopping · telecom or internet · banking ·
+     credit score or CIBIL · insurance · other"
+Wait for the user's reply. Map their answer to the closest domain label and store it:
+    store_memory(key="domain", value=<confirmed domain string>)
+    store_memory(key="domain_source", value="user_confirmed")
+
+*Case C — classify_domain() returns an error:*
+Ask the same clarifying question as in Case B.
+
+Once the domain is confirmed (from any case), continue with extract_entities() and the \
+remaining pipeline. Every downstream decision depends on the domain.
 
 **Rule 2 — One follow-up question at a time.**
 Collect the following six minimum required fields before drafting:

src/agent/tools.py

@@ -35,7 +35,12 @@ TOOL_DEFINITIONS: list[dict] = [
         "description": (
             "Classify a consumer complaint into one of six domains: "
             "ecommerce, telecom, banking, cibil, insurance, or general. "
-            "MUST be the very first tool called on every new complaint thread."
+            "MUST be the very first tool called on every new complaint thread. "
+            "The result includes a 'low_confidence' boolean field. "
+            "When low_confidence is true (model confidence < 0.50, or keyword "
+            "fallback was used), do NOT proceed with the suggested domain — "
+            "instead ask the user one clarifying question to confirm the domain "
+            "before continuing."
         ),
         "input_schema": {
             "type": "object",

src/classifier/model.py

@@ -42,7 +42,8 @@ class DomainResult:
     """Classification output for a single complaint."""
     domain: str
     confidence: float
-    all_probs: dict  # {domain_label: probability}
+    all_probs: dict       # {domain_label: probability}
+    low_confidence: bool = False  # True when confidence < DOMAIN_CONFIDENCE_THRESHOLD
 
 
 # ---------------------------------------------------------------------------

src/classifier/predict.py

@@ -22,6 +22,13 @@ logger = logging.getLogger(__name__)
 _DEFAULT_MODEL_DIR = "models/domain_classifier"
 _classifier: Optional[DomainClassifier] = None
 
+# Minimum model confidence required to trust the classification result.
+# Below this threshold the result is flagged as low_confidence=True and the
+# CMA agent is expected to ask the user a clarifying domain question instead
+# of proceeding automatically.  The keyword fallback (confidence=0.0 sentinel)
+# always sets low_confidence=True regardless of this threshold.
+DOMAIN_CONFIDENCE_THRESHOLD: float = 0.50
+
 # ---------------------------------------------------------------------------
 # Keyword fallback
 # ---------------------------------------------------------------------------
@@ -73,8 +80,9 @@ def _keyword_classify(text: str) -> DomainResult:
     logger.debug("Keyword fallback scores: %s → %s", scores, best)
     return DomainResult(
         domain=best,
-        confidence=0.0,  # sentinel: 0.0 signals keyword fallback, not a model score
+        confidence=0.0,   # sentinel: 0.0 signals keyword fallback, not a model score
         all_probs=all_probs,
+        low_confidence=True,  # always uncertain when falling back to keywords
     )
 
 
@@ -113,8 +121,12 @@ def classify(text: str, model_dir: str = _DEFAULT_MODEL_DIR) -> DomainResult:
     """
     Classify *text* and return a DomainResult.
 
-    If no checkpoint exists at *model_dir*, falls back to keyword heuristics
-    and returns a DomainResult with confidence=0.0 as a fallback signal.
+    Two sources of low_confidence=True:
+      1. No checkpoint exists → keyword fallback is used (confidence=0.0 sentinel).
+      2. Model exists but top-class probability < DOMAIN_CONFIDENCE_THRESHOLD.
+
+    Callers (and the CMA agent) must check low_confidence and ask the user a
+    clarifying domain question rather than proceeding with an uncertain result.
     """
     global _classifier
     if _classifier is None:
@@ -127,4 +139,15 @@ def classify(text: str, model_dir: str = _DEFAULT_MODEL_DIR) -> DomainResult:
             )
             return _keyword_classify(text)
         _classifier = DomainClassifier(model_dir)
-    return _classifier.predict(text)
+
+    result = _classifier.predict(text)
+
+    if result.confidence < DOMAIN_CONFIDENCE_THRESHOLD:
+        logger.info(
+            "DomainClassifier low confidence (%.2f < %.2f) for domain '%s' — "
+            "flagging for user clarification.",
+            result.confidence, DOMAIN_CONFIDENCE_THRESHOLD, result.domain,
+        )
+        result.low_confidence = True
+
+    return result