From ad47fee1812a90952de3b3c879067c48170003be Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 17 Jun 2024 12:48:46 -0700 Subject: [PATCH] feat add text completion config for mistral text --- litellm/__init__.py | 2 +- litellm/llms/openai.py | 79 +++++++++++++++++++++++++++ litellm/main.py | 2 +- litellm/tests/test_text_completion.py | 40 +++++++++----- litellm/utils.py | 16 +++++- 5 files changed, 120 insertions(+), 19 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 02e2c1f280b..bcf764f835b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -404,7 +404,6 @@ openai_compatible_providers: List = [ "mistral", "groq", "codestral", - "text-completion-codestral", "deepseek", "deepinfra", "perplexity", @@ -796,6 +795,7 @@ from .llms.openai import ( OpenAIConfig, OpenAITextCompletionConfig, MistralConfig, + MistralTextCompletionConfig, MistralEmbeddingConfig, DeepInfraConfig, AzureAIStudioConfig, diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index 1f2b836c3ac..976d1f5668e 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -208,6 +208,85 @@ class MistralEmbeddingConfig: return optional_params +class MistralTextCompletionConfig: + """ + Reference: https://docs.mistral.ai/api/#operation/createFIMCompletion + """ + + suffix: Optional[str] = None + temperature: Optional[int] = None + top_p: Optional[float] = None + max_tokens: Optional[int] = None + min_tokens: Optional[int] = None + stream: Optional[bool] = None + random_seed: Optional[int] = None + stop: Optional[str] = None + + def __init__( + self, + suffix: Optional[str] = None, + temperature: Optional[int] = None, + top_p: Optional[float] = None, + max_tokens: Optional[int] = None, + min_tokens: Optional[int] = None, + stream: Optional[bool] = None, + random_seed: Optional[int] = None, + stop: Optional[str] = None, + ) -> None: + locals_ = locals().copy() + for key, value in locals_.items(): + if key != "self" and value is not None: + setattr(self.__class__, key, value) + + @classmethod + def get_config(cls): + return { + k: v + for k, v in cls.__dict__.items() + if not k.startswith("__") + and not isinstance( + v, + ( + types.FunctionType, + types.BuiltinFunctionType, + classmethod, + staticmethod, + ), + ) + and v is not None + } + + def get_supported_openai_params(self): + return [ + "suffix", + "temperature", + "top_p", + "max_tokens", + "stream", + "seed", + "stop", + ] + + def map_openai_params(self, non_default_params: dict, optional_params: dict): + for param, value in non_default_params.items(): + if param == "suffix": + optional_params["suffix"] = value + if param == "temperature": + optional_params["temperature"] = value + if param == "top_p": + optional_params["top_p"] = value + if param == "max_tokens": + optional_params["max_tokens"] = value + if param == "stream" and value == True: + optional_params["stream"] = value + if param == "stop": + optional_params["stop"] = value + if param == "seed": + optional_params["extra_body"] = {"random_seed": value} + + return optional_params + + class AzureAIStudioConfig: def get_required_params(self) -> List[ProviderField]: """For a given provider, return it's required fields with a description""" diff --git a/litellm/main.py b/litellm/main.py index 91d44833ecd..648802620b3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1049,7 +1049,6 @@ def completion( or custom_llm_provider == "perplexity" or custom_llm_provider == "groq" or custom_llm_provider == "codestral" - or custom_llm_provider == "text-completion-codestral" or custom_llm_provider == "deepseek" or custom_llm_provider == "anyscale" or custom_llm_provider == "mistral" @@ -3711,6 +3710,7 @@ def text_completion( custom_llm_provider == "openai" or custom_llm_provider == "azure" or custom_llm_provider == "azure_text" + or custom_llm_provider == "text-completion-codestral" or custom_llm_provider == "text-completion-openai" ) and isinstance(prompt, list) diff --git a/litellm/tests/test_text_completion.py b/litellm/tests/test_text_completion.py index 5edfb1935aa..3ec3954fdea 100644 --- a/litellm/tests/test_text_completion.py +++ b/litellm/tests/test_text_completion.py @@ -4078,19 +4078,29 @@ async def test_async_text_completion_chat_model_stream(): # asyncio.run(test_async_text_completion_chat_model_stream()) -@pytest.mark.asyncio -async def test_completion_codestral_fim_api(): - try: - litellm.set_verbose = True - response = await litellm.atext_completion( - model="text-completion-codestral/codestral-2405", - prompt="def is_odd(n): \n return n % 2 == 1 \ndef test_is_odd():", - ) - # Add any assertions here to check the response - print(response) +# @pytest.mark.asyncio +# async def test_completion_codestral_fim_api(): +# try: +# litellm.set_verbose = True +# from litellm._logging import verbose_logger +# import logging +# verbose_logger.setLevel(level=logging.DEBUG) +# response = await litellm.atext_completion( +# model="text-completion-codestral/codestral-2405", +# prompt="def is_odd(n): \n return n % 2 == 1 \ndef test_is_odd():", +# suffix="return True", +# temperature=0, +# top_p=0.4, +# max_tokens=10, +# # min_tokens=10, +# seed=10, +# stop=["return"], +# ) +# # Add any assertions here to check the response +# print(response) - # cost = litellm.completion_cost(completion_response=response) - # print("cost to make mistral completion=", cost) - # assert cost > 0.0 - except Exception as e: - pytest.fail(f"Error occurred: {e}") +# # cost = litellm.completion_cost(completion_response=response) +# # print("cost to make mistral completion=", cost) +# # assert cost > 0.0 +# except Exception as e: +# pytest.fail(f"Error occurred: {e}") diff --git a/litellm/utils.py b/litellm/utils.py index c720a24cd03..054648825bd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2968,7 +2968,7 @@ def get_optional_params( optional_params["stream"] = stream if max_tokens: optional_params["max_tokens"] = max_tokens - elif custom_llm_provider == "mistral": + elif custom_llm_provider == "mistral" or custom_llm_provider == "codestral": supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider ) @@ -2976,6 +2976,15 @@ def get_optional_params( optional_params = litellm.MistralConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params ) + elif custom_llm_provider == "text-completion-codestral": + supported_params = get_supported_openai_params( + model=model, custom_llm_provider=custom_llm_provider + ) + _check_valid_arg(supported_params=supported_params) + optional_params = litellm.MistralTextCompletionConfig().map_openai_params( + non_default_params=non_default_params, optional_params=optional_params + ) + elif custom_llm_provider == "databricks": supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider @@ -3649,11 +3658,14 @@ def get_supported_openai_params( "tool_choice", "max_retries", ] - elif custom_llm_provider == "mistral": + elif custom_llm_provider == "mistral" or custom_llm_provider == "codestral": + # mistal and codestral api have the exact same params if request_type == "chat_completion": return litellm.MistralConfig().get_supported_openai_params() elif request_type == "embeddings": return litellm.MistralEmbeddingConfig().get_supported_openai_params() + elif custom_llm_provider == "text-completion-codestral": + return litellm.MistralTextCompletionConfig().get_supported_openai_params() elif custom_llm_provider == "replicate": return [ "stream",