scripts/test_silicon_llm_runtime.py

#!/usr/bin/env python3
"""Unit tests for dynamic LLM-backed silicon sampling runtime."""

from __future__ import annotations

import json
import sys
import unittest
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from silicon_llm import (  # noqa: E402
    build_silicon_llm_system_prompt,
    build_silicon_prompt_preview,
    build_silicon_llm_response_contract,
    parse_silicon_agent_response,
    run_silicon_llm_sampling,
)


QUESTIONS = [
    {
        "id": "q_likert_5",
        "title": "현 정부 국정 운영을 어떻게 평가하십니까?",
        "kind": "likert",
        "scale": 5,
        "lowLabel": "매우 부정",
        "highLabel": "매우 긍정",
    },
    {
        "id": "q_likert_4",
        "title": "지역 의료 접근성에 만족하십니까?",
        "kind": "likert",
        "scale": 4,
    },
    {
        "id": "q_open",
        "title": "가장 먼저 해결해야 할 문제는 무엇입니까?",
        "kind": "open",
    },
    {
        "id": "q_categorical",
        "title": "삶에서 가장 중요한 것을 하나만 고르십시오.",
        "kind": "categorical",
        "options": [
            {"id": "money", "label": "돈"},
            {"id": "honor", "label": "명예"},
            {"id": "health", "label": "건강"},
        ],
    },
]


class FakeLLM:
    def __init__(self) -> None:
        self.calls: list[dict[str, Any]] = []

    def complete_agent(self, *, agent: dict[str, Any], questions: list[dict[str, Any]], response_contract: dict[str, Any]) -> str:
        self.calls.append({"agent": agent, "questions": questions, "response_contract": response_contract})
        return json.dumps(
            {
                "answers": {
                    "q_likert_5": {"value": 4, "reason": "정책 평가는 대체로 긍정적입니다."},
                    "q_likert_4": {"value": 2, "reason": "의료 접근성은 지역에 따라 부족합니다."},
                    "q_open": {"text": "물가와 의료 접근성 개선이 가장 시급합니다.", "theme": "의료", "rationale": "생활비와 의료 이용 경험을 함께 고려했습니다."},
                    "q_categorical": {"optionId": "health", "label": "건강", "rationale": "고령 가족 돌봄 경험 때문에 건강을 우선했습니다."},
                }
            },
            ensure_ascii=False,
        )


