diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 3a2d9e13f77..45de410494b 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -27,6 +27,7 @@ from litellm.llms.azure.batches.handler import AzureBatchesAPI from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.moonshot.batches.handler import MoonshotBatchesAPI from litellm.llms.openai.openai import OpenAIBatchesAPI from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction from litellm.secret_managers.main import get_secret_str @@ -57,6 +58,7 @@ openai_batches_instance = OpenAIBatchesAPI() azure_batches_instance = AzureBatchesAPI() vertex_ai_batches_instance = VertexAIBatchPrediction(gcs_bucket_name="") anthropic_batches_instance = AnthropicBatchesHandler() +moonshot_batches_instance = MoonshotBatchesAPI() base_llm_http_handler = BaseLLMHTTPHandler() ################################################# @@ -105,7 +107,7 @@ async def acreate_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "moonshot"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -155,7 +157,7 @@ def create_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "moonshot"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -318,6 +320,18 @@ def create_batch( max_retries=optional_params.max_retries, create_batch_data=_create_batch_request, ) + elif custom_llm_provider == "moonshot": + api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + api_key = optional_params.api_key or litellm.api_key or get_secret_str("MOONSHOT_API_KEY") + + response = moonshot_batches_instance.create_batch( + _is_async=_is_async, + api_base=api_base, + api_key=api_key, + timeout=timeout, + max_retries=optional_params.max_retries, + create_batch_data=_create_batch_request, + ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support custom_llm_provider={} for 'create_batch'".format(custom_llm_provider), @@ -337,7 +351,7 @@ def create_batch( @client async def aretrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -383,7 +397,7 @@ def _handle_retrieve_batch_providers_without_provider_config( litellm_params: dict, _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot"] = "openai", logging_obj: Optional[Any] = None, ): api_base: Optional[str] = None @@ -485,11 +499,23 @@ def _handle_retrieve_batch_providers_without_provider_config( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "moonshot": + api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + api_key = optional_params.api_key or litellm.api_key or get_secret_str("MOONSHOT_API_KEY") + + response = moonshot_batches_instance.retrieve_batch( + _is_async=_is_async, + retrieve_batch_data=_retrieve_batch_request, + api_base=api_base, + api_key=api_key, + timeout=timeout, + max_retries=optional_params.max_retries, + ) else: raise litellm.exceptions.BadRequestError( message=( "LiteLLM doesn't support custom_llm_provider={} for 'retrieve_batch' without a `model` kwarg. " - "Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. " + "Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic', 'moonshot'. " "'bedrock' is supported but requires `model` to be passed so the provider config can be loaded." ).format(custom_llm_provider), model="n/a", @@ -506,7 +532,7 @@ def _handle_retrieve_batch_providers_without_provider_config( @client def retrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "moonshot"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -801,6 +827,19 @@ def list_batches( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "moonshot": + api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + api_key = optional_params.api_key or litellm.api_key or get_secret_str("MOONSHOT_API_KEY") + + response = moonshot_batches_instance.list_batches( + _is_async=_is_async, + after=after, + limit=limit, + api_base=api_base, + api_key=api_key, + timeout=timeout, + max_retries=optional_params.max_retries, + ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: {}.".format( @@ -823,7 +862,7 @@ def list_batches( async def acancel_batch( batch_id: str, model: Optional[str] = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "moonshot"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -869,7 +908,7 @@ async def acancel_batch( def cancel_batch( batch_id: str, model: Optional[str] = None, - custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai", + custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai", "moonshot"], str] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -990,9 +1029,21 @@ def cancel_batch( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "moonshot": + api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + api_key = optional_params.api_key or litellm.api_key or get_secret_str("MOONSHOT_API_KEY") + + response = moonshot_batches_instance.cancel_batch( + _is_async=_is_async, + cancel_batch_data=_cancel_batch_request, + api_base=api_base, + api_key=api_key, + timeout=timeout, + max_retries=optional_params.max_retries, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format( + message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', 'vertex_ai', and 'moonshot' are supported.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/llms/moonshot/batches/handler.py b/litellm/llms/moonshot/batches/handler.py new file mode 100644 index 00000000000..0390a51da48 --- /dev/null +++ b/litellm/llms/moonshot/batches/handler.py @@ -0,0 +1,250 @@ +""" +Moonshot AI (Kimi) Batches API Handler + +Moonshot's batch API is OpenAI-compatible; this handler wires the OpenAI SDK +to Moonshot's endpoint (https://api.moonshot.ai/v1) using the caller-supplied +or environment-sourced API key / base URL. +""" + +from typing import Any, Coroutine, Optional, Union, cast + +import httpx +from openai import AsyncOpenAI, OpenAI + +from litellm.llms.base import BaseLLM +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + CancelBatchRequest, + CreateBatchRequest, + RetrieveBatchRequest, +) +from litellm.types.utils import LiteLLMBatch + +MOONSHOT_DEFAULT_API_BASE = "https://api.moonshot.ai/v1" + + +class MoonshotBatchesAPI(BaseLLM): + """ + Moonshot AI batch operations — create, retrieve, cancel, list. + All four map 1-to-1 to the OpenAI SDK because the Kimi batch API is + OpenAI-compatible. + """ + + def __init__(self) -> None: + super().__init__() + + def _get_client( + self, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + _is_async: bool, + client: Optional[Union[OpenAI, AsyncOpenAI]] = None, + ) -> Optional[Union[OpenAI, AsyncOpenAI]]: + if client is not None: + return client + + resolved_key = api_key or get_secret_str("MOONSHOT_API_KEY") + resolved_base = ( + api_base + or get_secret_str("MOONSHOT_API_BASE") + or MOONSHOT_DEFAULT_API_BASE + ) + + kwargs: dict = {"base_url": resolved_base} + if resolved_key is not None: + kwargs["api_key"] = resolved_key + if max_retries is not None: + kwargs["max_retries"] = max_retries + if not isinstance(timeout, httpx.Timeout): + kwargs["timeout"] = float(timeout) + else: + kwargs["timeout"] = timeout + + if _is_async: + return AsyncOpenAI(**kwargs) + return OpenAI(**kwargs) + + # ------------------------------------------------------------------ create + + async def acreate_batch( + self, + create_batch_data: CreateBatchRequest, + client: AsyncOpenAI, + ) -> LiteLLMBatch: + response = await client.batches.create(**create_batch_data) # type: ignore[arg-type] + return LiteLLMBatch(**response.model_dump()) + + def create_batch( + self, + _is_async: bool, + create_batch_data: CreateBatchRequest, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + client: Optional[Union[OpenAI, AsyncOpenAI]] = None, + ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: + moonshot_client = self._get_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + _is_async=_is_async, + client=client, + ) + if moonshot_client is None: + raise ValueError( + "Moonshot client could not be initialised. " + "Pass api_key or set MOONSHOT_API_KEY in the environment." + ) + + if _is_async: + if not isinstance(moonshot_client, AsyncOpenAI): + raise ValueError( + "Expected an AsyncOpenAI client for async create_batch." + ) + return self.acreate_batch( + create_batch_data=create_batch_data, client=moonshot_client + ) + + response = cast(OpenAI, moonshot_client).batches.create(**create_batch_data) # type: ignore[arg-type] + return LiteLLMBatch(**response.model_dump()) + + # ---------------------------------------------------------------- retrieve + + async def aretrieve_batch( + self, + retrieve_batch_data: RetrieveBatchRequest, + client: AsyncOpenAI, + ) -> LiteLLMBatch: + response = await client.batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type] + return LiteLLMBatch(**response.model_dump()) + + def retrieve_batch( + self, + _is_async: bool, + retrieve_batch_data: RetrieveBatchRequest, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + client: Optional[Union[OpenAI, AsyncOpenAI]] = None, + ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: + moonshot_client = self._get_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + _is_async=_is_async, + client=client, + ) + if moonshot_client is None: + raise ValueError( + "Moonshot client could not be initialised. " + "Pass api_key or set MOONSHOT_API_KEY in the environment." + ) + + if _is_async: + if not isinstance(moonshot_client, AsyncOpenAI): + raise ValueError( + "Expected an AsyncOpenAI client for async retrieve_batch." + ) + return self.aretrieve_batch( + retrieve_batch_data=retrieve_batch_data, client=moonshot_client + ) + + response = cast(OpenAI, moonshot_client).batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type] + return LiteLLMBatch(**response.model_dump()) + + # ------------------------------------------------------------------ cancel + + async def acancel_batch( + self, + cancel_batch_data: CancelBatchRequest, + client: AsyncOpenAI, + ) -> LiteLLMBatch: + response = await client.batches.cancel(**cancel_batch_data) + return LiteLLMBatch(**response.model_dump()) + + def cancel_batch( + self, + _is_async: bool, + cancel_batch_data: CancelBatchRequest, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + client: Optional[Union[OpenAI, AsyncOpenAI]] = None, + ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: + moonshot_client = self._get_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + _is_async=_is_async, + client=client, + ) + if moonshot_client is None: + raise ValueError( + "Moonshot client could not be initialised. " + "Pass api_key or set MOONSHOT_API_KEY in the environment." + ) + + if _is_async: + if not isinstance(moonshot_client, AsyncOpenAI): + raise ValueError( + "Expected an AsyncOpenAI client for async cancel_batch." + ) + return self.acancel_batch( + cancel_batch_data=cancel_batch_data, client=moonshot_client + ) + + response = cast(OpenAI, moonshot_client).batches.cancel(**cancel_batch_data) + return LiteLLMBatch(**response.model_dump()) + + # -------------------------------------------------------------------- list + + async def alist_batches( + self, + client: AsyncOpenAI, + after: Optional[str] = None, + limit: Optional[int] = None, + ): + response = await client.batches.list(after=after, limit=limit) # type: ignore + return response + + def list_batches( + self, + _is_async: bool, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + after: Optional[str] = None, + limit: Optional[int] = None, + client: Optional[Union[OpenAI, AsyncOpenAI]] = None, + ): + moonshot_client = self._get_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + _is_async=_is_async, + client=client, + ) + if moonshot_client is None: + raise ValueError( + "Moonshot client could not be initialised. " + "Pass api_key or set MOONSHOT_API_KEY in the environment." + ) + + if _is_async: + if not isinstance(moonshot_client, AsyncOpenAI): + raise ValueError( + "Expected an AsyncOpenAI client for async list_batches." + ) + return self.alist_batches(client=moonshot_client, after=after, limit=limit) + + return cast(OpenAI, moonshot_client).batches.list(after=after, limit=limit) # type: ignore diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 279e9b15fe7..3d4fbd3a723 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3378,7 +3378,7 @@ OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = { LlmProviders.HOSTED_VLLM.value, } -ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai"] +ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai", "moonshot"] LIST_BATCHES_SUPPORTED_PROVIDERS: frozenset[str] = frozenset(get_args(ListBatchesSupportedProvider)) diff --git a/tests/test_litellm/llms/moonshot/batches/test_handler.py b/tests/test_litellm/llms/moonshot/batches/test_handler.py new file mode 100644 index 00000000000..74dadb945fd --- /dev/null +++ b/tests/test_litellm/llms/moonshot/batches/test_handler.py @@ -0,0 +1,326 @@ +"""Unit tests for MoonshotBatchesAPI (litellm/llms/moonshot/batches/handler.py). + +The handler is a thin credential-resolution + OpenAI-SDK-delegation layer. +We mock the two seams: + * ``MoonshotBatchesAPI._get_client`` – credential / client-construction. + * the returned client's ``batches.*`` methods – the network call. + +Pure logic (the _is_async branch, isinstance guards, model_dump parse) runs for +real. +""" + +from __future__ import annotations + +import asyncio +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from openai import AsyncOpenAI, OpenAI # noqa: E402 + +from litellm.llms.moonshot.batches.handler import MoonshotBatchesAPI # noqa: E402 +from litellm.types.utils import LiteLLMBatch # noqa: E402 + +GET_CLIENT = "litellm.llms.moonshot.batches.handler.MoonshotBatchesAPI._get_client" + +AUTH_KW = dict( + api_key="sk-moonshot-test", + api_base="https://api.moonshot.ai/v1", + timeout=600.0, + max_retries=3, +) + +CREATE_DATA = { + "completion_window": "24h", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-abc", +} +RETRIEVE_DATA = {"batch_id": "batch-123"} +CANCEL_DATA = {"batch_id": "batch-123"} + + +def _batch_dict(batch_id: str = "batch-123", status: str = "completed") -> dict: + return { + "id": batch_id, + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": "file-abc", + "object": "batch", + "status": status, + } + + +def _sdk_response(batch_dict: dict) -> MagicMock: + resp = MagicMock() + resp.model_dump.return_value = batch_dict + return resp + + +def _sync_client() -> MagicMock: + return MagicMock(spec=OpenAI) + + +def _async_client() -> MagicMock: + return MagicMock(spec=AsyncOpenAI) + + +# ====================================================================== create + + +class TestCreateBatch: + def test_sync_create_batch(self): + handler = MoonshotBatchesAPI() + sdk = _sync_client() + sdk.batches.create.return_value = _sdk_response(_batch_dict()) + + with patch(GET_CLIENT, return_value=sdk) as mock_get: + result = handler.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + **AUTH_KW, + ) + + mock_get.assert_called_once_with( + api_key=AUTH_KW["api_key"], + api_base=AUTH_KW["api_base"], + timeout=AUTH_KW["timeout"], + max_retries=AUTH_KW["max_retries"], + _is_async=False, + client=None, + ) + sdk.batches.create.assert_called_once_with(**CREATE_DATA) + assert isinstance(result, LiteLLMBatch) + assert result.id == "batch-123" + + def test_async_create_batch(self): + handler = MoonshotBatchesAPI() + sdk = _async_client() + sdk.batches.create = AsyncMock(return_value=_sdk_response(_batch_dict())) + + with patch(GET_CLIENT, return_value=sdk): + coro = handler.create_batch( + _is_async=True, + create_batch_data=CREATE_DATA, + **AUTH_KW, + ) + result = asyncio.get_event_loop().run_until_complete(coro) + + sdk.batches.create.assert_called_once_with(**CREATE_DATA) + assert isinstance(result, LiteLLMBatch) + + def test_sync_raises_when_client_none(self): + handler = MoonshotBatchesAPI() + with patch(GET_CLIENT, return_value=None): + with pytest.raises(ValueError, match="Moonshot client could not be initialised"): + handler.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + **AUTH_KW, + ) + + +# ==================================================================== retrieve + + +class TestRetrieveBatch: + def test_sync_retrieve_batch(self): + handler = MoonshotBatchesAPI() + sdk = _sync_client() + sdk.batches.retrieve.return_value = _sdk_response(_batch_dict()) + + with patch(GET_CLIENT, return_value=sdk): + result = handler.retrieve_batch( + _is_async=False, + retrieve_batch_data=RETRIEVE_DATA, + **AUTH_KW, + ) + + sdk.batches.retrieve.assert_called_once_with(**RETRIEVE_DATA) + assert isinstance(result, LiteLLMBatch) + + def test_async_retrieve_batch(self): + handler = MoonshotBatchesAPI() + sdk = _async_client() + sdk.batches.retrieve = AsyncMock(return_value=_sdk_response(_batch_dict())) + + with patch(GET_CLIENT, return_value=sdk): + coro = handler.retrieve_batch( + _is_async=True, + retrieve_batch_data=RETRIEVE_DATA, + **AUTH_KW, + ) + result = asyncio.get_event_loop().run_until_complete(coro) + + sdk.batches.retrieve.assert_called_once_with(**RETRIEVE_DATA) + assert isinstance(result, LiteLLMBatch) + + def test_sync_raises_when_client_none(self): + handler = MoonshotBatchesAPI() + with patch(GET_CLIENT, return_value=None): + with pytest.raises(ValueError, match="Moonshot client could not be initialised"): + handler.retrieve_batch( + _is_async=False, + retrieve_batch_data=RETRIEVE_DATA, + **AUTH_KW, + ) + + +# ====================================================================== cancel + + +class TestCancelBatch: + def test_sync_cancel_batch(self): + handler = MoonshotBatchesAPI() + sdk = _sync_client() + sdk.batches.cancel.return_value = _sdk_response(_batch_dict(status="cancelling")) + + with patch(GET_CLIENT, return_value=sdk): + result = handler.cancel_batch( + _is_async=False, + cancel_batch_data=CANCEL_DATA, + **AUTH_KW, + ) + + sdk.batches.cancel.assert_called_once_with(**CANCEL_DATA) + assert isinstance(result, LiteLLMBatch) + + def test_async_cancel_batch(self): + handler = MoonshotBatchesAPI() + sdk = _async_client() + sdk.batches.cancel = AsyncMock(return_value=_sdk_response(_batch_dict(status="cancelling"))) + + with patch(GET_CLIENT, return_value=sdk): + coro = handler.cancel_batch( + _is_async=True, + cancel_batch_data=CANCEL_DATA, + **AUTH_KW, + ) + result = asyncio.get_event_loop().run_until_complete(coro) + + sdk.batches.cancel.assert_called_once_with(**CANCEL_DATA) + assert isinstance(result, LiteLLMBatch) + + def test_sync_raises_when_client_none(self): + handler = MoonshotBatchesAPI() + with patch(GET_CLIENT, return_value=None): + with pytest.raises(ValueError, match="Moonshot client could not be initialised"): + handler.cancel_batch( + _is_async=False, + cancel_batch_data=CANCEL_DATA, + **AUTH_KW, + ) + + +# ======================================================================== list + + +class TestListBatches: + def test_sync_list_batches(self): + handler = MoonshotBatchesAPI() + sdk = _sync_client() + list_resp = MagicMock() + sdk.batches.list.return_value = list_resp + + with patch(GET_CLIENT, return_value=sdk): + result = handler.list_batches( + _is_async=False, + after=None, + limit=10, + **AUTH_KW, + ) + + sdk.batches.list.assert_called_once_with(after=None, limit=10) + assert result is list_resp + + def test_async_list_batches(self): + handler = MoonshotBatchesAPI() + sdk = _async_client() + list_resp = MagicMock() + sdk.batches.list = AsyncMock(return_value=list_resp) + + with patch(GET_CLIENT, return_value=sdk): + coro = handler.list_batches( + _is_async=True, + after="cursor-abc", + limit=5, + **AUTH_KW, + ) + result = asyncio.get_event_loop().run_until_complete(coro) + + sdk.batches.list.assert_called_once_with(after="cursor-abc", limit=5) + assert result is list_resp + + def test_sync_raises_when_client_none(self): + handler = MoonshotBatchesAPI() + with patch(GET_CLIENT, return_value=None): + with pytest.raises(ValueError, match="Moonshot client could not be initialised"): + handler.list_batches( + _is_async=False, + **AUTH_KW, + ) + + +# ============================================================= _get_client unit + + +class TestGetClient: + def test_uses_env_key_when_no_api_key(self, monkeypatch): + monkeypatch.setenv("MOONSHOT_API_KEY", "env-key") + handler = MoonshotBatchesAPI() + with patch("litellm.llms.moonshot.batches.handler.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock(spec=OpenAI) + client = handler._get_client( + api_key=None, + api_base="https://api.moonshot.ai/v1", + timeout=30.0, + max_retries=2, + _is_async=False, + ) + mock_cls.assert_called_once() + call_kwargs = mock_cls.call_args.kwargs + assert call_kwargs["api_key"] == "env-key" + + def test_defaults_to_moonshot_base_url(self): + handler = MoonshotBatchesAPI() + with patch("litellm.llms.moonshot.batches.handler.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock(spec=OpenAI) + handler._get_client( + api_key="sk-test", + api_base=None, + timeout=30.0, + max_retries=None, + _is_async=False, + ) + call_kwargs = mock_cls.call_args.kwargs + assert call_kwargs["base_url"] == "https://api.moonshot.ai/v1" + + def test_returns_async_client_when_is_async(self): + handler = MoonshotBatchesAPI() + with patch("litellm.llms.moonshot.batches.handler.AsyncOpenAI") as mock_cls: + mock_cls.return_value = MagicMock(spec=AsyncOpenAI) + client = handler._get_client( + api_key="sk-test", + api_base="https://api.moonshot.ai/v1", + timeout=30.0, + max_retries=None, + _is_async=True, + ) + mock_cls.assert_called_once() + + def test_returns_provided_client_unchanged(self): + handler = MoonshotBatchesAPI() + existing = MagicMock(spec=OpenAI) + result = handler._get_client( + api_key="sk-test", + api_base="https://api.moonshot.ai/v1", + timeout=30.0, + max_retries=None, + _is_async=False, + client=existing, + ) + assert result is existing