Update core/ai_engine.py

core/ai_engine.py

@@ -1,12 +1,13 @@
 """
 core/ai_engine.py
-All Groq API calls — task parsing, intelligent RAG-style scheduling, journaling.
+All Groq API calls — prompts taken verbatim from the tested Colab notebook.
+Covers: task parsing, scheduling, journaling Q&A, context synthesis.
 """
 
 import json
 import re
 import os
-from datetime import datetime, date, timedelta
+from datetime import datetime, date
 from copy import deepcopy
 
 from groq import Groq
@@ -19,7 +20,10 @@ def init_groq(api_key: str = None):
     global _client
     key = api_key or os.environ.get("GROQ_API_KEY", "")
     if not key:
-        raise ValueError("GROQ_API_KEY is not set. Add it in Space Settings -> Repository Secrets.")
+        raise ValueError(
+            "GROQ_API_KEY is not set. "
+            "Add it in HuggingFace Space → Settings → Repository secrets."
+        )
     _client = Groq(api_key=key)
 
 
@@ -32,6 +36,7 @@ def _groq() -> Groq:
 # ── Shared util ───────────────────────────────────────────────────────────────
 
 def safe_json_parse(text: str):
+    """Parse JSON, stripping markdown fences if present. Returns None on failure."""
     try:
         return json.loads(text)
     except json.JSONDecodeError:
@@ -50,34 +55,66 @@ def safe_json_parse(text: str):
 
 # ── Module 1: Task Capture ────────────────────────────────────────────────────
 
-TASK_CAPTURE_PROMPT = """You are a task classification assistant for a productivity app called Second Brain.
+TASK_CAPTURE_PROMPT = """You are a task classification assistant for a productivity app called The Second Brain.
+
+Your job is to take a user's raw task description and return a structured JSON object.
+
+Classify the task across these dimensions:
+
+1. title (string): A clean, concise, action-oriented task title. Fix grammar.
+
+2. life_area (string): Choose ONE from: Work, Health, Learning, Finance, Personal, Family, Other
+   - Work: job, meetings, deadlines, clients, projects
+   - Health: exercise, medical, diet, mental health
+   - Learning: courses, books, skills, studying
+   - Finance: bills, payments, budgeting, taxes, investments
+   - Personal: hobbies, errands, home maintenance
+   - Family: tasks involving family members
+   - Other: does not fit any category
 
-Take the user's raw task description and return a structured JSON object.
+3. urgency (string): Choose ONE from:
+   - Habit: recurring or routine task
+   - Urgent: hard deadline or time pressure
+   - Not Urgent: no specific deadline
 
-TODAY'S DATE: {today}
+4. importance (string): Choose ONE from:
+   - Move the Needle: very high impact
+   - Important: meaningful, should be done
+   - Not Important: low real impact
+
+5. state_of_mind (string): Choose ONE from:
+   - Quick: 5-10 mins, very low focus
+   - Easy: 10-20 mins, low focus
+   - Flow: deep concentration needed
+   - Personal: life admin, little goal impact
 
-Dimensions:
-1. title (string): Clean, action-oriented task title.
-2. life_area (string): ONE of: Work, Health, Learning, Finance, Personal, Family, Other
-3. urgency (string): ONE of: Habit | Urgent | Not Urgent
-4. importance (string): ONE of: Move the Needle | Important | Not Important
-5. state_of_mind (string): ONE of: Flow | Easy | Quick | Personal
 6. time_estimate (integer): Realistic minutes to complete.
-7. deadline_date (string | null): If the user mentions ANY deadline, due date, or "by X" — extract as YYYY-MM-DD.
-   Examples: "by Friday" -> next Friday's date, "due March 1" -> "2026-03-01", "end of month" -> last day of current month.
-   If no deadline mentioned, return null.
-8. clarifications_needed (array): Short specific questions if uncertain about any dimension. Return [] if confident.
 
