litellm/tests/logging_callback_tests/test_generic_api_callback.py

import io
import os



import asyncio
import litellm
import gzip
import httpx
import json
import logging
import time
from typing import Final
from unittest.mock import AsyncMock, patch

import pytest

from litellm import completion
from litellm._logging import verbose_logger
from litellm.integrations.gcs_pubsub.pub_sub import *
from datetime import datetime, timedelta
from litellm.types.utils import (
    StandardLoggingPayload,
    StandardLoggingModelInformation,
    StandardLoggingMetadata,
    StandardLoggingHiddenParams,
)

verbose_logger.setLevel(logging.DEBUG)
from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger


@pytest.mark.asyncio
async def test_generic_api_callback():
    """
    Test the GenericAPILogger callback with a standard logging payload.
    This test mocks the HTTP client and validates that the logger properly
    formats and sends the expected payload.
    """

    # Create a mock for the async_httpx_client's post method
    mock_post = AsyncMock()
    mock_post.return_value.status_code = 200
    mock_post.return_value.text = "OK"

    # Set up an endpoint for testing
    test_endpoint = "https://example.com/api/logs"
    test_headers = {"Authorization": "Bearer test_token"}
    os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint

    # Initialize the GenericAPILogger and set the mock
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint, headers=test_headers, flush_interval=1
    )
    generic_logger.async_httpx_client.post = mock_post
    litellm.callbacks = [generic_logger]

    # Make the completion call
    response = await litellm.acompletion(
        model="gpt-5.5",
        messages=[{"role": "user", "content": "Hello, world!"}],
        mock_response="hi",
        user="test_user",
    )

    # Wait for async flush
    await asyncio.sleep(3)

    # Assert httpx post was called
    mock_post.assert_called_once()

    # Get the actual request body from the mock
    actual_url = mock_post.call_args[1]["url"]
    print("##########\n")
    print(
        "logs were flushed to URL",
        actual_url,
        "with the following headers",
        mock_post.call_args[1]["headers"],
    )
    assert (
        actual_url == test_endpoint
    ), f"Expected URL {test_endpoint}, got {actual_url}"

    # Validate headers
    assert (
        mock_post.call_args[1]["headers"]["Content-Type"] == "application/json"
    ), "Content-Type should be application/json"

    # For the GenericAPILogger, it sends the payload directly as JSON in the data field
    json_data = mock_post.call_args[1]["data"]
    # Parse the JSON string

    print("##########\n")
    print("json_data", json_data)
    actual_request = json.loads(json_data)

    # The payload is a list of StandardLoggingPayload objects in the log queue
    assert isinstance(actual_request, list), "Request body should be a list"
    assert len(actual_request) > 0, "Request body list should not be empty"

    this_test_messages: Final = [{"role": "user", "content": "Hello, world!"}]
    mine: Final = [
        item for item in actual_request if item.get("messages") == this_test_messages
    ]
    assert (
        len(mine) == 1
    ), f"Expected this test's single call in the batch, got {len(mine)} of {len(actual_request)}"

    payload_item: StandardLoggingPayload = StandardLoggingPayload(**mine[0])
    print("##########\n")
    print(json.dumps(payload_item, indent=4))
    print("##########\n")

    # Basic assertions for standard logging payload
    assert payload_item["response_cost"] > 0, "Response cost should be greater than 0"
    assert payload_item["model"] == "gpt-5.5", "Model should be gpt-5.5"
    assert (
        payload_item["model_parameters"]["user"] == "test_user"
    ), "User should be test_user"
    assert payload_item["model"] == "gpt-5.5", "Model should be gpt-5.5"
    assert payload_item["messages"] == [
        {"role": "user", "content": "Hello, world!"}
    ], "Messages should be the same"
    assert (
        payload_item["response"]["choices"][0]["message"]["content"] == "hi"
    ), "Response should be hi"


