diff --git a/examples/async_jobs_chat.py b/examples/async_jobs_chat.py index 15e550c8..e5019148 100644 --- a/examples/async_jobs_chat.py +++ b/examples/async_jobs_chat.py @@ -54,6 +54,9 @@ async def main(): await client.files.delete(training_file.id) await client.files.delete(validation_file.id) + # Delete fine-tuned model + await client.delete_model(created_job.fine_tuned_model) + if __name__ == "__main__": asyncio.run(main()) diff --git a/pyproject.toml b/pyproject.toml index 68cfc961..4f498a6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "mistralai" -version = "0.4.0" +version = "0.4.1" description = "" authors = ["Bam4d "] readme = "README.md" diff --git a/src/mistralai/async_client.py b/src/mistralai/async_client.py index 4e80de84..24124486 100644 --- a/src/mistralai/async_client.py +++ b/src/mistralai/async_client.py @@ -29,7 +29,7 @@ ToolChoice, ) from mistralai.models.embeddings import EmbeddingResponse -from mistralai.models.models import ModelList +from mistralai.models.models import ModelDeleted, ModelList class MistralAsyncClient(ClientBase): @@ -304,6 +304,14 @@ async def list_models(self) -> ModelList: raise MistralException("No response received") + async def delete_model(self, model_id: str) -> ModelDeleted: + single_response = self._request("delete", {}, f"v1/models/{model_id}") + + async for response in single_response: + return ModelDeleted(**response) + + raise MistralException("No response received") + async def completion( self, model: str, diff --git a/src/mistralai/client.py b/src/mistralai/client.py index 1f70f4ea..c919ac71 100644 --- a/src/mistralai/client.py +++ b/src/mistralai/client.py @@ -22,7 +22,7 @@ ToolChoice, ) from mistralai.models.embeddings import EmbeddingResponse -from mistralai.models.models import ModelList +from mistralai.models.models import ModelDeleted, ModelList class MistralClient(ClientBase): @@ -298,6 +298,14 @@ def list_models(self) -> ModelList: raise MistralException("No response received") + def delete_model(self, model_id: str) -> ModelDeleted: + single_response = self._request("delete", {}, f"v1/models/{model_id}") + + for response in single_response: + return ModelDeleted(**response) + + raise MistralException("No response received") + def completion( self, model: str, diff --git a/src/mistralai/client_base.py b/src/mistralai/client_base.py index 9fb107b6..5cd996d9 100644 --- a/src/mistralai/client_base.py +++ b/src/mistralai/client_base.py @@ -10,7 +10,7 @@ ) from mistralai.models.chat_completion import ChatMessage, Function, ResponseFormat, ToolChoice -CLIENT_VERSION = "0.4.0" +CLIENT_VERSION = "0.4.1" class ClientBase(ABC): diff --git a/src/mistralai/models/models.py b/src/mistralai/models/models.py index 8ed152d3..f88033d4 100644 --- a/src/mistralai/models/models.py +++ b/src/mistralai/models/models.py @@ -31,3 +31,9 @@ class ModelCard(BaseModel): class ModelList(BaseModel): object: str data: List[ModelCard] + + +class ModelDeleted(BaseModel): + id: str + object: str + deleted: bool diff --git a/tests/test_delete_model.py b/tests/test_delete_model.py new file mode 100644 index 00000000..d050c21a --- /dev/null +++ b/tests/test_delete_model.py @@ -0,0 +1,26 @@ +from mistralai.models.models import ModelDeleted + +from .utils import mock_model_deleted_response_payload, mock_response + + +class TestDeleteModel: + def test_delete_model(self, client): + expected_response_model = ModelDeleted.model_validate_json(mock_model_deleted_response_payload()) + client._client.request.return_value = mock_response(200, expected_response_model.json()) + + response_model = client.delete_model("model_id") + + client._client.request.assert_called_once_with( + "delete", + "https://api.mistral.ai/v1/models/model_id", + headers={ + "User-Agent": f"mistral-client-python/{client._version}", + "Accept": "application/json", + "Content-Type": "application/json", + "Authorization": "Bearer test_api_key", + }, + json={}, + data=None, + ) + + assert response_model == expected_response_model diff --git a/tests/test_delete_model_async.py b/tests/test_delete_model_async.py new file mode 100644 index 00000000..9fa393e8 --- /dev/null +++ b/tests/test_delete_model_async.py @@ -0,0 +1,28 @@ +import pytest +from mistralai.models.models import ModelDeleted + +from .utils import mock_model_deleted_response_payload, mock_response + + +class TestAsyncDeleteModel: + @pytest.mark.asyncio + async def test_delete_model(self, async_client): + expected_response_model = ModelDeleted.model_validate_json(mock_model_deleted_response_payload()) + async_client._client.request.return_value = mock_response(200, expected_response_model.json()) + + response_model = await async_client.delete_model("model_id") + + async_client._client.request.assert_called_once_with( + "delete", + "https://api.mistral.ai/v1/models/model_id", + headers={ + "User-Agent": f"mistral-client-python/{async_client._version}", + "Accept": "application/json", + "Content-Type": "application/json", + "Authorization": "Bearer test_api_key", + }, + json={}, + data=None, + ) + + assert response_model == expected_response_model diff --git a/tests/utils.py b/tests/utils.py index 05fdc3df..5f7be152 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -318,3 +318,13 @@ def mock_file_deleted_response_payload() -> str: "deleted": True, } ) + + +def mock_model_deleted_response_payload() -> str: + return orjson.dumps( + { + "id": "model_id", + "object": "model", + "deleted": True, + } + )