-RETURN: Valid JSON only. No markdown, no prose. All 8 fields always present.
-CRITICAL: If user says "by [date]", "due [date]", "before [date]", "deadline [date]" — ALWAYS populate deadline_date."""
+7. deadline_date (string or null): Extract YYYY-MM-DD if user mentions any deadline.
+   Today is {today}. Examples: "by Friday" → next Friday's date, "due March 1" → "2026-03-01".
+   Return null if no deadline is mentioned.
+
+8. clarifications_needed (array of strings):
+   If NOT confident about a dimension, add a short specific question.
+   If everything is clear, return []
+
+STRICT RULES:
+- Return ONLY valid JSON. No markdown, no explanation.
+- Never guess if uncertain — ask a clarification question instead.
+- Always return all 8 fields.
+- ALWAYS extract deadline_date if user says "by X", "due X", "before X", "deadline X".
+
+Example: {"title": "Finish project proposal", "life_area": "Work", "urgency": "Urgent", "importance": "Move the Needle", "state_of_mind": "Flow", "time_estimate": 90, "deadline_date": null, "clarifications_needed": []}"""
 
 
 def parse_task_with_groq(raw_text: str, user_context: dict = None,
                          user_goals: list = None, life_areas: list = None) -> dict:
+    """Parse raw task text into structured dimensions using Groq."""
+    # Build context hint from AI memory + goals
     context_hint = ""
     if user_context and user_context.get("learned_patterns", {}).get("notes"):
         notes = user_context["learned_patterns"]["notes"]
-        context_hint += f"\n\nUser patterns: {'; '.join(notes[-3:])}"
+        context_hint += f"\n\nUser context notes (use to inform classification): {'; '.join(notes[-3:])}"
     if user_goals:
         context_hint += f"\nUser goals: {'; '.join(user_goals[:5])}"
     if life_areas:
@@ -102,233 +139,190 @@ def parse_task_with_groq(raw_text: str, user_context: dict = None,
             "life_area": None, "urgency": None, "importance": None,
             "state_of_mind": None, "time_estimate": None,
             "deadline_date": None,
-            "clarifications_needed": ["Could you give more details? Which area, how urgent, how long?"]
+            "clarifications_needed": [
+                "Could you give more details about this task?",
+                "Which area of your life does this belong to?",
+                "Is this urgent or flexible?"
+            ]
         }
     return result
 
 
-# ── Module 2: Intelligent RAG-style Task Scheduler ────────────────────────────
+# ── Module 2: Scheduling ──────────────────────────────────────────────────────
 
-SMART_SCHEDULER_PROMPT = """You are an intelligent task scheduler for Second Brain.
+SCHEDULING_SYSTEM_PROMPT = """You are an intelligent daily scheduler for a productivity app called The Second Brain.
 
-You assign UNSCHEDULED TASKS to specific future dates AND time blocks, like a smart personal assistant
-who understands the user's rhythms, goals, energy patterns, and current time.
+You receive a USER CONTEXT (preferences + learned patterns), a TASK LIST, and a SCHEDULING PROMPT.
+Return a time-blocked schedule as a JSON object.
 
