litellm/tests/test_spend_logs.py

import os
# What this tests?
## Tests /spend endpoints.

import pytest, uuid, json
import asyncio
import aiohttp


async def generate_key(session, models=[], team_id=None):
    url = "http://0.0.0.0:4000/key/generate"
    headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"}
    data = {
        "models": models,
        "duration": None,
    }
    if team_id is not None:
        data["team_id"] = team_id

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request did not return a 200 status code: {status}")
        return await response.json()


async def chat_completion(session, key, model="gpt-3.5-turbo"):
    url = "http://0.0.0.0:4000/chat/completions"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {
        "model": model,
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": f"Hello! {uuid.uuid4()}"},
        ],
    }

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request did not return a 200 status code: {status}")

        return await response.json()


async def get_spend_logs(session, request_id=None, api_key=None):
    if api_key is not None:
        url = f"http://0.0.0.0:4000/spend/logs?api_key={api_key}"
    else:
        url = f"http://0.0.0.0:4000/spend/logs?request_id={request_id}"
    headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"}

    async with session.get(url, headers=headers) as response:
        status = response.status
        response_text = await response.text()

        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request did not return a 200 status code: {status}")
        return await response.json()


async def generate_org(session: aiohttp.ClientSession) -> dict:
    """
    Generate a new organization using the API.

    Args:
        session: aiohttp client session

    Returns:
        dict: Response containing org_id
    """
    url = "http://0.0.0.0:4000/organization/new"
    headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"}

    request_body = {
        "organization_alias": f"test-org-{uuid.uuid4()}",
    }

    async with session.post(url, headers=headers, json=request_body) as response:
        return await response.json()


async def generate_team(session: aiohttp.ClientSession, org_id: str) -> dict:
    """
    Generate a new team within an organization using the API.

    Args:
        session: aiohttp client session
        org_id: Organization ID to create the team in

    Returns:
        dict: Response containing team_id
    """
    url = "http://0.0.0.0:4000/team/new"
    headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"}
    data = {"organization_id": org_id}

    async with session.post(url, headers=headers, json=data) as response:
        return await response.json()


@pytest.mark.skip(
    reason="Flaky in CI: /spend/logs?request_id=... returns 500 even after a 20s wait for the spend log to be written. Spend-log accuracy is covered by tests/unit/proxy/spend_tracking/ and the proxy_spend_accuracy_tests CircleCI job."
)
@pytest.mark.asyncio
async def test_spend_logs_with_org_id():
    """
    - Create Organization
    - Create Team in organization
    - Create Key in organization
    - Make call (makes sure it's in spend logs)
    - Get request id from logs
    - Assert spend logs have correct org_id and team_id
    """
    async with aiohttp.ClientSession() as session:
        org_gen = await generate_org(session=session)
        print("org_gen: ", json.dumps(org_gen, indent=4, default=str))
        org_id = org_gen["organization_id"]
        team_gen = await generate_team(session=session, org_id=org_id)
        print("team_gen: ", json.dumps(team_gen, indent=4, default=str))
        team_id = team_gen["team_id"]
        key_gen = await generate_key(session=session, team_id=team_id)
        print("key_gen: ", json.dumps(key_gen, indent=4, default=str))
        key = key_gen["key"]
        response = await chat_completion(session=session, key=key)
        await asyncio.sleep(20)
        spend_logs_response = await get_spend_logs(
            session=session, request_id=response["id"]
        )
        print(
            "spend_logs_response: ",
            json.dumps(spend_logs_response, indent=4, default=str),
        )
        spend_logs_response = spend_logs_response[0]
        assert spend_logs_response["metadata"]["user_api_key_org_id"] == org_id
        assert spend_logs_response["metadata"]["user_api_key_team_id"] == team_id
        assert spend_logs_response["team_id"] == team_id


async def get_predict_spend_logs(session):
    url = "http://0.0.0.0:4000/global/predict/spend/logs"
    headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"}
    data = {
        "data": [
            {
                "date": "2024-03-09",
                "spend": 200000,
                "api_key": "sk-test-mock-api-key-456",
            }
        ]
    }

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request did not return a 200 status code: {status}")
        return await response.json()


@pytest.mark.skip(reason="datetime in ci/cd gets set weirdly")
@pytest.mark.asyncio
async def test_get_predicted_spend_logs():
    """
    - Create key
    - Make call (makes sure it's in spend logs)
    - Get request id from logs
    """
    async with aiohttp.ClientSession() as session:
        result = await get_predict_spend_logs(session=session)
        print(result)

        assert "response" in result
        assert len(result["response"]) > 0