class SiliconLLMRuntimeTests(unittest.TestCase):
    def test_contract_is_dynamic_for_question_mix(self) -> None:
        contract = build_silicon_llm_response_contract(QUESTIONS)

        answer_properties = contract["schema"]["properties"]["answers"]["properties"]
        self.assertEqual(set(answer_properties), {"q_likert_5", "q_likert_4", "q_open", "q_categorical"})
        self.assertEqual(answer_properties["q_likert_5"]["properties"]["value"]["minimum"], 1)
        self.assertEqual(answer_properties["q_likert_5"]["properties"]["value"]["maximum"], 5)
        self.assertEqual(answer_properties["q_likert_4"]["properties"]["value"]["maximum"], 4)
        self.assertIn("rationale", answer_properties["q_likert_4"]["required"])
        self.assertIn("text", answer_properties["q_open"]["required"])
        self.assertIn("rationale", answer_properties["q_open"]["required"])
        self.assertEqual(answer_properties["q_categorical"]["properties"]["optionId"]["enum"], ["money", "honor", "health"])

    def test_prompt_preview_uses_personalized_backend_prompt(self) -> None:
        system_prompt = build_silicon_llm_system_prompt()
        self.assertIn("1인칭", system_prompt)
        self.assertIn("최소 두 가지", system_prompt)
        self.assertIn("일반론", system_prompt)
        preview = build_silicon_prompt_preview(
            {
                "config": {
                    "sampleSize": 1,
                    "genders": [{"id": "female", "enabled": True, "weight": 1}],
                    "ages": [{"id": "40s", "enabled": True, "weight": 1}],
                    "locations": [{"id": "seoul", "enabled": True, "weight": 1}],
                    "locationOptions": [{"id": "seoul", "label": "서울", "parentRegion": "seoul", "level": "sido"}],
                    "questions": QUESTIONS,
                }
            }
        )
        self.assertEqual(preview["systemContent"], system_prompt)
        self.assertIn("required_json_contract", preview["userContent"])
        self.assertIn("q_categorical", preview["userContent"])

    def test_parser_extracts_all_dynamic_answers_from_fenced_json(self) -> None:
        raw = """
        생각 과정은 생략합니다.
        ```json
        {
          "answers": {
            "q_likert_5": {"value": "5", "reason": "강한 긍정"},
            "q_likert_4": {"score": 3, "reason": "보통 이상"},
            "q_open": {"answer": "교통과 돌봄이 필요합니다", "topic": "교통"},
            "q_categorical": {"label": "건강", "rationale": "생활 안정의 전제라고 보았습니다."}
          }
        }
        ```
        """
        parsed = parse_silicon_agent_response(raw, QUESTIONS, respondent_id="R0001")

        self.assertEqual(len(parsed["likertAnswers"]), 2)
        self.assertEqual(len(parsed["categoricalAnswers"]), 1)
        self.assertEqual(len(parsed["openAnswers"]), 1)
        self.assertEqual(parsed["likertAnswers"][0]["value"], 5)
        self.assertEqual(parsed["likertAnswers"][1]["value"], 3)
        self.assertEqual(parsed["likertAnswers"][0]["rationale"], "강한 긍정")
        self.assertEqual(parsed["categoricalAnswers"][0]["optionId"], "health")
        self.assertEqual(parsed["categoricalAnswers"][0]["label"], "건강")
        self.assertEqual(parsed["openAnswers"][0]["theme"], "교통")
        self.assertTrue(parsed["openAnswers"][0]["rationale"])

    def test_runtime_calls_one_llm_agent_for_all_questions_and_builds_stats(self) -> None:
        fake = FakeLLM()
        payload = {
            "config": {
                "sampleSize": 2,
                "genders": [{"id": "male", "enabled": True, "weight": 1}, {"id": "female", "enabled": True, "weight": 1}],
                "ages": [{"id": "30s", "enabled": True, "weight": 2}],
                "locations": [{"id": "seoul", "enabled": True, "weight": 2}],
                "locationOptions": [{"id": "seoul", "label": "서울", "parentRegion": "seoul", "level": "sido", "defaultWeight": 100, "short": "서울", "group": "시도"}],
                "personaAttributes": [{"id": "occ_office", "enabled": True, "weight": 2}],
                "nemotronFields": ["persona", "sex", "age"],
                "questions": QUESTIONS,
                "seed": 42,
            },
            "execution": {"max_agents": 2},
        }

        result = run_silicon_llm_sampling(payload, llm=fake)

        self.assertEqual(len(fake.calls), 2)
        self.assertTrue(all(len(call["questions"]) == 4 for call in fake.calls))
        self.assertEqual(len(result["respondents"]), 2)
        self.assertEqual(len(result["likertAnswers"]), 4)
        self.assertEqual(len(result["categoricalAnswers"]), 2)
        self.assertEqual(len(result["openAnswers"]), 2)
        self.assertTrue(all(answer.get("rationale") for answer in result["likertAnswers"]))
        self.assertTrue(all(answer.get("rationale") for answer in result["categoricalAnswers"]))
        self.assertTrue(all(answer.get("rationale") for answer in result["openAnswers"]))
        self.assertEqual([stat["questionId"] for stat in result["questionStats"]], ["q_likert_5", "q_likert_4", "q_open", "q_categorical"])
        self.assertEqual(result["questionStats"][0]["distribution"][3]["count"], 2)
        self.assertEqual(result["questionStats"][3]["distribution"][2]["count"], 2)
        self.assertEqual(result["regionStats"][0]["respondents"], 2)
        self.assertEqual(result["llmTrace"]["calls"], 2)


if __name__ == "__main__":
    unittest.main(verbosity=2)