@pytest.mark.asyncio
async def test_generic_api_callback_multiple_logs():
    """
    Test the GenericAPILogger callback with multiple chat completions
    """
    # Create a mock for the async_httpx_client's post method
    mock_post = AsyncMock()
    mock_post.return_value.status_code = 200
    mock_post.return_value.text = "OK"

    # Set up an endpoint for testing
    test_endpoint = "https://example.com/api/logs"
    test_headers = {"Authorization": "Bearer test_token"}
    os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint

    # Initialize the GenericAPILogger and set the mock
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint, headers=test_headers, flush_interval=5
    )
    generic_logger.async_httpx_client.post = mock_post
    litellm.callbacks = [generic_logger]

    # Make the completion call
    for _ in range(10):
        response = await litellm.acompletion(
            model="gpt-5.5",
            messages=[{"role": "user", "content": "Hello, world!"}],
            mock_response="hi",
            user="test_user",
        )

    # Wait for async flush
    await asyncio.sleep(6)

    # Assert httpx post was called
    mock_post.assert_called_once()

    # Get the actual request body from the mock
    actual_url = mock_post.call_args[1]["url"]
    print("##########\n")
    print(
        "logs were flushed to URL",
        actual_url,
        "with the following headers",
        mock_post.call_args[1]["headers"],
    )
    assert (
        actual_url == test_endpoint
    ), f"Expected URL {test_endpoint}, got {actual_url}"

    # For the GenericAPILogger, it sends the payload directly as JSON in the data field
    json_data = mock_post.call_args[1]["data"]
    # Parse the JSON string

    print("##########\n")
    print("json_data", json_data)
    actual_request = json.loads(json_data)

    # The payload is a list of StandardLoggingPayload objects in the log queue
    assert isinstance(actual_request, list), "Request body should be a list"
    assert len(actual_request) > 0, "Request body list should not be empty"
    assert (
        len(actual_request) == 10
    ), "Request body list should be 10 items, since we made 10 calls"

    # Validate all payload items
    for payload_item in actual_request:
        payload_item: StandardLoggingPayload = StandardLoggingPayload(**payload_item)
        print("##########\n")
        print(json.dumps(payload_item, indent=4))
        print("##########\n")

        assert (
            payload_item["response_cost"] > 0
        ), "Response cost should be greater than 0"
        assert payload_item["model"] == "gpt-5.5", "Model should be gpt-5.5"
        assert (
            payload_item["model_parameters"]["user"] == "test_user"
        ), "User should be test_user"
        assert payload_item["model"] == "gpt-5.5", "Model should be gpt-5.5"
        assert payload_item["messages"] == [
            {"role": "user", "content": "Hello, world!"}
        ], "Messages should be the same"
        assert (
            payload_item["response"]["choices"][0]["message"]["content"] == "hi"
        ), "Response should be hi"


@pytest.mark.asyncio
async def test_generic_api_callback_ndjson_format():
    """
    Test the GenericAPILogger callback with ndjson log format.
    Validates that logs are sent as newline-delimited JSON.
    """
    # Create a mock for the async_httpx_client's post method
    mock_post = AsyncMock()
    mock_post.return_value.status_code = 200
    mock_post.return_value.text = "OK"

    # Set up an endpoint for testing
    test_endpoint = "https://example.com/api/logs"
    test_headers = {"Authorization": "Bearer test_token"}
    os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint

    # Initialize the GenericAPILogger with ndjson format
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint,
        headers=test_headers,
        flush_interval=1,
        log_format="ndjson",  # Set NDJSON format
    )
    generic_logger.async_httpx_client.post = mock_post
    litellm.callbacks = [generic_logger]

    # Make multiple completion calls to generate multiple logs
    for i in range(3):
        response = await litellm.acompletion(
            model="gpt-5.5",
            messages=[{"role": "user", "content": f"Hello, world! {i}"}],
            mock_response="hi",
            user="test_user",
        )

    # Wait for async flush
    await asyncio.sleep(3)

    # Assert httpx post was called
    mock_post.assert_called_once()

    # Get the actual request body from the mock
    actual_url = mock_post.call_args[1]["url"]
    assert (
        actual_url == test_endpoint
    ), f"Expected URL {test_endpoint}, got {actual_url}"

    # Get the data sent
    ndjson_data = mock_post.call_args[1]["data"]
    print("##########\n")
    print("ndjson_data:", ndjson_data)
    print("##########\n")

    # Validate it's NDJSON format (newline-delimited)
    assert isinstance(ndjson_data, str), "Data should be a string for NDJSON"

    # Split by newlines and parse each line
    lines = ndjson_data.strip().split("\n")
    assert len(lines) == 3, f"Expected 3 lines of NDJSON, got {len(lines)}"

    # Validate each line is valid JSON
    for i, line in enumerate(lines):
        payload_item = json.loads(line)
        payload_item = StandardLoggingPayload(**payload_item)

        # Basic assertions
        assert (
            payload_item["response_cost"] > 0
        ), "Response cost should be greater than 0"
        assert payload_item["model"] == "gpt-5.5", "Model should be gpt-5.5"
        assert (
            payload_item["model_parameters"]["user"] == "test_user"
        ), "User should be test_user"


