litellm/tests/integration/management/test_managed_file_create_ownership.py

from __future__ import annotations

import base64
import json
import uuid
from dataclasses import dataclass
from typing import Final

import httpx
import pytest
from integration._support.client import JSON_OBJECT, Gateway, Scenario, string_value
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
from pydantic import JsonValue

_FILE: Final = {
    "id": "file-in-$REQUEST_ID",
    "object": "file",
    "purpose": "batch",
    "bytes": 100,
    "created_at": 1,
    "filename": "input.jsonl",
    "status": "processed",
}
_BATCH: Final = {
    "id": "batch-$REQUEST_ID",
    "object": "batch",
    "endpoint": "/v1/chat/completions",
    "errors": None,
    "input_file_id": "file-in-$REQUEST_ID",
    "completion_window": "24h",
    "status": "validating",
    "output_file_id": None,
    "error_file_id": None,
    "created_at": 1,
    "expires_at": 1,
    "request_counts": {"total": 1, "completed": 0, "failed": 0},
    "metadata": None,
}
_JOB: Final = {
    "id": "ftjob-$REQUEST_ID",
    "object": "fine_tuning.job",
    "model": "gpt-4o-mini-2024-07-18",
    "created_at": 1,
    "fine_tuned_model": None,
    "organization_id": "org-integration",
    "result_files": [],
    "status": "validating_files",
    "validation_file": None,
    "training_file": "file-in-$REQUEST_ID",
    "hyperparameters": {"n_epochs": "auto"},
    "seed": 1,
}


@dataclass(frozen=True, slots=True)
class _Tenants:
    model: str
    handle: ScenarioHandle
    owner_key: str
    teammate_key: str
    outsider_key: str
    owner_file_id: str


def _routes() -> RoutedResponse:
    return RoutedResponse(
        content_type="application/x-routed",
        routes={
            "POST /files": JsonResponse(content_type="application/json", body=_FILE),
            "POST /batches": JsonResponse(content_type="application/json", body=_BATCH),
            "POST /fine_tuning/jobs": JsonResponse(content_type="application/json", body=_JOB),
        },
    )


def _upload(scenario: Scenario, key: str, model: str, purpose: str) -> str:
    line: Final = json.dumps(
        {
            "custom_id": "req-1",
            "method": "POST",
            "url": "/v1/chat/completions",
            "body": {"model": model, "messages": [{"role": "user", "content": "ping"}]},
        }
    ).encode()
    uploaded: Final = scenario.gateway.request_multipart(
        "/v1/files",
        {"purpose": purpose, "target_model_names": model},
        {"file": ("input.jsonl", line + b"\n", "application/jsonl")},
        key=key,
    )
    assert uploaded.status_code == 200, uploaded.text
    return string_value(JSON_OBJECT.validate_json(uploaded.content)["id"])


def _tenants(scenario: Scenario, purpose: str) -> _Tenants:
    handle: Final = register_scenario(f"managed-file-create-ownership-{uuid.uuid4().hex}", _routes())
    scenario.cleanups.callback(delete_scenario, handle)
    model: Final = scenario.model(api_base=handle.api_base())
    team_a: Final = scenario.team(models=[model])
    team_b: Final = scenario.team(models=[model])
    owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model])
    teammate_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a, models=[model])
    outsider_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, models=[model])
    return _Tenants(model, handle, owner_key, teammate_key, outsider_key, _upload(scenario, owner_key, model, purpose))


def _upstream_creates(gateway: Gateway, handle: ScenarioHandle, path: str) -> int:
    response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False)
    assert response.status_code == 200, response.text
    requests: Final = JSON_OBJECT.validate_json(response.content)["requests"]
    assert isinstance(requests, list), response.text
    return sum(
        1
        for request in requests
        if isinstance(request, dict)
        and request.get("method") == "POST"
        and handle.scenario_id in str(request.get("path"))
        and str(request.get("path")).endswith(path)
    )


def _create_batch(gateway: Gateway, key: str, model: str, input_file_id: JsonValue) -> httpx.Response:
    return gateway.request(
        "POST",
        "/v1/batches",
        {
            "input_file_id": input_file_id,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "model": model,
        },
        key=key,
    )


def _create_job(gateway: Gateway, key: str, model: str, files: dict[str, JsonValue]) -> httpx.Response:
    return gateway.request("POST", "/v1/fine_tuning/jobs", {"model": model, **files}, key=key)


def test_batch_create_from_another_teams_managed_file_is_403_and_never_reaches_the_provider(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        tenants: Final = _tenants(scenario, "batch")
        denied: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, tenants.owner_file_id)
        assert denied.status_code == 403, denied.text
        assert tenants.owner_file_id in denied.text, denied.text
        assert _upstream_creates(gateway, tenants.handle, "/batches") == 0


