brain-of-reese/tests/unit/test_llm_probe.py

"""Unit tests: scripts/llm_probe.py --tools response parsing (phase 37, task 01).

The live probe talks to aipi; the parsing and verdict-classification logic
is factored into pure functions and is what this module pins. The fixtures
mirror the exact wire shapes observed live against ``turbo`` on 2026-08-26
(reasoning_content first, then indexed delta.tool_calls fragments).
"""
from __future__ import annotations

import json

from scripts.llm_probe import (
    classify_tool_calling,
    parse_tool_response_nonstreaming,
    parse_tool_response_streaming,
)


def test_parse_nonstreaming_tool_calls() -> None:
    payload = {
        "choices": [
            {
                "finish_reason": "tool_calls",
                "message": {
                    "role": "assistant",
                    "content": "",
                    "reasoning_content": "Let me call it.\n",
                    "tool_calls": [
                        {
                            "id": "abc123",
                            "type": "function",
                            "function": {"name": "get_time", "arguments": "{}"},
                        }
                    ],
                },
            }
        ]
    }
    out = parse_tool_response_nonstreaming(payload)
    assert out == {"finish_reason": "tool_calls", "calls": [("get_time", "{}")]}


def test_parse_nonstreaming_plain_content_answer() -> None:
    payload = {
        "choices": [
            {"finish_reason": "stop", "message": {"role": "assistant", "content": "It is noon."}}
        ]
    }
    assert parse_tool_response_nonstreaming(payload) == {
        "finish_reason": "stop",
        "calls": [],
    }


def test_parse_nonstreaming_empty_or_malformed() -> None:
    assert parse_tool_response_nonstreaming({"choices": []}) == {
        "finish_reason": None,
        "calls": [],
    }
    assert parse_tool_response_nonstreaming({}) == {"finish_reason": None, "calls": []}
    assert parse_tool_response_nonstreaming(None) == {"finish_reason": None, "calls": []}


def _sse(payload: dict) -> str:
    return "data: " + json.dumps(payload)


def test_parse_streaming_accumulates_fragments() -> None:
    """The live wire shape: id+name+partial args in chunk 1, args in chunk 2."""
    lines = [
        _sse({"choices": [{"delta": {"reasoning_content": "Let me think."}, "index": 0}]}),
        _sse(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [
                                {
                                    "index": 0,
                                    "id": "call_1",
                                    "type": "function",
                                    "function": {"name": "get_time", "arguments": "{"},
                                }
                            ]
                        },
                        "index": 0,
                    }
                ]
            }
        ),
        _sse(
            {
                "choices": [
                    {
                        "delta": {"tool_calls": [{"index": 0, "function": {"arguments": "}"}}]},
                        "index": 0,
                    }
                ]
            }
        ),
        _sse({"choices": [{"delta": {}, "finish_reason": "tool_calls", "index": 0}]}),
        "data: [DONE]",
    ]
    out = parse_tool_response_streaming(lines)
    assert out["finish_reason"] == "tool_calls"
    assert out["calls"] == [("get_time", "{}")]
    assert out["delta_chunks"] == 2
    assert out["indexed"] is True
    assert out["had_id"] is True
    assert out["arguments_in_deltas"] is True


def test_parse_streaming_no_tool_calls() -> None:
    lines = [
        _sse({"choices": [{"delta": {"content": "It is "}, "index": 0}]}),
        _sse({"choices": [{"delta": {"content": "noon."}, "index": 0}]}),
        _sse({"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}]}),
        "data: [DONE]",
    ]
    out = parse_tool_response_streaming(lines)
    assert out["finish_reason"] == "stop"
    assert out["calls"] == []
    assert out["delta_chunks"] == 0
    assert out["had_id"] is False
    assert out["arguments_in_deltas"] is False


def test_parse_streaming_stops_at_done_and_skips_malformed() -> None:
    lines = [
        "data: not-json",
        "",
        _sse(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [
                                {
                                    "index": 0,
                                    "id": "x",
                                    "type": "function",
                                    "function": {"name": "get_time", "arguments": "{}"},
                                }
                            ]
                        },
                        "index": 0,
                    }
                ]
            }
        ),
        "data: [DONE]",
        # After [DONE] nothing must be parsed:
        _sse({"choices": [{"delta": {"content": "should not appear"}, "index": 0}]}),
    ]
    out = parse_tool_response_streaming(lines)
    assert out["calls"] == [("get_time", "{}")]
    assert out["delta_chunks"] == 1
    assert out["finish_reason"] is None


def test_parse_streaming_multiple_calls_by_index() -> None:
    lines = [
        _sse(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [
                                {
                                    "index": 1,
                                    "id": "b",
                                    "function": {"name": "second", "arguments": "{\"a\": "},
                                },
                                {
                                    "index": 0,
                                    "id": "a",
                                    "function": {"name": "first", "arguments": "{}"},
                                },
                            ]
                        },
                        "index": 0,
                    }
                ]
            }
        ),
        _sse(
            {
                "choices": [
                    {
                        "delta": {
                            "tool_calls": [{"index": 1, "function": {"arguments": "1}"}}]
                        },
                        "index": 0,
                    }
                ]
            }
        ),
    ]
    out = parse_tool_response_streaming(lines)
    assert out["calls"] == [("first", "{}"), ("second", '{"a": 1}')]
    assert out["delta_chunks"] == 3


def _ns(ok: bool = True) -> dict:
    return {
        "finish_reason": "tool_calls" if ok else "stop",
        "calls": [("get_time", "{}")] if ok else [],
    }


def _st(**overrides: object) -> dict:
    result: dict = {
        "finish_reason": "tool_calls",
        "calls": [("get_time", "{}")],
        "delta_chunks": 2,
        "indexed": True,
        "had_id": True,
        "arguments_in_deltas": True,
    }
    result.update(overrides)
    return result


def test_classify_supported() -> None:
    assert classify_tool_calling(_ns(), _st()) == "supported"


def test_classify_not_supported_variants() -> None:
    # Non-streaming did not call the tool.
    assert classify_tool_calling(_ns(False), _st()) == "not-supported"
    # Streaming did not call the tool.
    assert classify_tool_calling(_ns(), _st(finish_reason="stop", calls=[])) == "not-supported"
    # Streamed, but not as delta.tool_calls chunks.
    assert classify_tool_calling(_ns(), _st(delta_chunks=0)) == "not-supported"
    # Delta chunks without the OpenAI "index" field.
    assert classify_tool_calling(_ns(), _st(indexed=False)) == "not-supported"
    # Delta chunks without a call "id".
    assert classify_tool_calling(_ns(), _st(had_id=False)) == "not-supported"
    # Intermittent: non-stream ok, stream answered in content.
    assert classify_tool_calling(_ns(), _st(finish_reason="stop")) == "not-supported"