cot

app.py

@@ -290,50 +290,56 @@ def render_chain(s1_text, s2_text, s3_text, status="done"):
     s2_active = status in ("running2", "running3", "done")
     s3_active = status in ("running3", "done")
 
-    # ── Stage 1 ───────────────────────────────────────────────────────────────
+    # ── Stage 1: Location pills + temporal line ───────────────────────────────
     if status == "running1":
         s1_content = f'<div class="thinking">Extracting features {_dots()}</div>'
     elif s1_text:
         tags = []
-        # Parse LOCATION INVENTORY bullets: "- Name: N visits, description"
-        in_inventory = False
-        for line in s1_text.splitlines():
-            line = line.strip()
-            if "LOCATION INVENTORY" in line.upper():
-                in_inventory = True
+        lines = s1_text.splitlines()
+        # Find frequency section: "Frequency of visits" or direct "- Name: N visits"
+        freq_mode = False
+        for line in lines:
+            stripped = line.strip()
+            if not stripped:
                 continue
-            if in_inventory:
-                if line.startswith("TEMPORAL") or line.startswith("SEQUENCE") or (line and not line.startswith("-") and not line.startswith("*") and len(line) > 40):
-                    break
-                if line.startswith("-"):
-                    # "- Name: N visits, type"  or  "- Name (N visits)"
-                    clean = line.lstrip("-").strip()
-                    # Shorten: keep "Name (N visits)" style
-                    m = re.match(r'(.+?):\s*(\d+)\s*visit', clean, re.IGNORECASE)
-                    if m:
-                        name = m.group(1).strip()
-                        n    = m.group(2)
-                        tags.append(f"{name} · {n}x")
-                    elif clean:
-                        tags.append(clean[:55])
+            # Enter frequency section
+            if re.search(r'frequency of visits', stripped, re.IGNORECASE):
+                freq_mode = True
+                continue
+            # Exit on next section heading
+            if freq_mode and re.match(r'^[A-Z][A-Z\s]+:', stripped) and not stripped.startswith('-'):
+                break
+            if freq_mode and stripped.startswith('-'):
+                # "- Name: N visits" or "  - Name: N visits"
+                m = re.match(r'-+\s*(.+?):\s*(\d+)\s*visit', stripped, re.IGNORECASE)
+                if m:
+                    tags.append(f"{m.group(1).strip()} · {m.group(2)}x")
+                    if len(tags) >= 8:
+                        break
+                continue
+            # Fallback: direct "- Name: N visits, description" under LOCATION INVENTORY
+            if not freq_mode:
+                m = re.match(r'-\s*(.+?):\s*(\d+)\s*visit', stripped, re.IGNORECASE)
+                if m and 'LOCATION' not in stripped.upper():
+                    tags.append(f"{m.group(1).strip()} · {m.group(2)}x")
                     if len(tags) >= 8:
                         break
 
-        # Fallback: also grab temporal summary line
+        # Temporal summary line
         temporal_line = ""
-        for line in s1_text.splitlines():
-            line = line.strip()
-            if "weekly distribution" in line.lower() or "weekday" in line.lower():
-                temporal_line = line.lstrip("-").strip()[:70]
+        for line in lines:
+            stripped = line.strip().lstrip('-').strip()
+            if re.search(r'\d+%\s*weekday', stripped, re.IGNORECASE):
+                temporal_line = stripped[:80]
                 break
 
-        tag_html = "".join(f'<span class="tag">{t}</span>' for t in tags)
+        tag_html  = "".join(f'<span class="tag">{t}</span>' for t in tags)
         temp_html = f'<div class="temporal-line">⏱ {temporal_line}</div>' if temporal_line else ""
         s1_content = f'<div class="tag-row">{tag_html}</div>{temp_html}'
     else:
         s1_content = '<div class="empty-hint">Press ▶ to start</div>'
 