def test_batch_create_from_own_or_teammates_managed_file_still_dispatches(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        tenants: Final = _tenants(scenario, "batch")
        owner: Final = _create_batch(gateway, tenants.owner_key, tenants.model, tenants.owner_file_id)
        teammate: Final = _create_batch(gateway, tenants.teammate_key, tenants.model, tenants.owner_file_id)
        admin: Final = _create_batch(gateway, gateway.key, tenants.model, tenants.owner_file_id)
        assert (owner.status_code, teammate.status_code, admin.status_code) == (200, 200, 200), (
            owner.text,
            teammate.text,
            admin.text,
        )
        assert JSON_OBJECT.validate_json(owner.content)["input_file_id"] == tenants.owner_file_id, owner.text
        assert _upstream_creates(gateway, tenants.handle, "/batches") == 3


def test_batch_create_from_a_managed_file_id_with_no_owner_row_is_404(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        tenants: Final = _tenants(scenario, "batch")
        forged_raw: Final = (
            f"litellm_proxy:application/octet-stream;unified_id,{uuid.uuid4()};target_model_names,{tenants.model};"
            "llm_output_file_id,file-forged;llm_output_file_model_id,model-forged"
        )
        forged: Final = base64.urlsafe_b64encode(forged_raw.encode()).decode().rstrip("=")
        missing: Final = _create_batch(gateway, tenants.outsider_key, tenants.model, forged)
        assert missing.status_code == 404, missing.text
        assert _upstream_creates(gateway, tenants.handle, "/batches") == 0


@pytest.mark.parametrize("field", ["training_file", "validation_file"])
def test_fine_tuning_create_from_another_teams_managed_file_is_403(gateway: Gateway, field: str) -> None:
    with gateway.scenario() as scenario:
        tenants: Final = _tenants(scenario, "fine-tune")
        own_training: Final = _upload(scenario, tenants.outsider_key, tenants.model, "fine-tune")
        files: Final[dict[str, JsonValue]] = {"training_file": own_training, field: tenants.owner_file_id}
        denied: Final = _create_job(gateway, tenants.outsider_key, tenants.model, files)
        assert denied.status_code == 403, denied.text
        assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 0


def test_fine_tuning_create_from_own_managed_files_still_dispatches(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        tenants: Final = _tenants(scenario, "fine-tune")
        validation: Final = _upload(scenario, tenants.teammate_key, tenants.model, "fine-tune")
        created: Final = _create_job(
            gateway,
            tenants.owner_key,
            tenants.model,
            {"training_file": tenants.owner_file_id, "validation_file": validation},
        )
        assert created.status_code == 200, created.text
        assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 1


def _upstream_calls(gateway: Gateway, handle: ScenarioHandle, method: str, suffix: str) -> int:
    response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False)
    assert response.status_code == 200, response.text
    requests: Final = JSON_OBJECT.validate_json(response.content)["requests"]
    assert isinstance(requests, list), response.text
    return sum(
        1
        for request in requests
        if isinstance(request, dict)
        and request.get("method") == method
        and handle.scenario_id in str(request.get("path"))
        and str(request.get("path")).endswith(suffix)
    )


def test_batch_create_from_another_teams_file_is_403_before_the_rate_limiter_reads_it(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        scenario_id: Final = f"managed-file-create-ownership-{uuid.uuid4().hex}"
        handle: Final = register_scenario(scenario_id, _routes())
        scenario.cleanups.callback(delete_scenario, handle)
        model: Final = scenario.model(api_base=handle.api_base())
        line: Final = json.dumps(
            {
                "custom_id": "req-1",
                "method": "POST",
                "url": "/v1/chat/completions",
                "body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 5},
            }
        )
        register_scenario(
            scenario_id,
            RoutedResponse(
                content_type="application/x-routed",
                routes={
                    **_routes().routes,
                    "GET /files/file-in-$REQUEST_ID/content": TextResponse(
                        content_type="application/jsonl", body=line + "\n"
                    ),
                },
            ),
        )
        team_a: Final = scenario.team(models=[model])
        team_b: Final = scenario.team(models=[model])
        owner: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a)
        outsider: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, tpm_limit=1)
        owner_file_id: Final = _upload(scenario, owner, model, "batch")
        denied: Final = _create_batch(gateway, outsider, model, owner_file_id)
        assert denied.status_code == 403, denied.text
        assert owner_file_id in denied.text, denied.text
        assert _upstream_calls(gateway, handle, "GET", "/content") == 0


def test_regenerated_key_without_user_or_team_keeps_its_managed_files(gateway: Gateway) -> None:
    with gateway.scenario() as scenario:
        tenants: Final = _tenants(scenario, "batch")
        original: Final = string_value(gateway.post("/key/generate", {})["key"])
        file_id: Final = _upload(scenario, original, tenants.model, "batch")
        regenerated: Final = string_value(gateway.post("/key/regenerate", {"key": original})["key"])
        scenario.cleanups.callback(scenario.delete_key, regenerated)
        batch: Final = _create_batch(gateway, regenerated, tenants.model, file_id)
        retrieved: Final = gateway.request("GET", f"/v1/files/{file_id}", key=regenerated)
        stranger: Final = _create_batch(
            gateway, string_value(gateway.post("/key/generate", {})["key"]), tenants.model, file_id
        )
        assert (batch.status_code, retrieved.status_code, stranger.status_code) == (200, 200, 403), (
            batch.text,
            retrieved.text,
            stranger.text,
        )