assessment_agent.py

# assessment_agent.py
# Assesses Basil's developmental age_band based on full session transcript.

import os
import json
import re
from openai import OpenAI

from config import (
    ASSESSMENT_AGENT_MODEL,
    PROMPT_ASSESSMENT_AGENT,
)

from llm_client import create_smart_client
client = create_smart_client()


def load_prompt_template() -> str:
    """Load the assessment agent prompt template."""
    with open(PROMPT_ASSESSMENT_AGENT, "r") as f:
        return f.read()


def _extract_age_band_from_response(raw_output: str) -> int:
    """
    Extract age_band (0-7) from LLM response.
    
    Handles multiple formats:
    - JSON: {"assessed_age_band": 2}
    - Plain number: 2
    - Text with number: "age_band 2" or "band 2"
    
    Returns age_band or raises ValueError if not found.
    """
    text = raw_output.strip()
    
    # Try JSON format first
    try:
        # Remove markdown code blocks if present
        if text.startswith("```"):
            lines = text.split("\n")
            text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
        
        # Try to parse as JSON
        json_match = re.search(r'\{[^{}]*\}', text, re.DOTALL)
        if json_match:
            data = json.loads(json_match.group(0))
            if "assessed_age_band" in data:
                return int(data["assessed_age_band"])
    except (json.JSONDecodeError, KeyError, ValueError):
        pass
    
    # Try to extract plain number (0-7)
    # Look for standalone digits or "age_band X" or "band X"
    patterns = [
        r'\bage_band\s*[:=]?\s*(\d)',
        r'\bband\s*[:=]?\s*(\d)',
        r'\b(\d)\b',  # Any standalone digit
    ]
    
    for pattern in patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            age_band = int(match.group(1))
            if 0 <= age_band <= 7:
                return age_band
    
    raise ValueError(f"Could not extract age_band from: {raw_output[:100]}")


def assess_developmental_age_band(session_transcript: str, current_age_band: int) -> int:
    """
    Assess Basil's developmental age_band based on full session transcript.
    
    Args:
        session_transcript: Full session transcript (all turns)
        current_age_band: Current age_band (used as fallback on error)
    
    Returns:
        Assessed age_band (0-7), or current_age_band on error
    """
    template = load_prompt_template()
    
    # Format prompt with session transcript
    prompt = template.format(session_transcript=session_transcript)
    
    try:
        response = client.chat.completions.create(
            model=ASSESSMENT_AGENT_MODEL,
            messages=[
                {"role": "system", "content": "Answer with ONLY a number (0-7). No explanation."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,  # Low temperature for consistency
            max_tokens=50,  # Just need a number
        )
        
        raw_output = response.choices[0].message.content.strip()
        
        # Extract age_band from response
        assessed_age_band = _extract_age_band_from_response(raw_output)
        
        # Clamp to valid range
        assessed_age_band = max(0, min(7, assessed_age_band))
        
        return assessed_age_band
        
    except Exception as e:
        print(f"[Assessment Agent] Error assessing age_band: {e}")
        print(f"[Assessment Agent] Falling back to current_age_band: {current_age_band}")
        return current_age_band  # Safe fallback: no promotion on error


if __name__ == "__main__":
    # Test the assessment agent
    print("Testing Assessment Agent...")
    
    test_transcript = """Tutor: Hello everyone! Today we're going to learn about Animals.
Sophie: That sounds fun! Can we learn about cats?
Tutor: Yes! Cats are furry animals that say meow.
Sophie: And they have whiskers!
Tutor: Basil, can you say 'cat'?
Basil: cat
Tutor: Great job!
Sophie: Good try, Basil!"""
    
    result = assess_developmental_age_band(test_transcript, current_age_band=0)
    print(f"Assessed age_band: {result}")