@pytest.mark.asyncio
async def test_generic_api_callback_single_format():
    """
    Test the GenericAPILogger callback with single log format.
    Validates that each log is sent as an individual request in parallel.
    """
    # Create a mock for the async_httpx_client's post method
    mock_post = AsyncMock()
    mock_post.return_value.status_code = 200
    mock_post.return_value.text = "OK"

    # Set up an endpoint for testing
    test_endpoint = "https://example.com/api/logs"
    test_headers = {"Authorization": "Bearer test_token"}
    os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint

    # Initialize the GenericAPILogger with single format
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint,
        headers=test_headers,
        flush_interval=1,  # Quick flush to trigger batch send
        log_format="single",  # Set single format
    )
    generic_logger.async_httpx_client.post = mock_post
    litellm.callbacks = [generic_logger]

    # Make 3 completion calls
    for i in range(3):
        response = await litellm.acompletion(
            model="gpt-5.5",
            messages=[{"role": "user", "content": f"Hello, world! {i}"}],
            mock_response="hi",
            user="test_user",
        )

    # Wait for async flush
    await asyncio.sleep(3)

    # Assert httpx post was called 3 times (once per log in batch)
    assert mock_post.call_count == 3, f"Expected 3 calls, got {mock_post.call_count}"

    # Validate each call sent a single log object (not an array)
    for call_idx in range(3):
        call_args = mock_post.call_args_list[call_idx]
        json_data = call_args[1]["data"]

        print(f"########## Call {call_idx} ##########")
        print("json_data:", json_data)

        # Parse and validate - should be a single object, not an array
        actual_request = json.loads(json_data)
        assert isinstance(
            actual_request, dict
        ), f"Call {call_idx}: Expected dict, got {type(actual_request)}"

        # Validate it's a valid StandardLoggingPayload
        payload_item = StandardLoggingPayload(**actual_request)
        assert (
            payload_item["response_cost"] > 0
        ), "Response cost should be greater than 0"
        assert payload_item["model"] == "gpt-5.5", "Model should be gpt-5.5"


@pytest.mark.asyncio
async def test_generic_api_callback_json_array_format_explicit():
    """
    Test the GenericAPILogger callback with explicit json_array format.
    Validates backward compatibility when explicitly set to json_array.
    """
    # Create a mock for the async_httpx_client's post method
    mock_post = AsyncMock()
    mock_post.return_value.status_code = 200
    mock_post.return_value.text = "OK"

    # Set up an endpoint for testing
    test_endpoint = "https://example.com/api/logs"
    test_headers = {"Authorization": "Bearer test_token"}
    os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint

    # Initialize the GenericAPILogger with explicit json_array format
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint,
        headers=test_headers,
        flush_interval=1,
        log_format="json_array",  # Explicitly set json_array
    )
    generic_logger.async_httpx_client.post = mock_post
    litellm.callbacks = [generic_logger]

    # Make multiple completion calls
    for i in range(5):
        response = await litellm.acompletion(
            model="gpt-5.5",
            messages=[{"role": "user", "content": f"Hello, world! {i}"}],
            mock_response="hi",
            user="test_user",
        )

    # Wait for async flush
    await asyncio.sleep(3)

    # Assert httpx post was called once (batched)
    mock_post.assert_called_once()

    # Get the data and validate it's a JSON array
    json_data = mock_post.call_args[1]["data"]
    actual_request = json.loads(json_data)

    assert isinstance(
        actual_request, list
    ), "Request body should be a list (JSON array)"
    assert len(actual_request) == 5, f"Expected 5 items, got {len(actual_request)}"

    # Validate each item
    for payload_item in actual_request:
        payload_item = StandardLoggingPayload(**payload_item)
        assert (
            payload_item["response_cost"] > 0
        ), "Response cost should be greater than 0"
        assert payload_item["model"] == "gpt-5.5", "Model should be gpt-5.5"


