AgentEvalTool/tests/unit/test_utils_llm.py

"""Unit tests for agenteval.utils.llm — the shared LLM utility functions."""

import json
import pytest

from agenteval.utils.llm import extract_reply_text, extract_content_from_llm_response, parse_json_from_llm_text


# ── extract_reply_text ───────────────────────────────────────────────────

def test_extract_reply_text_none():
    assert extract_reply_text(None) == ""


def test_extract_reply_text_str():
    assert extract_reply_text("hello") == "hello"


def test_extract_reply_text_dict_msgbody_dict():
    msg = {"msgBody": {"content": "inner text"}}
    assert extract_reply_text(msg) == "inner text"


def test_extract_reply_text_dict_msgbody_str():
    msg = {"msgBody": "plain body"}
    assert extract_reply_text(msg) == "plain body"


def test_extract_reply_text_dict_content_key():
    msg = {"content": "direct content"}
    assert extract_reply_text(msg) == "direct content"


def test_extract_reply_text_dict_msgbody_none_fallback_content():
    # msgBody is falsy → falls back to content key
    msg = {"msgBody": None, "content": "fallback"}
    assert extract_reply_text(msg) == "fallback"


def test_extract_reply_text_int_coerces_to_str():
    assert extract_reply_text(42) == "42"


def test_extract_reply_text_empty_dict():
    assert extract_reply_text({}) == ""


# ── extract_content_from_llm_response ────────────────────────────────────

def test_extract_content_string_format():
    data = {"choices": [{"message": {"content": "answer text"}}]}
    assert extract_content_from_llm_response(data) == "answer text"


def test_extract_content_block_array_text():
    data = {
        "choices": [{
            "message": {
                "content": [
                    {"type": "text", "text": "block one"},
                    {"type": "text", "text": "block two"},
                ]
            }
        }]
    }
    result = extract_content_from_llm_response(data)
    assert "block one" in result
    assert "block two" in result


def test_extract_content_block_array_skips_non_text():
    data = {
        "choices": [{
            "message": {
                "content": [
                    {"type": "tool_use", "id": "t1", "input": {}},
                    {"type": "text", "text": "real answer"},
                ]
            }
        }]
    }
    result = extract_content_from_llm_response(data)
    assert result == "real answer"


def test_extract_content_block_uses_content_key_fallback():
    # Some providers use "content" instead of "text" inside blocks
    data = {
        "choices": [{
            "message": {
                "content": [{"type": "text", "content": "via content key"}]
            }
        }]
    }
    assert extract_content_from_llm_response(data) == "via content key"


def test_extract_content_missing_choices():
    assert extract_content_from_llm_response({}) == ""


def test_extract_content_empty_choices():
    assert extract_content_from_llm_response({"choices": []}) == ""


def test_extract_content_integer_coerced():
    data = {"choices": [{"message": {"content": 123}}]}
    assert extract_content_from_llm_response(data) == "123"


def test_extract_content_empty_block_list():
    data = {"choices": [{"message": {"content": []}}]}
    assert extract_content_from_llm_response(data) == ""


# ── parse_json_from_llm_text ─────────────────────────────────────────────

def test_parse_json_object():
    result = parse_json_from_llm_text('{"score": 8, "reason": "good"}')
    assert result["score"] == 8


def test_parse_json_array():
    result = parse_json_from_llm_text('["a", "b", "c"]')
    assert result == ["a", "b", "c"]


def test_parse_json_with_surrounding_text():
    text = 'Here is the result: {"score": 7} and nothing else.'
    result = parse_json_from_llm_text(text)
    assert result["score"] == 7


def test_parse_json_array_with_preamble():
    text = 'Generated questions: ["q1", "q2", "q3"]'
    result = parse_json_from_llm_text(text)
    assert result == ["q1", "q2", "q3"]


def test_parse_json_raises_on_no_json():
    with pytest.raises((ValueError, json.JSONDecodeError)):
        parse_json_from_llm_text("no json here at all")


def test_parse_json_markdown_wrapped():
    text = '```json\n{"key": "value"}\n```'
    # The fallback bracket-search finds the { in the markdown
    result = parse_json_from_llm_text(text)
    assert result["key"] == "value"