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