Update core/ai_engine.py

core/ai_engine.py

@@ -1,13 +1,12 @@
 """
 core/ai_engine.py
-All Groq API calls — prompts taken verbatim from the tested Colab notebook.
-Covers: task parsing, scheduling, journaling Q&A, context synthesis.
+All Groq API calls — task parsing, intelligent RAG-style scheduling, journaling.
 """
 
 import json
 import re
 import os
-from datetime import datetime, date
+from datetime import datetime, date, timedelta
 from copy import deepcopy
 
 from groq import Groq
@@ -20,10 +19,7 @@ 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 HuggingFace Space → Settings → Repository secrets."
-        )
+        raise ValueError("GROQ_API_KEY is not set. Add it in Space Settings -> Repository Secrets.")
     _client = Groq(api_key=key)
 
 
@@ -36,7 +32,6 @@ 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:
@@ -55,61 +50,28 @@ def safe_json_parse(text: str):
 
 # ── Module 1: Task Capture ────────────────────────────────────────────────────
 
-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
+TASK_CAPTURE_PROMPT = """You are a task classification assistant for a productivity app called Second Brain.
 
-3. urgency (string): Choose ONE from:
-   - Habit: recurring or routine task
-   - Urgent: hard deadline or time pressure
-   - Not Urgent: no specific deadline
-
-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
+Take the user's raw task description and return a structured JSON object.
 
+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. clarifications_needed (array): Short specific questions if uncertain about any dimension. Return [] if confident.
 
-7. 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 7 fields.
-
-Example: {"title": "Finish project proposal", "life_area": "Work", "urgency": "Urgent", "importance": "Move the Needle", "state_of_mind": "Flow", "time_estimate": 90, "clarifications_needed": []}"""
+RETURN: Valid JSON only. No markdown, no prose. All 7 fields always present."""
 
 
 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 context notes (use to inform classification): {'; '.join(notes[-3:])}"
+        context_hint += f"\n\nUser patterns: {'; '.join(notes[-3:])}"
     if user_goals:
         context_hint += f"\nUser goals: {'; '.join(user_goals[:5])}"
     if life_areas:
@@ -131,190 +93,227 @@ def parse_task_with_groq(raw_text: str, user_context: dict = None,
             "title": raw_text,
             "life_area": None, "urgency": None, "importance": None,
             "state_of_mind": None, "time_estimate": None,
-            "clarifications_needed": [
-                "Could you give more details about this task?",
-                "Which area of your life does this belong to?",
-                "Is this urgent or flexible?"
-            ]
+            "clarifications_needed": ["Could you give more details? Which area, how urgent, how long?"]
         }
     return result
 
 
-# ── Module 2: Scheduling ──────────────────────────────────────────────────────
+# ── Module 2: Intelligent RAG-style Task Scheduler ────────────────────────────
 
-SCHEDULING_SYSTEM_PROMPT = """You are an intelligent daily scheduler for a productivity app called The Second Brain.
+SMART_SCHEDULER_PROMPT = """You are an intelligent task scheduler for Second Brain.
 
-You receive a USER CONTEXT (preferences + learned patterns), a TASK LIST, and a SCHEDULING PROMPT.
-Return a time-blocked schedule as a JSON object.
+You assign UNSCHEDULED TASKS to specific future dates, like a smart personal assistant who understands the
+user's rhythms, goals, energy patterns, and the current time.
 
-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
+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 days by state_of_mind:
+   - Flow = peak days (Mon-Thu mornings if peak=Morning)
+   - Quick/Easy = any day, fill gaps
+   - Habit = today or tomorrow
+6. Spread load — don't stack everything on one day
+7. Tasks with deadlines must land BEFORE that deadline
 
 RETURN FORMAT (JSON only, no markdown):
 {
-  "schedule_date": "YYYY-MM-DD",
-  "scheduled_tasks": [
+  "assignments": [
     {
-      "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"
+      "task_id": 123,
+      "title": "Task title",
+      "assigned_date": "YYYY-MM-DD",
+      "reasoning": "1 sentence why this date"
     }
   ],
-  "deferred_tasks": [{"task_id": "...", "title": "...", "reason": "..."}],
-  "day_summary": "2-3 sentences on day structure and reasoning",
-  "warnings": ["any concerns e.g. day overloaded"]
-}"""
-
-
-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."""
+  "skipped": [
+    {
+      "task_id": 456,
+      "title": "Task title",
+      "reason": "why not assigned"
+    }
+  ],
+  "summary": "2-3 sentence plain-English explanation of what was scheduled and why"
+}
+
+NEVER assign to a past date. NEVER ignore an explicit user instruction about when to schedule."""
+
+
+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", {})
+
+    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:
+{chr(10).join(f'- {g}' for g in (user_goals or [])) or '(none set)'}"""
+
+    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."""
 
