refactor: refactor testing
This commit is contained in:
Krrish Dholakia
parent
6df995d0d1
commit
5cd8ca2365
@ -123,78 +123,6 @@ async def test_create_fine_tune_jobs_async():
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_create_fine_tune_jobs_async():
|
||||
try:
|
||||
verbose_logger.setLevel(logging.DEBUG)
|
||||
file_name = "azure_fine_tune.jsonl"
|
||||
_current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
file_path = os.path.join(_current_dir, file_name)
|
||||
|
||||
file_id = "file-5e4b20ecbd724182b9964f3cd2ab7212"
|
||||
|
||||
create_fine_tuning_response = await litellm.acreate_fine_tuning_job(
|
||||
model="gpt-35-turbo-1106",
|
||||
training_file=file_id,
|
||||
custom_llm_provider="azure",
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app",
|
||||
)
|
||||
|
||||
print(
|
||||
"response from litellm.create_fine_tuning_job=", create_fine_tuning_response
|
||||
)
|
||||
|
||||
assert create_fine_tuning_response.id is not None
|
||||
|
||||
# response from Example/mocked endpoint
|
||||
assert create_fine_tuning_response.model == "davinci-002"
|
||||
|
||||
# list fine tuning jobs
|
||||
print("listing ft jobs")
|
||||
ft_jobs = await litellm.alist_fine_tuning_jobs(
|
||||
limit=2,
|
||||
custom_llm_provider="azure",
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app",
|
||||
)
|
||||
print("response from litellm.list_fine_tuning_jobs=", ft_jobs)
|
||||
|
||||
# cancel ft job
|
||||
response = await litellm.acancel_fine_tuning_job(
|
||||
fine_tuning_job_id=create_fine_tuning_response.id,
|
||||
custom_llm_provider="azure",
|
||||
api_key=os.getenv("AZURE_SWEDEN_API_KEY"),
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app",
|
||||
)
|
||||
|
||||
print("response from litellm.cancel_fine_tuning_job=", response)
|
||||
|
||||
assert response.status == "cancelled"
|
||||
assert response.id == create_fine_tuning_response.id
|
||||
except openai.RateLimitError:
|
||||
pass
|
||||
except Exception as e:
|
||||
if "Job has already completed" in str(e):
|
||||
pass
|
||||
else:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
pass
|
||||
|
||||
|
||||
def test_azure_trainingtype_defaults_to_one():
|
||||
"""
|
||||
Azure requires trainingType in extra_body. When omitted, AzureOpenAIFineTuningAPI defaults it to 1.
|
||||
"""
|
||||
from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI
|
||||
|
||||
handler = AzureOpenAIFineTuningAPI()
|
||||
create_data = {"model": "gpt-4o-mini", "training_file": "file-test"}
|
||||
|
||||
handler._ensure_training_type(create_data)
|
||||
|
||||
assert "extra_body" in create_data
|
||||
assert create_data["extra_body"]["trainingType"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_create_vertex_fine_tune_jobs_mocked():
|
||||
load_vertex_ai_credentials()
|
||||
|
||||
@ -148,48 +148,6 @@ async def test_basic_rerank_together_ai(sync_mode):
|
||||
raise e
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.skip(reason="Skipping test due to Cohere RBAC issues")
|
||||
async def test_basic_rerank_azure_ai(sync_mode):
|
||||
import os
|
||||
|
||||
litellm.set_verbose = True
|
||||
|
||||
if sync_mode is True:
|
||||
response = litellm.rerank(
|
||||
model="azure_ai/Cohere-rerank-v3-multilingual-ko",
|
||||
query="hello",
|
||||
documents=["hello", "world"],
|
||||
top_n=3,
|
||||
api_key=os.getenv("AZURE_AI_COHERE_API_KEY"),
|
||||
api_base=os.getenv("AZURE_AI_COHERE_API_BASE"),
|
||||
)
|
||||
|
||||
print("re rank response: ", response)
|
||||
|
||||
assert response.id is not None
|
||||
assert response.results is not None
|
||||
|
||||
assert_response_shape(response, custom_llm_provider="together_ai")
|
||||
else:
|
||||
response = await litellm.arerank(
|
||||
model="azure_ai/Cohere-rerank-v3-multilingual-ko",
|
||||
query="hello",
|
||||
documents=["hello", "world"],
|
||||
top_n=3,
|
||||
api_key=os.getenv("AZURE_AI_COHERE_API_KEY"),
|
||||
api_base=os.getenv("AZURE_AI_COHERE_API_BASE"),
|
||||
)
|
||||
|
||||
print("async re rank response: ", response)
|
||||
|
||||
assert response.id is not None
|
||||
assert response.results is not None
|
||||
|
||||
assert_response_shape(response, custom_llm_provider="together_ai")
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
@pytest.mark.parametrize("version", ["v1", "v2"])
|
||||
async def test_rerank_custom_api_base(version):
|
||||
|
||||
@ -0,0 +1,20 @@
|
||||
{
|
||||
"id": "ftjob-azure-create-123",
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1735689600,
|
||||
"model": "davinci-002",
|
||||
"status": "cancelled",
|
||||
"fine_tuned_model": null,
|
||||
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
|
||||
"hyperparameters": {
|
||||
"n_epochs": 3,
|
||||
"batch_size": null,
|
||||
"learning_rate_multiplier": null
|
||||
},
|
||||
"organization_id": "",
|
||||
"result_files": [],
|
||||
"validation_file": null,
|
||||
"trained_tokens": null,
|
||||
"estimated_finish": null,
|
||||
"error": null
|
||||
}
|
||||
@ -0,0 +1,18 @@
|
||||
{
|
||||
"id": "ftjob-azure-create-123",
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1735689600,
|
||||
"model": "davinci-002",
|
||||
"status": "canceled",
|
||||
"fine_tuned_model": null,
|
||||
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
|
||||
"hyperparameters": {
|
||||
"n_epochs": 3
|
||||
},
|
||||
"organization_id": null,
|
||||
"result_files": null,
|
||||
"validation_file": null,
|
||||
"trained_tokens": null,
|
||||
"estimated_finish": null,
|
||||
"error": null
|
||||
}
|
||||
@ -0,0 +1,3 @@
|
||||
{
|
||||
"fine_tuning_job_id": "ftjob-azure-create-123"
|
||||
}
|
||||
@ -0,0 +1,20 @@
|
||||
{
|
||||
"id": "ftjob-azure-create-123",
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1735689600,
|
||||
"model": "davinci-002",
|
||||
"status": "running",
|
||||
"fine_tuned_model": null,
|
||||
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
|
||||
"hyperparameters": {
|
||||
"n_epochs": 3,
|
||||
"batch_size": null,
|
||||
"learning_rate_multiplier": null
|
||||
},
|
||||
"organization_id": "",
|
||||
"result_files": [],
|
||||
"validation_file": null,
|
||||
"trained_tokens": null,
|
||||
"estimated_finish": null,
|
||||
"error": null
|
||||
}
|
||||
@ -0,0 +1,18 @@
|
||||
{
|
||||
"id": "ftjob-azure-create-123",
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1735689600,
|
||||
"model": "davinci-002",
|
||||
"status": "running",
|
||||
"fine_tuned_model": null,
|
||||
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
|
||||
"hyperparameters": {
|
||||
"n_epochs": 3
|
||||
},
|
||||
"organization_id": null,
|
||||
"result_files": null,
|
||||
"validation_file": null,
|
||||
"trained_tokens": null,
|
||||
"estimated_finish": null,
|
||||
"error": null
|
||||
}
|
||||
@ -0,0 +1,8 @@
|
||||
{
|
||||
"model": "gpt-35-turbo-1106",
|
||||
"training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212",
|
||||
"hyperparameters": {},
|
||||
"extra_body": {
|
||||
"trainingType": 1
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,20 @@
|
||||
{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "ftjob-azure-create-123",
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1735689600,
|
||||
"model": "davinci-002",
|
||||
"status": "running"
|
||||
},
|
||||
{
|
||||
"id": "ftjob-azure-prev-000",
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1735603200,
|
||||
"model": "davinci-002",
|
||||
"status": "succeeded"
|
||||
}
|
||||
],
|
||||
"has_more": false
|
||||
}
|
||||
@ -0,0 +1,4 @@
|
||||
{
|
||||
"after": "ftjob-azure-prev-000",
|
||||
"limit": 2
|
||||
}
|
||||
150
tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py
Normal file
150
tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py
Normal file
@ -0,0 +1,150 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from openai import AsyncAzureOpenAI
|
||||
|
||||
import litellm
|
||||
from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI
|
||||
|
||||
|
||||
def _expected_dir() -> Path:
|
||||
return Path(__file__).resolve().parent.parent.parent / "expected_fine_tuning_api"
|
||||
|
||||
|
||||
def _load_json(file_name: str) -> dict:
|
||||
path = _expected_dir() / file_name
|
||||
assert path.exists(), f"Expected fixture file not found: {path}"
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
class _MockSDKResponse:
|
||||
def __init__(self, payload: dict):
|
||||
self._payload = payload
|
||||
|
||||
def model_dump(self) -> dict:
|
||||
return self._payload
|
||||
|
||||
|
||||
def _mock_azure_client(
|
||||
create_payload: dict | None = None,
|
||||
list_payload: dict | None = None,
|
||||
cancel_payload: dict | None = None,
|
||||
):
|
||||
client = AsyncAzureOpenAI(
|
||||
api_key="test-key",
|
||||
api_version="2024-10-21",
|
||||
azure_endpoint="https://exampleopenaiendpoint-production.up.railway.app",
|
||||
)
|
||||
client.fine_tuning.jobs.create = AsyncMock(
|
||||
return_value=(
|
||||
_MockSDKResponse(create_payload) if create_payload is not None else None
|
||||
)
|
||||
) # type: ignore[method-assign]
|
||||
client.fine_tuning.jobs.list = AsyncMock(
|
||||
return_value=list_payload
|
||||
) # type: ignore[method-assign]
|
||||
client.fine_tuning.jobs.cancel = AsyncMock(
|
||||
return_value=(
|
||||
_MockSDKResponse(cancel_payload) if cancel_payload is not None else None
|
||||
)
|
||||
) # type: ignore[method-assign]
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_acreate_fine_tuning_job_request_and_output_match_expected_json():
|
||||
expected_request = _load_json("azure_create_request.json")
|
||||
raw_response = _load_json("azure_create_raw_response.json")
|
||||
expected_output = _load_json("azure_create_expected_output.json")
|
||||
|
||||
mock_client = _mock_azure_client(create_payload=raw_response)
|
||||
|
||||
with patch.object(
|
||||
AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client
|
||||
):
|
||||
response = await litellm.acreate_fine_tuning_job(
|
||||
model="gpt-35-turbo-1106",
|
||||
training_file="file-5e4b20ecbd724182b9964f3cd2ab7212",
|
||||
custom_llm_provider="azure",
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app",
|
||||
api_key="test-key",
|
||||
api_version="2024-10-21",
|
||||
)
|
||||
|
||||
request_kwargs = mock_client.fine_tuning.jobs.create.call_args.kwargs
|
||||
assert request_kwargs == expected_request
|
||||
|
||||
response_dict = response.model_dump(exclude={"_hidden_params"})
|
||||
for key, expected_value in expected_output.items():
|
||||
assert key in response_dict, f"Missing key in response: {key}"
|
||||
assert response_dict[key] == expected_value
|
||||
|
||||
assert response.id is not None
|
||||
assert response.model == "davinci-002"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_alist_fine_tuning_jobs_request_matches_expected_json():
|
||||
expected_request = _load_json("azure_list_request.json")
|
||||
raw_list_response = _load_json("azure_list_raw_response.json")
|
||||
|
||||
mock_client = _mock_azure_client(list_payload=raw_list_response)
|
||||
|
||||
with patch.object(
|
||||
AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client
|
||||
):
|
||||
response = await litellm.alist_fine_tuning_jobs(
|
||||
after=expected_request["after"],
|
||||
limit=expected_request["limit"],
|
||||
custom_llm_provider="azure",
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app",
|
||||
api_key="test-key",
|
||||
api_version="2024-10-21",
|
||||
)
|
||||
|
||||
request_kwargs = mock_client.fine_tuning.jobs.list.call_args.kwargs
|
||||
assert request_kwargs == expected_request
|
||||
assert response == raw_list_response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_acancel_fine_tuning_job_request_and_output_match_expected_json():
|
||||
expected_request = _load_json("azure_cancel_request.json")
|
||||
raw_response = _load_json("azure_cancel_raw_response.json")
|
||||
expected_output = _load_json("azure_cancel_expected_output.json")
|
||||
|
||||
mock_client = _mock_azure_client(cancel_payload=raw_response)
|
||||
|
||||
with patch.object(
|
||||
AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client
|
||||
):
|
||||
response = await litellm.acancel_fine_tuning_job(
|
||||
fine_tuning_job_id=expected_request["fine_tuning_job_id"],
|
||||
custom_llm_provider="azure",
|
||||
api_base="https://exampleopenaiendpoint-production.up.railway.app",
|
||||
api_key="test-key",
|
||||
api_version="2024-10-21",
|
||||
)
|
||||
|
||||
request_kwargs = mock_client.fine_tuning.jobs.cancel.call_args.kwargs
|
||||
assert request_kwargs == expected_request
|
||||
|
||||
response_dict = response.model_dump(exclude={"_hidden_params"})
|
||||
for key, expected_value in expected_output.items():
|
||||
assert key in response_dict, f"Missing key in response: {key}"
|
||||
assert response_dict[key] == expected_value
|
||||
|
||||
assert response.status == "cancelled"
|
||||
|
||||
|
||||
def test_azure_trainingtype_defaults_to_one():
|
||||
handler = AzureOpenAIFineTuningAPI()
|
||||
create_data = {"model": "gpt-4o-mini", "training_file": "file-test"}
|
||||
|
||||
handler._ensure_training_type(create_data)
|
||||
|
||||
assert "extra_body" in create_data
|
||||
assert create_data["extra_body"]["trainingType"] == 1
|
||||
Loading…
Reference in New Issue
Block a user