breakpilot-lehrer/backend-lehrer/story_generator.py

"""
Story Generator — Creates short stories using vocabulary words.

Generates age-appropriate mini-stories (3-5 sentences) that incorporate
the given vocabulary words, marked with <mark> tags for highlighting.

Uses Ollama (local LLM) for generation.
"""

import os
import json
import logging
import requests
from typing import List, Dict, Any, Optional

logger = logging.getLogger(__name__)

OLLAMA_URL = os.getenv("OLLAMA_BASE_URL", "http://host.docker.internal:11434")
STORY_MODEL = os.getenv("STORY_MODEL", "llama3.1:8b")


def generate_story(
    vocabulary: List[Dict[str, str]],
    language: str = "en",
    grade_level: str = "5-8",
    max_words: int = 5,
) -> Dict[str, Any]:
    """
    Generate a short story incorporating vocabulary words.

    Args:
        vocabulary: List of dicts with 'english' and 'german' keys
        language: 'en' for English story, 'de' for German story
        grade_level: Target grade level
        max_words: Maximum vocab words to include (to keep story short)

    Returns:
        Dict with 'story_html', 'story_text', 'vocab_used', 'language'
    """
    # Select subset of vocabulary
    words = vocabulary[:max_words]
    word_list = [w.get("english", "") if language == "en" else w.get("german", "") for w in words]
    word_list = [w for w in word_list if w.strip()]

    if not word_list:
        return {"story_html": "", "story_text": "", "vocab_used": [], "language": language}

    lang_name = "English" if language == "en" else "German"
    words_str = ", ".join(word_list)

    prompt = f"""Write a short story (3-5 sentences) in {lang_name} for a grade {grade_level} student.
The story MUST use these vocabulary words: {words_str}

Rules:
1. The story should be fun and age-appropriate
2. Each vocabulary word must appear at least once
3. Keep sentences simple and clear
4. The story should make sense and be engaging

Write ONLY the story, nothing else. No title, no introduction."""

    try:
        resp = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": STORY_MODEL,
                "prompt": prompt,
                "stream": False,
                "options": {"temperature": 0.8, "num_predict": 300},
            },
            timeout=30,
        )
        resp.raise_for_status()
        story_text = resp.json().get("response", "").strip()
    except Exception as e:
        logger.error(f"Story generation failed: {e}")
        # Fallback: simple template story
        story_text = _fallback_story(word_list, language)

    # Mark vocabulary words in the story
    story_html = story_text
    vocab_found = []
    for word in word_list:
        if word.lower() in story_html.lower():
            # Case-insensitive replacement preserving original case
            import re
            pattern = re.compile(re.escape(word), re.IGNORECASE)
            story_html = pattern.sub(
                lambda m: f'<mark class="vocab-highlight">{m.group()}</mark>',
                story_html,
                count=1,
            )
            vocab_found.append(word)

    return {
        "story_html": story_html,
        "story_text": story_text,
        "vocab_used": vocab_found,
        "vocab_total": len(word_list),
        "language": language,
    }


def _fallback_story(words: List[str], language: str) -> str:
    """Simple fallback when LLM is unavailable."""
    if language == "de":
        return f"Heute habe ich neue Woerter gelernt: {', '.join(words)}. Es war ein guter Tag zum Lernen."
    return f"Today I learned new words: {', '.join(words)}. It was a great day for learning."