-    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,
-    )
+    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
 
-    result = safe_json_parse(response.choices[0].message.content.strip())
     if result is None:
-        result = {
-            "error": "Could not parse schedule response.",
-            "raw": response.choices[0].message.content
+        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["schedule_date"] = schedule_date
-    result["generated_at"] = datetime.now().isoformat()
+    # 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
     return result
 
 
 # ── Module 3: Journaling ──────────────────────────────────────────────────────
 
-JOURNAL_QUESTION_PROMPT = """You are a reflective journaling coach for a productivity app called The Second Brain.
+JOURNAL_QUESTION_PROMPT = """You are a reflective journaling coach for Second Brain.
 
-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.
+Given user context, today's tasks, and the conversation so far — decide what to ask next.
 
-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
+FOCUS: completion reasons, energy patterns, time estimate accuracy, flow placement, overall balance.
 
-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.
+RULES: ONE question at a time. Build on prior answers. After 5-7 exchanges, signal session_complete.
+Warm, efficient tone — 2-minute check-in.
 
-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}"""
+RETURN (JSON only):
+{"question": "...", "question_focus": "...", "session_complete": false}
+OR: {"question": null, "question_focus": null, "session_complete": true}"""
 
 
-SYNTHESIS_PROMPT = """You are a pattern recognition engine for a productivity app called The Second Brain.
+SYNTHESIS_PROMPT = """You are a pattern recognition engine for Second Brain.
 
-You have a completed journaling conversation. Extract learnings and return the UPDATED user context JSON.
+Given a completed journaling conversation, return the UPDATED user context JSON.
 
-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)
+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).
 
-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
+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).
 
-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."""
+RETURN: Complete updated context JSON only. No markdown. Conservative — only update on clear evidence."""
 
 
 def build_opening_question(context: dict, tasks_today: list) -> dict:
-    """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))
+    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 = "It looks like you didn't have any tasks scheduled today — was that intentional or did things go sideways?"
+        q = "No tasks were scheduled today — intentional, or did things go sideways?"
     elif completed == 0:
-        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?"
+        q = f"None of today's {total} tasks got marked complete — derailed, or the plan didn't fit?"
     elif completed == total:
-        q = f"You completed all {total} tasks today — great day! Did the schedule feel natural, or were you pushing through?"
+        q = f"You completed all {total} tasks — great day! Did it feel natural, or were you grinding through it?"
     elif len(incomplete) == 1:
-        q = f'You got almost everything done — the one task left was "{incomplete[0]["title"]}". What got in the way?'
+        q = f'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 — was that time, energy, or something else?"
+        q = f"You completed {completed}/{total} tasks ({rate}%). Tasks like {titles} didn't get done — 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 SCHEDULE (with completion):
+TODAY'S TASKS:
 {json.dumps(tasks_today, indent=2)}
 
-CONVERSATION SO FAR ({len(conversation_history)} exchanges):
+CONVERSATION ({len(conversation_history)} exchanges):
 {json.dumps(conversation_history, indent=2)}
 
-What should I ask next? Return session_complete: true if enough has been covered."""
+What should I ask next?"""
 
     response = _groq().chat.completions.create(
         model=GROQ_MODEL,
@@ -328,30 +327,25 @@ What should I ask next? Return session_complete: true if enough has been covered
 
     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:
-    """Synthesize conversation into updated context. Fallback to manual stats update if LLM fails."""
-    total     = len(tasks_today)
+def synthesize_journal(context: dict, tasks_today: list, conversation_history: list) -> dict:
+    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 (current):
+    user_message = f"""USER CONTEXT:
 {json.dumps(context, indent=2)}
 
-TODAY'S SCHEDULE + COMPLETION:
+TODAY'S TASKS:
 {json.dumps(tasks_today, indent=2)}
 
-Today's completion rate: {completion_rate} ({completed}/{total})
+Completion rate: {completion_rate} ({completed}/{total})
 
-FULL JOURNALING CONVERSATION:
+CONVERSATION:
 {json.dumps(conversation_history, indent=2)}
 
 Return the complete updated context JSON."""
@@ -367,9 +361,7 @@ 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
@@ -415,4 +407,4 @@ def create_blank_context(user_id, preferences: dict = None) -> dict:
             "avg_completion_rate":         0.0,
             "last_7_day_completion_rates": []
         }
-    }
+    }