"""Autonomous study agent — plans and executes sessions with minimal user input.""" from __future__ import annotations import json import logging import time from collections.abc import Iterator from dataclasses import dataclass, field from typing import Any from plane_mode_scholar.core import llm from plane_mode_scholar.core.config import DEFAULT_USER_ID from plane_mode_scholar.core.orchestration import Orchestrator from plane_mode_scholar.gradio_ui.citations import linkify_citations from plane_mode_scholar.memory.scheduler import get_due_reviews, pick_study_focus from plane_mode_scholar.study.planner import SessionPlanner from plane_mode_scholar.study.quiz import QuizGenerator from plane_mode_scholar.storage.sqlite_store import SQLiteStore logger = logging.getLogger(__name__) AGENT_TOOLS = [ {"name": "ensure_pack", "description": "Load or create demo study materials"}, {"name": "start_session", "description": "Open a timed study session with memory recall"}, {"name": "explain_topic", "description": "Grounded explanation with [1] citations"}, {"name": "run_quiz", "description": "Generate and present a quiz question"}, {"name": "surface_memory", "description": "Show what the coach remembers about the learner"}, {"name": "plan_next", "description": "Decide the next study action from mastery + SRS"}, ] @dataclass class AgentEvent: phase: str message: str tool: str | None = None data: dict[str, Any] = field(default_factory=dict) cycle_id: int = 0 def to_dict(self) -> dict[str, Any]: return { "phase": self.phase, "message": self.message, "tool": self.tool, "data": self.data, "cycle_id": self.cycle_id, "timestamp": time.time(), } class StudyAgent: """Monitor → plan → act loop for autonomous study coaching.""" def __init__( self, store: SQLiteStore | None = None, orchestrator: Orchestrator | None = None, ) -> None: self.store = store or SQLiteStore() self.orchestrator = orchestrator or Orchestrator(self.store) self.planner = SessionPlanner(self.store) self.quiz_gen = QuizGenerator(self.store) self._cycle = 0 def _next_cycle(self) -> int: self._cycle += 1 return self._cycle def _emit(self, phase: str, message: str, tool: str | None = None, **data: Any) -> dict: return AgentEvent( phase=phase, message=message, tool=tool, data=data, cycle_id=self._next_cycle(), ).to_dict() def plan_next_action(self, pack_id: str, user_id: str, session_id: str | None) -> dict: due = get_due_reviews(self.store, pack_id, user_id) weak = self.orchestrator.mastery.get_weak_topics(pack_id) focus = pick_study_focus(self.store, pack_id, user_id, weak) turns = self.store.list_turns(session_id) if session_id else [] if not session_id: return {"action": "start_session", "topic": focus, "reason": "No active session"} if due: return { "action": "explain_topic", "topic": due[0].content[:80], "reason": f"{len(due)} review(s) due from last session", } if len(turns) < 2: return { "action": "explain_topic", "topic": focus, "reason": "Opening explanation for weak focus area", } if len(turns) < 6: return { "action": "run_quiz", "topic": focus, "reason": "Check understanding before moving on", } return { "action": "surface_memory", "topic": focus, "reason": "Session depth reached — show memory + suggest wrap-up", } def ensure_pack(self, user_id: str, pack_id: str | None = None) -> str: if pack_id: pack = self.store.get_pack(pack_id) if pack: return pack_id packs = self.store.list_packs(user_id=user_id) if packs: return packs[0].pack_id from plane_mode_scholar.gradio_ui.layout import load_demo_pack load_demo_pack(user_id) packs = self.store.list_packs(user_id=user_id) if not packs: raise RuntimeError("Could not create demo study pack") return packs[0].pack_id def start_session( self, pack_id: str, user_id: str, goals: list[str] | None = None ) -> dict[str, Any]: result = self.planner.start_session( pack_id, goals=goals or ["Review key exam topics"], duration_minutes=45, user_id=user_id ) return result def run_autopilot( self, user_id: str = DEFAULT_USER_ID, pack_id: str | None = None, max_steps: int = 5, ) -> Iterator[dict[str, Any]]: """One-click autonomous study flow — yields telemetry for SwarmGrid-style UI.""" yield self._emit("BOOT", "Plane Mode Scholar agent online", tool="ensure_pack") try: pack_id = self.ensure_pack(user_id, pack_id) pack = self.store.get_pack(pack_id) pack_name = pack.name if pack else pack_id yield self._emit( "PACK", f"Study pack ready: {pack_name}", tool="ensure_pack", pack_id=pack_id, ) session_result = self.start_session(pack_id, user_id) session_id = session_result["session"]["session_id"] due_count = session_result.get("due_count", 0) recall = session_result.get("retrieved_memories", []) yield self._emit( "SESSION", f"Session started — {due_count} due review(s), {len(recall)} memories loaded", tool="start_session", session_id=session_id, due_count=due_count, recall_preview=[m.get("content", "")[:80] for m in recall[:3]], recommended=session_result.get("recommended_start", ""), ) steps_done = 0 while steps_done < max_steps: plan = self.plan_next_action(pack_id, user_id, session_id) action = plan["action"] topic = plan.get("topic", "") yield self._emit( "PLAN", f"Next: {action} — {plan.get('reason', '')}", tool="plan_next", plan=plan, ) if action == "explain_topic": query = f"Explain {topic} clearly using my course materials" prepared = self.orchestrator.prepare_turn( query, pack_id, session_id, topic, user_id=user_id ) yield self._emit( "RETRIEVE", f"Retrieved {len(prepared.chunks)} chunks, {len(prepared.memories)} memories", tool="explain_topic", chunks=len(prepared.chunks), memories=len(prepared.memories), ) accumulated = "" t0 = time.time() for partial in llm.generate_stream(prepared.prompt): accumulated = partial yield self._emit( "STREAM", accumulated, tool="explain_topic", delta=partial, topic=topic, ) result = self.orchestrator.finalize_turn(prepared, accumulated) cited = linkify_citations(accumulated, prepared.chunks) yield self._emit( "EXPLAIN", cited, tool="explain_topic", response=cited, memory_writes=len(result.memory_writes), inference_ms=int((time.time() - t0) * 1000), ) elif action == "run_quiz": questions = self.quiz_gen.generate_quiz( pack_id, num_questions=1, topic=topic, llm_generate=llm.generate, user_id=user_id, ) if questions: q = questions[0] yield self._emit( "QUIZ", q.get("question", "Quiz ready"), tool="run_quiz", question=q, topic=topic, ) else: yield self._emit("QUIZ", "Could not generate quiz", tool="run_quiz") elif action == "surface_memory": active = self.store.list_memories( pack_id=pack_id, user_id=user_id, status="active" ) yield self._emit( "MEMORY", f"Coach remembers {len(active)} active memories", tool="surface_memory", memories=[m.to_dict() for m in active[:6]], ) steps_done += 1 yield self._emit( "DONE", "Autopilot segment complete — continue chatting or tap Fly again", tool=None, session_id=session_id, pack_id=pack_id, ) except Exception as e: logger.exception("Autopilot failed") yield self._emit("ERROR", str(e), tool=None) def agent_plan_json(self, pack_id: str, user_id: str, context: str) -> dict: """Optional LLM planner for Nemotron tool-calling style decisions.""" prompt = f"""You are a study coach agent. Pick the single best next action. Tools: {json.dumps([t["name"] for t in AGENT_TOOLS])} Context: {context} Pack has materials indexed. Reply JSON only: {{"action": "", "topic": "", "reason": ""}}""" raw = llm.generate(prompt, max_new_tokens=120) try: start = raw.find("{") end = raw.rfind("}") + 1 if start >= 0 and end > start: return json.loads(raw[start:end]) except json.JSONDecodeError: pass return self.plan_next_action(pack_id, user_id, None)