From edfc35ddac1eff638142a33833596953e5b4d4b4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Nov 2025 13:08:03 -0800 Subject: [PATCH] [Feature]: Add Provider publicai.co (#17230) * init PublicAIChatConfig * add publicai * init public ai * add publicai * add publicai/swiss-ai models etc --- docs/my-website/docs/providers/publicai.md | 209 ++++++++++++++++++ docs/my-website/sidebars.js | 1 + litellm/__init__.py | 6 + litellm/constants.py | 4 + .../get_llm_provider_logic.py | 10 + litellm/llms/publicai/chat/transformation.py | 114 ++++++++++ ...odel_prices_and_context_window_backup.json | 110 +++++++++ litellm/types/utils.py | 1 + model_prices_and_context_window.json | 110 +++++++++ provider_endpoints_support.json | 16 ++ .../test_publicai_chat_transformation.py | 141 ++++++++++++ 11 files changed, 722 insertions(+) create mode 100644 docs/my-website/docs/providers/publicai.md create mode 100644 litellm/llms/publicai/chat/transformation.py create mode 100644 tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py diff --git a/docs/my-website/docs/providers/publicai.md b/docs/my-website/docs/providers/publicai.md new file mode 100644 index 0000000000..1ab8bd5a06 --- /dev/null +++ b/docs/my-website/docs/providers/publicai.md @@ -0,0 +1,209 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# PublicAI + +## Overview + +| Property | Details | +|-------|-------| +| Description | PublicAI provides large language models including essential models like the swiss-ai apertus model. | +| Provider Route on LiteLLM | `publicai/` | +| Link to Provider Doc | [PublicAI ↗](https://platform.publicai.co/) | +| Base URL | `https://platform.publicai.co/` | +| Supported Operations | [`/chat/completions`](#sample-usage) | + +
+
+ +https://platform.publicai.co/ + +**We support ALL PublicAI models, just set `publicai/` as a prefix when sending completion requests** + +## Required Variables + +```python showLineNumbers title="Environment Variables" +os.environ["PUBLICAI_API_KEY"] = "" # your PublicAI API key +``` + +You can overwrite the base url with: + +``` +os.environ["PUBLICAI_API_BASE"] = "https://platform.publicai.co/v1" +``` + +## Usage - LiteLLM Python SDK + +### Non-streaming + +```python showLineNumbers title="PublicAI Non-streaming Completion" +import os +import litellm +from litellm import completion + +os.environ["PUBLICAI_API_KEY"] = "" # your PublicAI API key + +messages = [{"content": "Hello, how are you?", "role": "user"}] + +# PublicAI call +response = completion( + model="publicai/swiss-ai/apertus-8b-instruct", + messages=messages +) + +print(response) +``` + +### Streaming + +```python showLineNumbers title="PublicAI Streaming Completion" +import os +import litellm +from litellm import completion + +os.environ["PUBLICAI_API_KEY"] = "" # your PublicAI API key + +messages = [{"content": "Hello, how are you?", "role": "user"}] + +# PublicAI call with streaming +response = completion( + model="publicai/swiss-ai/apertus-8b-instruct", + messages=messages, + stream=True +) + +for chunk in response: + print(chunk) +``` + +## Usage - LiteLLM Proxy + +Add the following to your LiteLLM Proxy configuration file: + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: swiss-ai-apertus-8b + litellm_params: + model: publicai/swiss-ai/apertus-8b-instruct + api_key: os.environ/PUBLICAI_API_KEY + + - model_name: swiss-ai-apertus-70b + litellm_params: + model: publicai/swiss-ai/apertus-70b-instruct + api_key: os.environ/PUBLICAI_API_KEY +``` + +Start your LiteLLM Proxy server: + +```bash showLineNumbers title="Start LiteLLM Proxy" +litellm --config config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + + + + +```python showLineNumbers title="PublicAI via Proxy - Non-streaming" +from openai import OpenAI + +# Initialize client with your proxy URL +client = OpenAI( + base_url="http://localhost:4000", # Your proxy URL + api_key="your-proxy-api-key" # Your proxy API key +) + +# Non-streaming response +response = client.chat.completions.create( + model="swiss-ai-apertus-8b", + messages=[{"role": "user", "content": "hello from litellm"}] +) + +print(response.choices[0].message.content) +``` + +```python showLineNumbers title="PublicAI via Proxy - Streaming" +from openai import OpenAI + +# Initialize client with your proxy URL +client = OpenAI( + base_url="http://localhost:4000", # Your proxy URL + api_key="your-proxy-api-key" # Your proxy API key +) + +# Streaming response +response = client.chat.completions.create( + model="swiss-ai-apertus-8b", + messages=[{"role": "user", "content": "hello from litellm"}], + stream=True +) + +for chunk in response: + if chunk.choices[0].delta.content is not None: + print(chunk.choices[0].delta.content, end="") +``` + + + + + +```python showLineNumbers title="PublicAI via Proxy - LiteLLM SDK" +import litellm + +# Configure LiteLLM to use your proxy +response = litellm.completion( + model="litellm_proxy/swiss-ai-apertus-8b", + messages=[{"role": "user", "content": "hello from litellm"}], + api_base="http://localhost:4000", + api_key="your-proxy-api-key" +) + +print(response.choices[0].message.content) +``` + +```python showLineNumbers title="PublicAI via Proxy - LiteLLM SDK Streaming" +import litellm + +# Configure LiteLLM to use your proxy with streaming +response = litellm.completion( + model="litellm_proxy/swiss-ai-apertus-8b", + messages=[{"role": "user", "content": "hello from litellm"}], + api_base="http://localhost:4000", + api_key="your-proxy-api-key", + stream=True +) + +for chunk in response: + if hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta.content is not None: + print(chunk.choices[0].delta.content, end="") +``` + + + + + +```bash showLineNumbers title="PublicAI via Proxy - cURL" +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer your-proxy-api-key" \ + -d '{ + "model": "swiss-ai-apertus-8b", + "messages": [{"role": "user", "content": "hello from litellm"}] + }' +``` + +```bash showLineNumbers title="PublicAI via Proxy - cURL Streaming" +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer your-proxy-api-key" \ + -d '{ + "model": "swiss-ai-apertus-8b", + "messages": [{"role": "user", "content": "hello from litellm"}], + "stream": true + }' +``` + + + + +For more detailed information on using the LiteLLM Proxy, see the [LiteLLM Proxy documentation](../providers/litellm_proxy). diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 917fdca602..802ffdd5bb 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -622,6 +622,7 @@ const sidebars = { "providers/ovhcloud", "providers/perplexity", "providers/petals", + "providers/publicai", "providers/predibase", "providers/recraft", "providers/replicate", diff --git a/litellm/__init__.py b/litellm/__init__.py index 71be5113e2..bbe17c97b7 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -555,6 +555,7 @@ deepgram_models: Set = set() elevenlabs_models: Set = set() dashscope_models: Set = set() moonshot_models: Set = set() +publicai_models: Set = set() v0_models: Set = set() morph_models: Set = set() lambda_ai_models: Set = set() @@ -781,6 +782,8 @@ def add_known_models(): dashscope_models.add(key) elif value.get("litellm_provider") == "moonshot": moonshot_models.add(key) + elif value.get("litellm_provider") == "publicai": + publicai_models.add(key) elif value.get("litellm_provider") == "v0": v0_models.add(key) elif value.get("litellm_provider") == "morph": @@ -899,6 +902,7 @@ model_list = list( | elevenlabs_models | dashscope_models | moonshot_models + | publicai_models | v0_models | morph_models | lambda_ai_models @@ -992,6 +996,7 @@ models_by_provider: dict = { "heroku": heroku_models, "dashscope": dashscope_models, "moonshot": moonshot_models, + "publicai": publicai_models, "v0": v0_models, "morph": morph_models, "lambda_ai": lambda_ai_models, @@ -1370,6 +1375,7 @@ from .llms.nebius.chat.transformation import NebiusConfig from .llms.wandb.chat.transformation import WandbConfig from .llms.dashscope.chat.transformation import DashScopeChatConfig from .llms.moonshot.chat.transformation import MoonshotChatConfig +from .llms.publicai.chat.transformation import PublicAIChatConfig from .llms.docker_model_runner.chat.transformation import DockerModelRunnerChatConfig from .llms.v0.chat.transformation import V0ChatConfig from .llms.oci.chat.transformation import OCIChatConfig diff --git a/litellm/constants.py b/litellm/constants.py index 9235916dd4..f1126b1909 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -384,6 +384,7 @@ LITELLM_CHAT_PROVIDERS = [ "nebius", "dashscope", "moonshot", + "publicai", "v0", "heroku", "oci", @@ -526,6 +527,7 @@ openai_compatible_endpoints: List = [ "api.studio.nebius.ai/v1", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", "https://api.moonshot.ai/v1", + "https://platform.publicai.co/v1", "https://api.v0.dev/v1", "https://api.morphllm.com/v1", "https://api.lambda.ai/v1", @@ -571,6 +573,7 @@ openai_compatible_providers: List = [ "nebius", "dashscope", "moonshot", + "publicai", "v0", "morph", "lambda_ai", @@ -593,6 +596,7 @@ openai_text_completion_compatible_providers: List = ( "nebius", "dashscope", "moonshot", + "publicai", "v0", "lambda_ai", "hyperbolic", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 1efea63beb..3bb1e4afb9 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -258,6 +258,9 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "api.moonshot.ai/v1": custom_llm_provider = "moonshot" dynamic_api_key = get_secret_str("MOONSHOT_API_KEY") + elif endpoint == "platform.publicai.co/v1": + custom_llm_provider = "publicai" + dynamic_api_key = get_secret_str("PUBLICAI_API_KEY") elif endpoint == "https://api.v0.dev/v1": custom_llm_provider = "v0" dynamic_api_key = get_secret_str("V0_API_KEY") @@ -759,6 +762,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ) = litellm.MoonshotChatConfig()._get_openai_compatible_provider_info( api_base, api_key ) + elif custom_llm_provider == "publicai": + ( + api_base, + dynamic_api_key, + ) = litellm.PublicAIChatConfig()._get_openai_compatible_provider_info( + api_base, api_key + ) elif custom_llm_provider == "docker_model_runner": ( api_base, diff --git a/litellm/llms/publicai/chat/transformation.py b/litellm/llms/publicai/chat/transformation.py new file mode 100644 index 0000000000..93921e48d5 --- /dev/null +++ b/litellm/llms/publicai/chat/transformation.py @@ -0,0 +1,114 @@ +""" +Translates from OpenAI's `/v1/chat/completions` to PublicAI's `/v1/chat/completions` +""" + +from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, overload + +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + handle_messages_with_content_list_to_str_conversion, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues + +from ...openai.chat.gpt_transformation import OpenAIGPTConfig + + +class PublicAIChatConfig(OpenAIGPTConfig): + @overload + def _transform_messages( + self, messages: List[AllMessageValues], model: str, is_async: Literal[True] + ) -> Coroutine[Any, Any, List[AllMessageValues]]: + ... + + @overload + def _transform_messages( + self, + messages: List[AllMessageValues], + model: str, + is_async: Literal[False] = False, + ) -> List[AllMessageValues]: + ... + + def _transform_messages( + self, messages: List[AllMessageValues], model: str, is_async: bool = False + ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: + """ + PublicAI does not support content in list format. + """ + messages = handle_messages_with_content_list_to_str_conversion(messages) + if is_async: + return super()._transform_messages( + messages=messages, model=model, is_async=True + ) + else: + return super()._transform_messages( + messages=messages, model=model, is_async=False + ) + + def _get_openai_compatible_provider_info( + self, api_base: Optional[str], api_key: Optional[str] + ) -> Tuple[Optional[str], Optional[str]]: + api_base = ( + api_base + or get_secret_str("PUBLICAI_API_BASE") + or "https://platform.publicai.co/v1" + ) # type: ignore + dynamic_api_key = api_key or get_secret_str("PUBLICAI_API_KEY") + return api_base, dynamic_api_key + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + If api_base is not provided, use the default PublicAI /chat/completions endpoint. + """ + if not api_base: + api_base = "https://platform.publicai.co/v1" + + if not api_base.endswith("/chat/completions"): + api_base = f"{api_base}/chat/completions" + + return api_base + + def get_supported_openai_params(self, model: str) -> list: + """ + Get the supported OpenAI params for PublicAI models + + PublicAI limitations: + - functions parameter is not supported (use tools instead) + """ + excluded_params: List[str] = ["functions"] + + base_openai_params = super().get_supported_openai_params(model=model) + final_params: List[str] = [] + for param in base_openai_params: + if param not in excluded_params: + final_params.append(param) + + return final_params + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to PublicAI parameters + """ + supported_openai_params = self.get_supported_openai_params(model) + for param, value in non_default_params.items(): + if param == "max_completion_tokens": + optional_params["max_tokens"] = value + elif param in supported_openai_params: + optional_params[param] = value + + return optional_params + diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f07f1cbe5e..8ef4d66381 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21570,6 +21570,116 @@ "mode": "chat", "output_cost_per_token": 2.8e-07 }, + "publicai/swiss-ai/apertus-8b-instruct": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/swiss-ai/apertus-70b-instruct": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/aisingapore/Gemma-SEA-LION-v4-27B-IT": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/BSC-LT/salamandra-7b-instruct-tools-16k": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 16384, + "max_output_tokens": 4096, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/BSC-LT/ALIA-40b-instruct_Q8_0": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/allenai/Olmo-3-7B-Instruct": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/allenai/Olmo-3-7B-Think": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true + }, + "publicai/allenai/Olmo-3-32B-Think": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true + }, "qwen.qwen3-coder-480b-a35b-v1:0": { "input_cost_per_token": 2.2e-07, "litellm_provider": "bedrock_converse", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b0d081d8f8..37a476b33f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2598,6 +2598,7 @@ class LlmProviders(str, Enum): TEXT_COMPLETION_CODESTRAL = "text-completion-codestral" DASHSCOPE = "dashscope" MOONSHOT = "moonshot" + PUBLICAI = "publicai" V0 = "v0" MORPH = "morph" LAMBDA_AI = "lambda_ai" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f07f1cbe5e..8ef4d66381 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21570,6 +21570,116 @@ "mode": "chat", "output_cost_per_token": 2.8e-07 }, + "publicai/swiss-ai/apertus-8b-instruct": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/swiss-ai/apertus-70b-instruct": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/aisingapore/Gemma-SEA-LION-v4-27B-IT": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/BSC-LT/salamandra-7b-instruct-tools-16k": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 16384, + "max_output_tokens": 4096, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/BSC-LT/ALIA-40b-instruct_Q8_0": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 8192, + "max_output_tokens": 4096, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/allenai/Olmo-3-7B-Instruct": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "publicai/allenai/Olmo-3-7B-Think": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true + }, + "publicai/allenai/Olmo-3-32B-Think": { + "input_cost_per_token": 0.0, + "litellm_provider": "publicai", + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://platform.publicai.co/docs", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true + }, "qwen.qwen3-coder-480b-a35b-v1:0": { "input_cost_per_token": 2.2e-07, "litellm_provider": "bedrock_converse", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index e05d6d8d00..5eab130cdd 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1329,6 +1329,22 @@ "rerank": false } }, + "publicai": { + "display_name": "PublicAI (`publicai`)", + "url": "https://docs.litellm.ai/docs/providers/publicai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } + }, "predibase": { "display_name": "Predibase (`predibase`)", "url": "https://docs.litellm.ai/docs/providers/predibase", diff --git a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py b/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py new file mode 100644 index 0000000000..47722686c5 --- /dev/null +++ b/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py @@ -0,0 +1,141 @@ +""" +Unit tests for PublicAI configuration. + +These tests validate the PublicAIChatConfig class which extends OpenAIGPTConfig. +PublicAI is an OpenAI-compatible provider with minor customizations. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) + +import pytest + +import litellm +import litellm.utils +from litellm import completion +from litellm.llms.publicai.chat.transformation import PublicAIChatConfig + + +class TestPublicAIConfig: + """Test class for PublicAI functionality""" + + def test_default_api_base(self): + """ + Test that default API base is used when none is provided + """ + config = PublicAIChatConfig() + headers = {} + api_key = "fake-publicai-key" + + result = config.validate_environment( + headers=headers, + model="swiss-ai-apertus", + messages=[{"role": "user", "content": "Hey"}], + optional_params={}, + litellm_params={}, + api_key=api_key, + api_base=None, + ) + + assert result["Authorization"] == f"Bearer {api_key}" + assert result["Content-Type"] == "application/json" + + def test_get_supported_openai_params(self): + """ + Test that get_supported_openai_params returns correct params + """ + config = PublicAIChatConfig() + + supported_params = config.get_supported_openai_params(model="swiss-ai-apertus") + + assert "tools" in supported_params + assert "tool_choice" in supported_params + assert "temperature" in supported_params + assert "max_tokens" in supported_params + assert "stream" in supported_params + + assert "functions" not in supported_params + + def test_map_openai_params_excludes_functions(self): + """ + Test that functions parameter is not mapped + """ + config = PublicAIChatConfig() + + non_default_params = { + "functions": [{"name": "test_function", "description": "Test function"}], + "temperature": 0.7, + "max_tokens": 1000 + } + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model="swiss-ai-apertus", + drop_params=False + ) + + assert "functions" not in result + assert result.get("temperature") == 0.7 + assert result.get("max_tokens") == 1000 + + def test_map_openai_params_max_completion_tokens_mapping(self): + """ + Test that max_completion_tokens is mapped to max_tokens + """ + config = PublicAIChatConfig() + + non_default_params = { + "max_completion_tokens": 1000, + "temperature": 0.7 + } + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model="swiss-ai-apertus", + drop_params=False + ) + + assert result.get("max_tokens") == 1000 + assert "max_completion_tokens" not in result + assert result.get("temperature") == 0.7 + + def test_get_complete_url(self): + """ + Test that get_complete_url constructs the correct endpoint URL + """ + config = PublicAIChatConfig() + + url = config.get_complete_url( + api_base=None, + api_key="fake-key", + model="swiss-ai-apertus", + optional_params={}, + litellm_params={}, + stream=False + ) + + assert url == "https://platform.publicai.co/v1/chat/completions" + + def test_get_complete_url_with_custom_base(self): + """ + Test that get_complete_url works with custom api_base + """ + config = PublicAIChatConfig() + + url = config.get_complete_url( + api_base="https://custom.publicai.co/v1", + api_key="fake-key", + model="swiss-ai-apertus", + optional_params={}, + litellm_params={}, + stream=False + ) + + assert url == "https://custom.publicai.co/v1/chat/completions" +