llm-connect/tests/test_payload.py

from llm_connect._payload import merge_gemini_model_params, merge_openai_chat_model_params


STRUCTURED_SCHEMA = {
    "type": "object",
    "properties": {
        "summary": {"type": "string"},
        "recommendations": {"type": "array", "items": {"type": "string"}},
    },
    "required": ["summary", "recommendations"],
}


ACTIVITY_CORE_MODEL_PARAMS = {
    "reasoning_effort": "medium",
    "max_depth": 4,
    "json_schema": STRUCTURED_SCHEMA,
    "top_p": 0.8,
}


def test_openai_chat_model_params_translate_activity_core_shape():
    payload = {
        "model": "gpt-4.1-mini",
        "messages": [{"role": "user", "content": "triage"}],
        "temperature": 0.2,
        "max_tokens": 200,
    }

    merge_openai_chat_model_params(payload, ACTIVITY_CORE_MODEL_PARAMS)

    assert payload["response_format"] == {
        "type": "json_schema",
        "json_schema": {
            "name": "structured_output",
            "schema": STRUCTURED_SCHEMA,
            "strict": False,
        },
    }
    assert payload["top_p"] == 0.8
    assert "reasoning_effort" not in payload
    assert "max_depth" not in payload
    assert "json_schema" not in payload


def test_openai_chat_model_params_preserve_explicit_response_format():
    explicit = {
        "type": "json_schema",
        "json_schema": {
            "name": "custom",
            "schema": STRUCTURED_SCHEMA,
            "strict": True,
        },
    }
    payload = {"model": "gpt-4.1-mini", "messages": []}

    merge_openai_chat_model_params(
        payload,
        {"json_schema": STRUCTURED_SCHEMA, "response_format": explicit},
    )

    assert payload["response_format"] == explicit


def test_gemini_model_params_translate_activity_core_shape():
    payload = {
        "contents": [{"role": "user", "parts": [{"text": "triage"}]}],
        "generationConfig": {
            "temperature": 0.2,
            "maxOutputTokens": 200,
        },
    }

    merge_gemini_model_params(payload, ACTIVITY_CORE_MODEL_PARAMS)

    assert payload["generationConfig"]["responseMimeType"] == "application/json"
    assert payload["generationConfig"]["responseSchema"] == STRUCTURED_SCHEMA
    assert payload["generationConfig"]["topP"] == 0.8
    assert "reasoning_effort" not in payload
    assert "max_depth" not in payload
    assert "json_schema" not in payload