-REASONING PROCESS:
-1. Read the user request carefully — honour it above all else
-   - "clear my day" / "nothing today" / "free today" = assign NOTHING to today
-   - "schedule for tomorrow" = assign to tomorrow
-   - "this week" = spread across next 5 days
-   - No explicit date = use next 1-7 days intelligently
-2. Never assign to a date/time in the past (current datetime is given)
-3. If current time is past 18:00, treat today as unavailable unless user explicitly asks
-4. Prioritise by: deadline proximity > urgency > importance > goal alignment
-5. Match tasks to time blocks by state_of_mind:
-   - Flow = peak focus window (morning if peak=Morning), 60-90 min blocks
-   - Easy = mid-morning or early afternoon, 20-45 min
-   - Quick = any gap, 10-20 min
-   - Personal = end of day or lunch, flexible
-6. Respect wake/sleep times. Add 10min breaks between tasks. Don't exceed max_flow_block.
-7. Spread load sensibly — don't overstack one day
-8. Tasks with deadlines must land BEFORE that deadline
+SCHEDULING RULES:
+- Respect wake_time and sleep_time from context
+- Place Flow tasks during the user's peak focus time
+- If avg_task_overrun_pct > 0, add buffer proportionally to time estimates
+- If flow_batch_capable is true, group Flow tasks; otherwise space them out
+- Place Quick and Easy tasks around transitions and low-energy windows
+- Place Personal/Habit tasks at day boundaries (start or end of day)
+- Urgent tasks are scheduled before Not Urgent ones
+- Move the Needle tasks get the best time slots
+- Add 5-10 min breaks between tasks
+- Respect any fixed commitments mentioned in the scheduling prompt
+- Do NOT schedule past sleep_time
+- If tasks won't realistically fit, put them in deferred_tasks
+- If context is minimal (new user), use sensible defaults
 
 RETURN FORMAT (JSON only, no markdown):
 {
-  "assignments": [
+  "schedule_date": "YYYY-MM-DD",
+  "scheduled_tasks": [
     {
-      "task_id": 123,
-      "title": "Task title",
-      "assigned_date": "YYYY-MM-DD",
-      "start_time": "09:00",
-      "end_time": "10:30",
-      "reasoning": "1 sentence why this slot"
+      "task_id": "(id from input or index)",
+      "title": "...",
+      "life_area": "...",
+      "start_time": "HH:MM",
+      "end_time": "HH:MM",
+      "duration_minutes": 60,
+      "state_of_mind": "...",
+      "scheduling_reason": "1-sentence explanation"
     }
   ],
-  "skipped": [
-    {
-      "task_id": 456,
-      "title": "Task title",
-      "reason": "why not assigned"
-    }
-  ],
-  "summary": "2-3 sentence plain-English explanation of what was scheduled and why",
-  "warnings": ["any overload or conflict warnings"]
-}
-
-NEVER assign to a past date/time. NEVER ignore explicit user scheduling instructions."""
-
-
-def smart_schedule_tasks(
-    tasks: list,
-    user_context: dict,
-    user_goals: list,
-    scheduling_prompt: str,
-    current_dt: datetime = None,
-) -> dict:
-    """
-    RAG-style scheduler: reads context + goals + patterns + current time + user request,
-    assigns each unscheduled task to a specific future date.
-    """
-    if current_dt is None:
-        current_dt = datetime.now()
-
-    today    = current_dt.date()
-    tomorrow = today + timedelta(days=1)
-    next_7   = [(today + timedelta(days=i)).isoformat() for i in range(8)]
-
-    prefs    = user_context.get("preferences", {})
-    patterns = user_context.get("learned_patterns", {})
-    feedback = user_context.get("scheduling_feedback", {})
-
-    goals_lines = "\n".join(f"- {g}" for g in (user_goals or [])) or "(none set)"
-    context_block = f"""CURRENT DATE/TIME: {current_dt.strftime('%Y-%m-%d %H:%M')} ({current_dt.strftime('%A')})
-TODAY: {today.isoformat()}  |  TOMORROW: {tomorrow.isoformat()}
-NEXT 7 DAYS: {', '.join(next_7)}
-
-USER PREFERENCES:
-- Wake: {prefs.get('wake_time', '08:00')} | Sleep: {prefs.get('sleep_time', '23:00')}
-- Peak focus: {prefs.get('focus_peak', 'Morning')}
-- Max flow block: {prefs.get('max_flow_block_minutes', 90)} min
-
-LEARNED PATTERNS:
-- Productive times: {patterns.get('productive_times', 'unknown')}
-- Low energy times: {patterns.get('low_energy_times', 'unknown')}
-- Avg task overrun: {patterns.get('avg_task_overrun_pct', 0)}%
-- Flow batching: {patterns.get('flow_batch_capable', 'unknown')}
-- Commonly skipped: {patterns.get('common_skipped_task_types', [])}
-- Notes: {'; '.join(patterns.get('notes', [])[-3:])}
-
-HISTORY:
-- Days tracked: {feedback.get('total_days_scheduled', 0)}
-- Avg completion: {round(feedback.get('avg_completion_rate', 0) * 100)}%
-
-GOALS:
-{goals_lines}"""
-
-    tasks_block = json.dumps([{
-        "task_id":       t.get("id", t.get("task_id")),
-        "title":         t.get("title"),
-        "life_area":     t.get("life_area"),
-        "urgency":       t.get("urgency"),
-        "importance":    t.get("importance"),
-        "state_of_mind": t.get("state_of_mind"),
-        "time_estimate": t.get("time_estimate"),
-        "deadline_date": t.get("deadline_date") or "none",
-    } for t in tasks], indent=2)
-
-    user_message = f"""{context_block}
-
-UNSCHEDULED TASKS ({len(tasks)} tasks):
-{tasks_block}
-
-USER REQUEST: "{scheduling_prompt}"
-
-Assign each task to the best date. Follow the user request precisely."""
+  "deferred_tasks": [{"task_id": "...", "title": "...", "reason": "..."}],
+  "day_summary": "2-3 sentences on day structure and reasoning",
+  "warnings": ["any concerns e.g. day overloaded"]
+}"""
 
-    try:
-        response = _groq().chat.completions.create(
-            model=GROQ_MODEL,
-            messages=[
-                {"role": "system", "content": SMART_SCHEDULER_PROMPT},
-                {"role": "user",   "content": user_message}
-            ],
-            max_tokens=2048,
-            temperature=0.15,
-        )
-        result = safe_json_parse(response.choices[0].message.content.strip())
-    except Exception as e:
-        result = None
 
+def generate_schedule(context: dict, tasks: list, scheduling_prompt: str,
+                      goals: list = None, schedule_date: str = None) -> dict:
+    if not schedule_date:
+        schedule_date = str(date.today())
+
+    goals_section = ""
+    if goals:
+        goals_section = "\nUSER GOALS:\n" + "\n".join(f"- {g}" for g in goals)
+
+    user_message = f"""Schedule Date: {schedule_date}
+
+USER CONTEXT:
+{json.dumps(context, indent=2)}
+{goals_section}
+TASKS TO SCHEDULE ({len(tasks)} tasks):
+{json.dumps(tasks, indent=2)}
+
+USER SCHEDULING PROMPT:
+{scheduling_prompt}
+
+Generate the optimal schedule."""
+
+    response = _groq().chat.completions.create(
+        model=GROQ_MODEL,
+        messages=[
+            {"role": "system", "content": SCHEDULING_SYSTEM_PROMPT},
+            {"role": "user",   "content": user_message}
+        ],
+        max_tokens=2048,
+        temperature=0.2,
+    )
+
+    result = safe_json_parse(response.choices[0].message.content.strip())
     if result is None:
-        return {
-            "assignments": [],
-            "skipped": [{"task_id": t.get("id", t.get("task_id")), "title": t.get("title"),
-                         "reason": "AI scheduling failed"} for t in tasks],
-            "summary": "Scheduling failed — please try again or rephrase your request."
+        result = {
+            "error": "Could not parse schedule response.",
+            "raw": response.choices[0].message.content
         }
 
-    # Safety pass: strip any assignments set in the past
-    safe_assignments = []
-    for a in result.get("assignments", []):
-        try:
-            assigned = date.fromisoformat(a["assigned_date"])
-            if assigned >= today:
-                safe_assignments.append(a)
-            else:
-                result.setdefault("skipped", []).append({
-                    "task_id": a.get("task_id"),
-                    "title":   a.get("title", ""),
-                    "reason":  f"AI tried to assign to past date {a['assigned_date']} — blocked"
-                })
-        except (ValueError, KeyError):
-            pass
-
-    result["assignments"] = safe_assignments
+    result["schedule_date"] = schedule_date
+    result["generated_at"] = datetime.now().isoformat()
     return result
 
 
 # ── Module 3: Journaling ──────────────────────────────────────────────────────
 
-JOURNAL_QUESTION_PROMPT = """You are a reflective journaling coach for Second Brain.
+JOURNAL_QUESTION_PROMPT = """You are a reflective journaling coach for a productivity app called The Second Brain.
 
-Given user context, today's tasks, and the conversation so far — decide what to ask next.
+You receive the user's context, today's schedule with completion status, and the conversation so far.
+Your job: decide what targeted question to ask NEXT.
 
-FOCUS: completion reasons, energy patterns, time estimate accuracy, flow placement, overall balance.
+FOCUS AREAS (cover what's most relevant, don't ask all):
+- Tasks not completed — why? wrong time? too tired? overestimated?
+- Tasks that took much longer than estimated
+- Energy levels — when were they sharp vs drained?
+- Whether Flow tasks were placed well or hard to start
+- Whether the day felt balanced or overloaded
+- Patterns the user noticed about themselves
 
-RULES: ONE question at a time. Build on prior answers. After 5-7 exchanges, signal session_complete.
-Warm, efficient tone — 2-minute check-in.
+RULES:
+- Ask ONE question at a time. Short and specific.
+- Build on previous answers — don't repeat covered ground.
+- After 4-6 good exchanges, signal completion.
+- Keep tone warm and efficient — 2-minute check-in, not therapy.
 
-RETURN (JSON only):
-{"question": "...", "question_focus": "...", "session_complete": false}
-OR: {"question": null, "question_focus": null, "session_complete": true}"""
+RETURN FORMAT (JSON only):
+{"question": "Your next question", "question_focus": "what aspect this targets", "session_complete": false}
+OR when done:
+{"question": null, "question_focus": null, "session_complete": true}"""
 
 
-SYNTHESIS_PROMPT = """You are a pattern recognition engine for Second Brain.
+SYNTHESIS_PROMPT = """You are a pattern recognition engine for a productivity app called The Second Brain.
 
-Given a completed journaling conversation, return the UPDATED user context JSON.
+You have a completed journaling conversation. Extract learnings and return the UPDATED user context JSON.
 
-UPDATE learned_patterns based on evidence: productive_times, low_energy_times, avg_task_overrun_pct,
-flow_batch_capable, best_life_areas_morning, common_skipped_task_types.
-Append 1-2 new insight notes (never remove existing).
+UPDATE these fields in learned_patterns based on conversation evidence:
+- productive_times: when user felt sharp/focused
+- low_energy_times: when they felt drained or skipped tasks
+- avg_task_overrun_pct: recalculate from actual vs estimated times mentioned
+- flow_batch_capable: update if user gave clear evidence
+- best_life_areas_morning: what they completed well before noon
+- common_skipped_task_types: patterns in what gets consistently skipped
+- notes: append 1-2 new insight notes (keep existing ones)
 
-ALWAYS UPDATE: scheduling_feedback (total_days_scheduled +1, rolling avg, last_7_rates),
-history_summary (append today, keep last 14), last_updated (now), version (+1).
+ALWAYS UPDATE:
+- scheduling_feedback.total_days_scheduled: +1
+- scheduling_feedback.avg_completion_rate: rolling average
+- scheduling_feedback.last_7_day_completion_rates: append today, keep last 7
+- history_summary: append brief today summary, keep last 14
+- last_updated: now
+- version: +1
 
-RETURN: Complete updated context JSON only. No markdown. Conservative — only update on clear evidence."""
+RULES:
+- Return ONLY the complete updated context JSON. Nothing else.
+- Never remove existing patterns — only update or append.
+- Be conservative — only update if there is clear evidence in the conversation."""
 
 
 def build_opening_question(context: dict, tasks_today: list) -> dict:
-    total      = len(tasks_today)
-    completed  = sum(1 for t in tasks_today if t.get("completed", False))
+    """Generate the first journal question based on task completion at a glance."""
+    total     = len(tasks_today)
+    completed = sum(1 for t in tasks_today if t.get("completed", False))
     incomplete = [t for t in tasks_today if not t.get("completed", False)]
 
     if total == 0:
-        q = "No tasks were scheduled today — intentional, or did things go sideways?"
+        q = "It looks like you didn't have any tasks scheduled today — was that intentional or did things go sideways?"
     elif completed == 0:
-        q = f"None of today's {total} tasks got marked complete — derailed, or the plan didn't fit?"
+        q = f"None of today's {total} tasks got marked complete — was the day unexpectedly derailed, or did the plan just not fit how your day went?"
     elif completed == total:
-        q = f"You completed all {total} tasks — great day! Did it feel natural, or were you grinding through it?"
+        q = f"You completed all {total} tasks today — great day! Did the schedule feel natural, or were you pushing through?"
     elif len(incomplete) == 1:
-        q = f'Almost everything done — the one task left was "{incomplete[0]["title"]}". What got in the way?'
+        q = f'You got almost everything done — the one task left was "{incomplete[0]["title"]}". What got in the way?'
     else:
         rate   = round(completed / total * 100)
         titles = ", ".join(f'"{t["title"]}"' for t in incomplete[:2])
-        q = f"You completed {completed}/{total} tasks ({rate}%). Tasks like {titles} didn't get done — time, energy, or something else?"
+        q = f"You completed {completed}/{total} tasks ({rate}%). Tasks like {titles} didn't get done — was that time, energy, or something else?"
 
     return {"question": q, "question_focus": "completion_overview", "session_complete": False}
 
 
-def get_next_journal_question(context: dict, tasks_today: list, conversation_history: list) -> dict:
+def get_next_journal_question(context: dict, tasks_today: list,
+                              conversation_history: list) -> dict:
     user_message = f"""USER CONTEXT:
 {json.dumps(context, indent=2)}
 
-TODAY'S TASKS:
+TODAY'S SCHEDULE (with completion):
 {json.dumps(tasks_today, indent=2)}
 
-CONVERSATION ({len(conversation_history)} exchanges):
+CONVERSATION SO FAR ({len(conversation_history)} exchanges):
 {json.dumps(conversation_history, indent=2)}
 
-What should I ask next?"""
+What should I ask next? Return session_complete: true if enough has been covered."""
 
     response = _groq().chat.completions.create(
         model=GROQ_MODEL,
@@ -342,25 +336,30 @@ What should I ask next?"""
 
     result = safe_json_parse(response.choices[0].message.content.strip())
     if result is None:
-        result = {"question": response.choices[0].message.content.strip(),
-                  "question_focus": "general", "session_complete": False}
+        result = {
+            "question": response.choices[0].message.content.strip(),
+            "question_focus": "general",
+            "session_complete": False
+        }
     return result
 
 
-def synthesize_journal(context: dict, tasks_today: list, conversation_history: list) -> dict:
-    total = len(tasks_today)
+def synthesize_journal(context: dict, tasks_today: list,
+                       conversation_history: list) -> dict:
+    """Synthesize conversation into updated context. Fallback to manual stats update if LLM fails."""
+    total     = len(tasks_today)
     completed = sum(1 for t in tasks_today if t.get("completed", False))
     completion_rate = round(completed / total, 2) if total > 0 else 0.0
 
-    user_message = f"""USER CONTEXT:
+    user_message = f"""USER CONTEXT (current):
 {json.dumps(context, indent=2)}
 
-TODAY'S TASKS:
+TODAY'S SCHEDULE + COMPLETION:
 {json.dumps(tasks_today, indent=2)}
 
-Completion rate: {completion_rate} ({completed}/{total})
+Today's completion rate: {completion_rate} ({completed}/{total})
 
-CONVERSATION:
+FULL JOURNALING CONVERSATION:
 {json.dumps(conversation_history, indent=2)}
 
 Return the complete updated context JSON."""
@@ -376,7 +375,9 @@ Return the complete updated context JSON."""
     )
 
     updated = safe_json_parse(response.choices[0].message.content.strip())
+
     if updated is None:
+        # Fallback: update stats manually if synthesis fails
         updated = deepcopy(context)
         updated["last_updated"] = datetime.now().isoformat()
         updated["version"]      = context.get("version", 1) + 1