tests/test_generator.py

import os
import pytest
from mediastorm.rag.prompts import build_system_prompt, build_context
from mediastorm.rag.generator import generate_response
from mediastorm.rag.retriever import RetrievalResult


GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "")


def test_prompt_includes_citation_format():
    """System prompt should instruct link-based citations."""
    prompt = build_system_prompt()
    assert "mediastorm.com" in prompt
    assert "slug" in prompt.lower() or "{slug}" in prompt


def test_prompt_no_transcript_references():
    """System prompt should not reference transcripts or segments."""
    prompt = build_system_prompt()
    assert "transcript" not in prompt.lower()
    assert "segment" not in prompt.lower()
    assert "Speaker Name" not in prompt
    assert "Timestamp" not in prompt


def test_context_formats_stories_with_enrichment():
    result = RetrievalResult(
        stories=[{
            "id": "s1",
            "text": "A story about elephants in Kenya.",
            "metadata": {
                "name": "Kenya Wildlife",
                "slug": "kenya-wildlife",
                "media_type": "video documentary",
                "published_year": 2015,
                "director": "Tim McLaughlin",
                "themes": "wildlife & conservation",
                "countries": "Kenya",
                "event_period": "2014-2015",
                "awards": "Emmy Winner (2015)",
            },
            "distance": 0.1,
        }],
    )
    context = build_context(result)
    assert "Kenya Wildlife" in context
    assert "video documentary" in context
    assert "Tim McLaughlin" in context
    assert "wildlife & conservation" in context
    assert "kenya-wildlife" in context


def test_context_handles_empty():
    result = RetrievalResult(stories=[])
    context = build_context(result)
    assert "no relevant" in context.lower()


@pytest.mark.asyncio
@pytest.mark.skipif(not GEMINI_KEY, reason="GEMINI_API_KEY not set")
async def test_generate_response_mentions_story():
    result = RetrievalResult(
        stories=[{
            "id": "s1",
            "text": "Photographer documents aging father's mortality.",
            "metadata": {
                "name": "A Shadow Remains",
                "slug": "a-shadow-remains",
                "media_type": "photo essay",
                "published_year": 2012,
                "director": "Phillip Toledano",
                "themes": "family & aging",
                "countries": "United States",
                "event_period": "2010-2012",
                "awards": "",
            },
            "distance": 0.1,
        }],
    )
    response = await generate_response(
        query="What is A Shadow Remains about?",
        retrieval_result=result,
    )
    assert "Shadow" in response or "Toledano" in response or "father" in response