-    # ── Stage 2 ───────────────────────────────────────────────────────────────
+    # ── Stage 2: 4 behavior dimensions ───────────────────────────────────────
     KEYS = [
         ("SCHEDULE",  ["ROUTINE", "SCHEDULE"]),
         ("ECONOMIC",  ["ECONOMIC", "SPENDING", "FINANCIAL"]),
@@ -343,77 +349,82 @@ def render_chain(s1_text, s2_text, s3_text, status="done"):
     if status == "running2":
         s2_content = f'<div class="thinking" style="color:#a06030">Analyzing behavior {_dots()}</div>'
     elif s2_text:
-        # Parse numbered sections
+        # Parse both "1. HEADING:" and "STEP 1: HEADING" formats
         sections = {}
         current_key = None
-        current_bullets = []
+        current_sents = []
         for line in s2_text.splitlines():
-            line = line.strip()
-            m = re.match(r'^\d+\.\s+(.+?)(?:\s+ANALYSIS)?(?:\s+PATTERNS)?(?:\s+INDICATORS)?:\s*$', line, re.IGNORECASE)
+            stripped = line.strip()
+            # Match "1. TITLE ANALYSIS:" or "STEP 1: TITLE ANALYSIS"
+            m = re.match(r'^(?:\d+\.|STEP\s*\d+:)\s+(.+?)(?:\s+ANALYSIS|\s+PATTERNS|\s+INDICATORS|\s+BEHAVIOR|\s+CHARACTERISTICS)?:\s*$', stripped, re.IGNORECASE)
             if m:
-                if current_key:
-                    sections[current_key] = current_bullets
-                current_key = m.group(1).upper()
-                current_bullets = []
-            elif current_key and line.startswith("-"):
-                bullet = line.lstrip("-").strip()
-                if bullet:
-                    current_bullets.append(bullet)
-        if current_key:
-            sections[current_key] = current_bullets
+                if current_key and current_sents:
+                    sections[current_key] = " ".join(current_sents)
+                current_key = m.group(1).upper().strip()
+                current_sents = []
+            elif current_key:
+                # Grab sub-bullet text (skip sub-headings like "1.1")
+                if re.match(r'^\d+\.\d+', stripped):
+                    # sub-heading: extract its inline content
+                    sub = re.sub(r'^\d+\.\d+[^:]*:\s*', '', stripped)
+                    if sub:
+                        current_sents.append(sub)
+                elif stripped.startswith('-'):
+                    current_sents.append(stripped.lstrip('-').strip())
+
+        if current_key and current_sents:
+            sections[current_key] = " ".join(current_sents)
 
         rows_html = ""
         for label, search_words in KEYS:
             val = "—"
-            for k, bullets in sections.items():
-                if any(w in k for w in search_words) and bullets:
-                    # Take first bullet, truncate at 2 sentences
-                    text = bullets[0]
-                    sentences = re.split(r'(?<=[.!?])\s+', text)
-                    val = " ".join(sentences[:2])
-                    if len(val) > 100:
-                        val = val[:97] + "..."
+            for k, text in sections.items():
+                if any(w in k for w in search_words) and text:
+                    # First 2 sentences
+                    sents = re.split(r'(?<=[.!?])\s+', text.strip())
+                    val = " ".join(sents[:2])
+                    if len(val) > 110:
+                        val = val[:107] + "..."
                     break
             rows_html += f'<div class="bkey">{label}</div><div class="bval">{val}</div>'
         s2_content = f'<div class="behavior-row">{rows_html}</div>'
     else:
         s2_content = '<div class="empty-hint">Waiting...</div>'
 
-    # ── Stage 3 ───────────────────────────────────────────────────────────────
+    # ── Stage 3: prediction + full reasoning ─────────────────────────────────
     if status == "running3":
         s3_content = f'<div class="thinking" style="color:#c0392b">Inferring demographics {_dots()}</div>'
     elif s3_text:
-        pred = reasoning = ""
-        lines = s3_text.splitlines()
-        i = 0
-        while i < len(lines):
-            line = lines[i].strip()
-            if line.startswith("INCOME_PREDICTION:"):
-                pred = line.replace("INCOME_PREDICTION:", "").strip()
-            elif line.startswith("INCOME_REASONING:"):
-                reasoning = line.replace("INCOME_REASONING:", "").strip()
-                # Collect continuation lines until blank or next key
-                i += 1
-                while i < len(lines):
-                    nxt = lines[i].strip()
-                    if not nxt or nxt.startswith("INCOME_") or re.match(r'^\d+\.', nxt):
-                        break
-                    reasoning += " " + nxt
-                    i += 1
-                continue
-            i += 1
+        pred = ""
+        reasoning_lines = []
+        in_reasoning = False
+        for line in s3_text.splitlines():
+            stripped = line.strip()
+            if stripped.startswith("INCOME_PREDICTION:"):
+                pred = stripped.replace("INCOME_PREDICTION:", "").strip()
+                in_reasoning = False
+            elif stripped.startswith("INCOME_REASONING:"):
+                in_reasoning = True
+                reasoning_lines.append(stripped.replace("INCOME_REASONING:", "").strip())
+            elif in_reasoning:
+                # Stop at second prediction block
+                if re.match(r'^2\.', stripped) or stripped.startswith("INCOME_CONFIDENCE"):
+                    break
+                if stripped:
+                    reasoning_lines.append(stripped)
 
-        # Truncate reasoning to ~2 sentences
-        sentences = re.split(r'(?<=[.!?])\s+', reasoning.strip())
-        short_reasoning = " ".join(sentences[:2])
-        if len(short_reasoning) > 160:
-            short_reasoning = short_reasoning[:157] + "..."
+        reasoning = " ".join(reasoning_lines).strip()
+        # Truncate to 3 sentences
+        sents = re.split(r'(?<=[.!?])\s+', reasoning)
+        short = " ".join(sents[:3])
+        if len(short) > 220:
+            short = short[:217] + "..."
 
         s3_content = f"""
         <div class="pred-block">
             <div class="pred-label">Income Prediction</div>
             <div class="pred-value">{pred or "—"}</div>
-            <div class="reasoning-text">{short_reasoning}</div>
+            <div class="reasoning-text">{short}</div>
         </div>"""
     else:
         s3_content = '<div class="empty-hint">Waiting...</div>'