@pytest.mark.asyncio
async def test_generic_api_callback_sumologic_uses_ndjson():
    """
    Test that the sumologic callback uses ndjson format by default
    when loaded from generic_api_compatible_callbacks.json
    """
    # Create a mock for the async_httpx_client's post method
    mock_post = AsyncMock()
    mock_post.return_value.status_code = 200
    mock_post.return_value.text = "OK"

    # Set environment variable for sumologic
    os.environ["SUMOLOGIC_WEBHOOK_URL"] = (
        "https://collectors.sumologic.com/receiver/v1/http/test123"
    )

    # Initialize using callback_name (loads from JSON config)
    generic_logger = GenericAPILogger(callback_name="sumologic", flush_interval=1)
    generic_logger.async_httpx_client.post = mock_post
    litellm.callbacks = [generic_logger]

    # Verify the logger has ndjson format
    assert generic_logger.log_format == "ndjson", "Sumologic should use ndjson format"

    # Make completion calls
    for i in range(2):
        await litellm.acompletion(
            model="gpt-5.5",
            messages=[{"role": "user", "content": f"Test {i}"}],
            mock_response="response",
            user="test_user",
        )

    # Wait for async flush
    await asyncio.sleep(3)

    # Assert httpx post was called
    mock_post.assert_called_once()

    # Verify NDJSON format
    ndjson_data = mock_post.call_args[1]["data"]
    assert isinstance(ndjson_data, str), "Data should be a string for NDJSON"

    lines = ndjson_data.strip().split("\n")
    records: Final = [json.loads(line) for line in lines]

    this_test_messages: Final = [
        [{"role": "user", "content": f"Test {i}"}] for i in range(2)
    ]
    mine: Final = [
        record for record in records if record.get("messages") in this_test_messages
    ]
    assert (
        len(mine) == 2
    ), f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}"


@pytest.mark.asyncio
async def test_generic_api_callback_invalid_log_format():
    """
    Test that invalid log_format values raise a ValueError
    """
    test_endpoint = "https://example.com/api/logs"
    os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint

    with pytest.raises(ValueError, match="Invalid log_format"):
        GenericAPILogger(
            endpoint=test_endpoint,
            log_format="invalid_format",  # type: ignore  # Intentionally invalid for testing
        )


@pytest.mark.asyncio
async def test_generic_api_callback_retries_timeout_then_succeeds():
    """
    Test that GenericAPILogger retries LiteLLM timeout errors when configured.
    """
    test_endpoint = "https://example.com/api/logs"
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint,
        max_retries=1,
        retry_delay=0,
        timeout=0.2,
    )

    mock_post = AsyncMock()
    mock_post.side_effect = [
        litellm.Timeout(
            message="Connection timed out",
            model="default-model-name",
            llm_provider="litellm-httpx-handler",
        ),
        type("Response", (), {"status_code": 200})(),
    ]
    generic_logger.async_httpx_client.post = mock_post
    generic_logger.log_queue = [{"event": "timeout-retry"}]

    await generic_logger.async_send_batch()

    assert mock_post.call_count == 2
    first_call = mock_post.call_args_list[0][1]
    assert first_call["url"] == test_endpoint
    assert first_call["timeout"] == 0.2
    assert json.loads(first_call["data"]) == [{"event": "timeout-retry"}]


@pytest.mark.asyncio
async def test_generic_api_callback_retries_5xx_then_succeeds():
    """
    Test that GenericAPILogger retries transient HTTP 5xx errors when configured.
    """
    test_endpoint = "https://example.com/api/logs"
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint,
        max_retries=1,
        retry_delay=0,
    )

    request = httpx.Request("POST", test_endpoint)
    response = httpx.Response(status_code=503, request=request)
    mock_post = AsyncMock()
    mock_post.side_effect = [
        httpx.HTTPStatusError(
            "Server error",
            request=request,
            response=response,
        ),
        type("Response", (), {"status_code": 200})(),
    ]
    generic_logger.async_httpx_client.post = mock_post
    generic_logger.log_queue = [{"event": "5xx-retry"}]

    await generic_logger.async_send_batch()

    assert mock_post.call_count == 2


@pytest.mark.asyncio
async def test_generic_api_callback_does_not_retry_4xx():
    """
    Test that GenericAPILogger does not retry non-transient HTTP 4xx errors.
    """
    test_endpoint = "https://example.com/api/logs"
    generic_logger = GenericAPILogger(
        endpoint=test_endpoint,
        max_retries=2,
        retry_delay=0,
    )

    request = httpx.Request("POST", test_endpoint)
    response = httpx.Response(status_code=401, request=request)
    mock_post = AsyncMock()
    mock_post.side_effect = httpx.HTTPStatusError(
        "Unauthorized",
        request=request,
        response=response,
    )
    generic_logger.async_httpx_client.post = mock_post
    generic_logger.log_queue = [{"event": "4xx-no-retry"}]

    await generic_logger.async_send_batch()

    mock_post.assert_called_once()