agents.py

"""
agents.py
=========
The specialised agents (roadmap section 9). Each agent is a small, focused
unit with a single responsibility. `graph.py` wires them into a workflow.

Agents
------
ClassifierAgent : routes a question to an intent (subnet/linux/packet/concept).
RetrieverAgent  : pulls the most relevant notes from the vector store.
ReasoningAgent  : turns retrieved context into a clear explanation.
QuizAgent       : creates practice questions from the explanation.
FeedbackAgent   : grades responses and recommends next steps.
DebugAgent      : explains Linux/networking commands and subnet/packet queries.
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional

from .llm import get_llm
from .prompts import (
    CLASSIFIER_PROMPT,
    DEBUG_PROMPT,
    REASONING_PROMPT,
)
from .quiz import Quiz, generate_quiz
from .retriever import RetrievalResult, get_retriever
from .tools import linux_commands, packet_flow, subnet


# ---------------------------------------------------------------------------
# Classifier
# ---------------------------------------------------------------------------
class ClassifierAgent:
    """Decide the question's intent. Uses cheap rules first, LLM as backup."""

    LABELS = {"subnet", "linux", "packet_flow", "concept"}

    def classify(self, question: str) -> str:
        q = question.lower()
        # Packet-flow first: a question with two IPs + a movement word is about
        # how traffic travels, even though it also contains addresses.
        if packet_flow.try_parse_flow_question(question):
            return "packet_flow"
        if subnet.try_parse_subnet_question(question) or any(
            k in q for k in ["subnet", "cidr", "broadcast address", "usable host", "netmask"]
        ):
            return "subnet"
        if any(c in q.split() for c in linux_commands.list_commands()) or any(
            k in q for k in ["command", "ping", "traceroute", "tcpdump", "ip addr", "netstat", "ss ", "scp"]
        ):
            return "linux"
        # LLM tie-breaker (optional)
        llm = get_llm()
        if getattr(llm, "enabled", False):
            try:
                label = llm.complete(CLASSIFIER_PROMPT, question).strip().lower()
                label = re.sub(r"[^a-z_]", "", label)
                if label in self.LABELS:
                    return label
            except Exception:
                pass
        return "concept"


# ---------------------------------------------------------------------------
# Retriever
# ---------------------------------------------------------------------------
class RetrieverAgent:
    def __init__(self):
        self._retriever = get_retriever()

    def run(self, question: str, topic: Optional[str] = None) -> RetrievalResult:
        return self._retriever.retrieve(question, topic=topic)


# ---------------------------------------------------------------------------
# Reasoning
# ---------------------------------------------------------------------------
class ReasoningAgent:
    def run(self, question: str, retrieval: RetrievalResult) -> str:
        llm = get_llm()
        if retrieval.is_empty:
            base = (
                "I couldn't find supporting notes for that question in the knowledge "
                "base. Add relevant PDFs to data/raw_pdfs and run ingestion, or ask "
                "about a topic covered by your notes."
            )
            return base
        context = retrieval.context_block()
        user = f"CONTEXT:\n{context}\n\nQUESTION: {question}"
        return llm.complete(system=REASONING_PROMPT, user=user)


# ---------------------------------------------------------------------------
# Quiz
# ---------------------------------------------------------------------------
class QuizAgent:
    def run(self, topic: str, answer: str, context: str = "") -> Quiz:
        return generate_quiz(topic=topic, answer=answer, context=context)


# ---------------------------------------------------------------------------
# Debug (deterministic tools + LLM fallback)
# ---------------------------------------------------------------------------
class DebugAgent:
    """Handles subnet, packet-flow and Linux-command questions directly."""

    def run(self, question: str, intent: str) -> Optional[str]:
        if intent == "subnet":
            sol = subnet.try_parse_subnet_question(question)
            if sol:
                return subnet.format_solution(sol)
        if intent == "packet_flow":
            flow = packet_flow.try_parse_flow_question(question)
            if flow:
                return packet_flow.format_trace(flow)
        if intent == "linux":
            info = linux_commands.explain(question)
            if info:
                return linux_commands.format_explanation(info)
            llm = get_llm()
            if getattr(llm, "enabled", False):
                try:
                    return llm.complete(DEBUG_PROMPT, question)
                except Exception:
                    return None
        return None


# Convenience singletons
classifier_agent = ClassifierAgent()
retriever_agent = RetrieverAgent()
reasoning_agent = ReasoningAgent()
quiz_agent = QuizAgent()
debug_agent = DebugAgent()