science_domain_generator.py

# science_domain_generator.py
# Dynamically generates a list of science/curiosity domains for HowItWorks sessions.
#
# Mirrors the pattern in story_genre_generator.py:
#   1. LLM generates 50 domains at high temperature
#   2. Python picks a random domain from the list
#   3. The list is cached to disk and refreshed when exhausted or stale
#
# Key functions:
#   generate_science_domains(basil_assessment) -> list[str]
#   pick_domain(basil_assessment) -> str

import os
import json
import random
from datetime import datetime
from openai import OpenAI

from config import (
    TASK_AGENT_MODEL,
    PROMPT_SCIENCE_DOMAIN_GENERATOR,
    SCIENCE_DOMAINS_FILE,
    HOWITWORKS_DIR,
)
from file_lock_utils import get_lock

from llm_client import create_smart_client
client = create_smart_client()

FALLBACK_DOMAINS = [
    "physics and forces",
    "chemistry and materials",
    "biology and living things",
    "astronomy and space",
    "weather and climate",
    "engineering and machines",
    "medicine and the human body",
    "geology and earth science",
    "ecology and ecosystems",
    "technology and computers",
    "food science and cooking",
    "ocean science and marine life",
    "animal behavior and adaptation",
    "energy and electricity",
    "light and optics",
    "sound and acoustics",
    "transportation and vehicles",
    "construction and architecture",
    "water and fluid dynamics",
    "fire and combustion",
    "magnetism and electromagnetism",
    "genetics and heredity",
    "robotics and automation",
    "timekeeping and clocks",
    "navigation and cartography",
    "agriculture and farming",
    "textiles and clothing manufacture",
    "communication technology",
    "printing and publishing",
    "mining and metallurgy",
]


def _load_prompt_template() -> str:
    """Load the science domain generator prompt template."""
    if os.path.exists(PROMPT_SCIENCE_DOMAIN_GENERATOR):
        with open(PROMPT_SCIENCE_DOMAIN_GENERATOR, "r") as f:
            return f.read()
    raise FileNotFoundError(
        f"Science domain generator prompt not found: {PROMPT_SCIENCE_DOMAIN_GENERATOR}"
    )


def _load_cached_domains() -> dict:
    """Load cached domains from disk. Returns full data dict or empty dict."""
    if not os.path.exists(SCIENCE_DOMAINS_FILE):
        return {}
    try:
        with open(SCIENCE_DOMAINS_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, IOError):
        return {}


def _save_domains(domains: list, age_band: int):
    """Save domains list to the cache file."""
    data = {
        "generated_at": datetime.now().isoformat(),
        "age_band": age_band,
        "domains": domains,
    }
    os.makedirs(os.path.dirname(SCIENCE_DOMAINS_FILE), exist_ok=True)
    with open(SCIENCE_DOMAINS_FILE, "w") as f:
        json.dump(data, f, indent=2)


def generate_science_domains(
    basil_assessment: dict,
    recent_domains: list = None,
    n: int = 50,
) -> list:
    """
    Generate a diverse list of science/curiosity domains via LLM.

    Args:
        basil_assessment: Dict with age_band, capabilities, etc.
        recent_domains: List of recently used domain names to avoid.
        n: Number of domains to request.

    Returns:
        List of domain name strings (deduplicated, filtered).
    """
    age_band = basil_assessment.get("age_band", 0)

    if recent_domains:
        recent_text = ", ".join(recent_domains[-30:])
    else:
        recent_text = "(none)"

    template = _load_prompt_template()
    prompt = template.format(recent_domains=recent_text)

    system_msg = f"Generate exactly {n} science domains. Output valid JSON only."

    try:
        response = client.chat.completions.create(
            model=TASK_AGENT_MODEL,
            messages=[
                {"role": "system", "content": system_msg},
                {"role": "user", "content": prompt},
            ],
            temperature=1.0,
            max_tokens=2000,
        )

        raw_output = response.choices[0].message.content.strip()

        if raw_output.startswith("```"):
            lines = raw_output.split("\n")
            raw_output = "\n".join(lines[1:-1])

        data = json.loads(raw_output)
        domains_list = data.get("domains", [])

        seen = set()
        recent_lower = set(d.lower().strip() for d in (recent_domains or []))
        filtered = []

        for domain in domains_list:
            if isinstance(domain, str):
                domain = domain.strip()
                if not domain:
                    continue
                key = domain.lower()
                if key in seen or key in recent_lower:
                    continue
                seen.add(key)
                filtered.append(domain)

        print(f"[ScienceDomain] Generated {len(filtered)} domains for age_band={age_band}")

        _save_domains(filtered, age_band)

        return filtered

    except json.JSONDecodeError as e:
        print(f"[ScienceDomain] JSON parse error: {e}")
        return FALLBACK_DOMAINS[:]
    except Exception as e:
        print(f"[ScienceDomain] Error: {e}")
        return FALLBACK_DOMAINS[:]


def pick_domain(basil_assessment: dict, used_domains_this_session: list = None) -> str:
    """
    Pick a random science domain for the current HowItWorks session.

    Loads the cached domain list, regenerating it if empty or missing.
    Avoids domains already used in this session (if provided).

    Args:
        basil_assessment: Dict with age_band, capabilities, etc.
        used_domains_this_session: Domains already tried this session (for retries).

    Returns:
        A domain string, e.g. "physics and forces".
    """
    used_this_session = set(
        d.lower().strip() for d in (used_domains_this_session or [])
    )

    with get_lock(SCIENCE_DOMAINS_FILE):
        cached = _load_cached_domains()
        domains = cached.get("domains", [])

        available = [d for d in domains if d.lower().strip() not in used_this_session]

        if len(available) < 5:
            print(f"[ScienceDomain] Only {len(available)} domains available, regenerating...")
            domains = generate_science_domains(
                basil_assessment,
                recent_domains=list(used_this_session),
            )
            available = [d for d in domains if d.lower().strip() not in used_this_session]

    if not available:
        available = [
            d for d in FALLBACK_DOMAINS
            if d.lower().strip() not in used_this_session
        ]
        if not available:
            available = FALLBACK_DOMAINS[:]

    chosen = random.choice(available)
    return chosen


if __name__ == "__main__":
    print("Testing Science Domain Generator...")
    print()

    test_assessment = {
        "age_band": 2,
        "capabilities": ["Some word production", "Basic pattern recognition"],
    }

    domains = generate_science_domains(test_assessment)
    print(f"\nGenerated {len(domains)} domains:")
    for i, domain in enumerate(domains, 1):
        print(f"  {i:>3}. {domain}")

    print(f"\n--- Random picks ---")
    for _ in range(5):
        domain = pick_domain(test_assessment)
        print(f"  Picked: {domain}")