diff --git a/README.md b/README.md index 9987f13d127..0ae0d68d250 100644 --- a/README.md +++ b/README.md @@ -342,6 +342,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | | | [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | | | [CompactifAI (`compactifai`)](https://docs.litellm.ai/docs/providers/compactifai) | ✅ | ✅ | ✅ | | | | | | | | +| [CoralBricks (`coralbricks`)](https://docs.litellm.ai/docs/providers/coralbricks) | ✅ | ✅ | ✅ | | | | | | | | | [Custom (`custom`)](https://docs.litellm.ai/docs/providers/custom_llm_server) | ✅ | ✅ | ✅ | | | | | | | | | [Custom OpenAI (`custom_openai`)](https://docs.litellm.ai/docs/providers/openai_compatible) | ✅ | ✅ | ✅ | | | ✅ | ✅ | ✅ | ✅ | | | [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | diff --git a/litellm/__init__.py b/litellm/__init__.py index 0a27c9e26fb..d29a0e03d5b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -721,6 +721,7 @@ aws_polly_models: Set = set() transcribe_models: Set = set() gigachat_models: Set = set() llamagate_models: Set = set() +coralbricks_models: Set[str] = set() reducto_models: Set = set() bedrock_mantle_models: Set = set() @@ -1012,6 +1013,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje gigachat_models.add(key) elif value.get("litellm_provider") == "llamagate": llamagate_models.add(key) + elif value.get("litellm_provider") == "coralbricks": + coralbricks_models.add(key) elif value.get("litellm_provider") == "reducto": reducto_models.add(key) elif value.get("litellm_provider") == "bedrock_mantle": @@ -1146,6 +1149,7 @@ model_list = list( | docker_model_runner_models | reducto_models | bedrock_mantle_models + | coralbricks_models | set(clarifai_models) ) @@ -1265,6 +1269,7 @@ def _build_models_by_provider() -> dict: "transcribe": transcribe_models, "gigachat": gigachat_models, "llamagate": llamagate_models, + "coralbricks": coralbricks_models, "reducto": reducto_models, "bedrock_mantle": bedrock_mantle_models, } diff --git a/litellm/constants.py b/litellm/constants.py index 38d16d61d7b..d9441f36a1f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1007,6 +1007,7 @@ openai_compatible_endpoints: Final[list] = [ "https://api.meta.ai/v1", "https://api.sailresearch.com/v1", "https://api.cognition.ai/v1", + "https://inference.coralbricks.ai/v1", "https://api.cortecs.ai/v1", "https://api.scx.ai/v1", "https://api.prisminference.com/v1", @@ -1083,6 +1084,7 @@ openai_compatible_providers: Final[list] = [ "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider "cognition", + "coralbricks", "cortecs", "scx-ai", "prism", diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index f3aa7768928..d25257b0e81 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -180,6 +180,12 @@ "api_key_env": "COGNITION_API_KEY", "api_base_env": "COGNITION_API_BASE" }, + "coralbricks": { + "base_url": "https://inference.coralbricks.ai/v1", + "api_key_env": "CORALBRICKS_API_KEY", + "api_base_env": "CORALBRICKS_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + }, "cortecs": { "base_url": "https://api.cortecs.ai/v1", "api_key_env": "CORTECS_API_KEY", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d8459174eb8..1a1a035191c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16634,6 +16634,144 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "coralbricks/glm-5.3-fast": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3-fp4": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/deepseek-v4.1-flash-fast": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash-fast-fp4": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "computer-use-preview": { "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index a0e75b1f02d..8dbdac37717 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -618,6 +618,24 @@ "interactions": true } }, + "coralbricks": { + "display_name": "CoralBricks (`coralbricks`)", + "url": "https://docs.litellm.ai/docs/providers/coralbricks", + "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, + "a2a": false, + "interactions": false + } + }, "cortecs": { "display_name": "Cortecs (`cortecs`)", "url": "https://docs.litellm.ai/docs/providers/cortecs", diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 633e9edf71c..cb7d099cafb 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1033,6 +1033,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "CORALBRICKS", + "provider_display_name": "CoralBricks", + "litellm_provider": "coralbricks", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://inference.coralbricks.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "coralbricks/glm-5.3-fast" + }, { "provider": "CORTECS", "provider_display_name": "Cortecs", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c486fc4f4d0..e6103270ee3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4309,6 +4309,7 @@ class LlmProviders(str, Enum): LIBERTAI = "libertai" PINSTRIPES = "pinstripes" COGNITION = "cognition" + CORALBRICKS = "coralbricks" CORTECS = "cortecs" SCX_AI = "scx-ai" PRISM = "prism" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d8459174eb8..1a1a035191c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16634,6 +16634,144 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "coralbricks/glm-5.3-fast": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3-fp4": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/glm-5.3": { + "cache_creation_input_token_cost": 1.68e-06, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1.12e-06, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "coralbricks/deepseek-v4.1-flash-fast": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash-fast-fp4": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "coralbricks/deepseek-v4.1-flash": { + "cache_creation_input_token_cost": 9e-08, + "cache_read_input_token_cost": 0.0, + "input_cost_per_token": 1e-08, + "litellm_provider": "coralbricks", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.coralbricks.ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "computer-use-preview": { "deprecation_date": "2026-07-23", "input_cost_per_token": 3e-06, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 5b179b2e8d8..8263c05a4d9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -672,6 +672,24 @@ "interactions": true } }, + "coralbricks": { + "display_name": "CoralBricks (`coralbricks`)", + "url": "https://docs.litellm.ai/docs/providers/coralbricks", + "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, + "a2a": false, + "interactions": false + } + }, "cortecs": { "display_name": "Cortecs (`cortecs`)", "url": "https://docs.litellm.ai/docs/providers/cortecs", diff --git a/tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts b/tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts new file mode 100644 index 00000000000..1fea7eacbe3 --- /dev/null +++ b/tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts @@ -0,0 +1,172 @@ +import { test, expect, APIRequestContext, Page as PlaywrightPage } from "@playwright/test"; +import { randomUUID } from "node:crypto"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody } from "../../helpers/roundTrip"; + +const master = process.env.LITELLM_MASTER_KEY ?? "sk-integration-master"; +const headers = { Authorization: `Bearer ${master}` }; +const upstream = (process.env.INTEGRATION_UPSTREAM_URL ?? "http://127.0.0.1:8190").replace(/\/+$/, ""); +const CORALBRICKS_MODEL = "deepseek-v4.1-flash-fast"; +const UI_KEY = "cb_ui_integration_key"; +const UPSTREAM_PROMPT_TOKENS = 20; +const UPSTREAM_COMPLETION_TOKENS = 20; + +type DeploymentRow = { + model_name: string; + litellm_params?: { model?: string; custom_llm_provider?: string }; + model_info?: { id?: string }; +}; + +type CatalogRow = { input_cost_per_token?: number; output_cost_per_token?: number }; + +type SpendRow = { + request_id: string; + custom_llm_provider: string; + model: string; + model_group: string; + prompt_tokens: number; + completion_tokens: number; + spend: number | string; +}; + +const optionNamed = (label: string): RegExp => + new RegExp(`(^|\\s)${label.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`); + +async function loginAsAdmin(page: PlaywrightPage): Promise { + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill("admin"); + await page.getByPlaceholder("Enter your password").fill(master); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page).toHaveURL( + (url) => url.pathname.startsWith("/ui") && !url.pathname.includes("login"), + ); +} + +async function selectProvider(page: PlaywrightPage, providerName: string): Promise { + const providerDropdown = page.getByRole("combobox", { name: "Provider", exact: true }); + await providerDropdown.click(); + await providerDropdown.fill(providerName); + await page.getByRole("option", { name: optionNamed(providerName) }).click(); + await expect(providerDropdown).toHaveValue(providerName); +} + +async function deploymentNamed( + request: APIRequestContext, + modelName: string, +): Promise { + const response = await request.get("/v1/model/info", { headers }); + expect(response.ok(), await response.text()).toBe(true); + const rows = (await response.json()).data as DeploymentRow[]; + return rows.find((row) => row.model_name === modelName); +} + +async function deleteDeployment(request: APIRequestContext, modelName: string): Promise { + const stored = await deploymentNamed(request, modelName); + const id = stored?.model_info?.id; + if (!id) return; + const deleted = await request.post("/model/delete", { headers, data: { id } }); + expect(deleted.ok(), await deleted.text()).toBe(true); +} + +async function chatOnce(request: APIRequestContext, modelName: string): Promise<{ status: number; body: string }> { + const response = await request.post("/v1/chat/completions", { + headers, + data: { model: modelName, messages: [{ role: "user", content: `hello from ${modelName}` }] }, + }); + return { status: response.status(), body: await response.text() }; +} + +async function spendRow(request: APIRequestContext, requestId: string): Promise { + const response = await request.get(`/spend/logs?request_id=${encodeURIComponent(requestId)}`, { headers }); + expect(response.ok(), await response.text()).toBe(true); + const rows = (await response.json()) as SpendRow[]; + return rows[0]; +} + +test("the Add Model form offers CoralBricks and the deployment it creates serves and bills a chat", async ({ + page, + request, +}) => { + const alias = `integration-ui-coralbricks-${randomUUID().replace(/-/g, "").slice(0, 12)}`; + try { + await loginAsAdmin(page); + await navigateToPage(page, Page.Models); + const catalogLoaded = page.waitForResponse( + (response) => response.url().includes("/public/litellm_model_cost_map?catalog_only=true") && response.ok(), + ); + await page.getByRole("tab", { name: "Add Model" }).click(); + const catalog = (await (await catalogLoaded).json()) as Record; + const catalogRow = catalog[`coralbricks/${CORALBRICKS_MODEL}`]; + + await selectProvider(page, "CoralBricks"); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "Custom Model Name (Enter below)" }).click(); + await page.keyboard.press("Escape"); + await page.getByPlaceholder("Enter custom model name").fill(CORALBRICKS_MODEL); + + const publicName = page.getByRole("row").filter({ hasText: CORALBRICKS_MODEL }).getByRole("textbox"); + await expect(publicName).toHaveValue(CORALBRICKS_MODEL); + await publicName.fill(alias); + + await page.locator("#api_base").fill(`${upstream}/v1`); + await page.locator("#api_key").fill(UI_KEY); + + const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { + await page.getByTestId("add-model-btn").click(); + }); + expect(created.model_name, "the public name typed into the mapping row").toBe(alias); + expect(created.litellm_params?.model, "the custom model name typed into the form").toBe(CORALBRICKS_MODEL); + expect(created.litellm_params?.custom_llm_provider, "the provider the picker maps CoralBricks to").toBe( + "coralbricks", + ); + expect(created.litellm_params?.api_base).toBe(`${upstream}/v1`); + expect(created.litellm_params?.api_key).toBe(UI_KEY); + await expect(page.getByText("created successfully")).toBeVisible({ timeout: 15_000 }); + + let served = { status: 0, body: "" }; + await expect + .poll( + async () => { + served = await chatOnce(request, alias); + return served.status; + }, + { message: `the deployment added through the UI never served a chat: ${served.body}`, timeout: 60_000 }, + ) + .toBe(200); + const requestId = JSON.parse(served.body).id as string; + expect(requestId, served.body).toBeTruthy(); + + await expect + .poll(async () => (await deploymentNamed(request, alias))?.litellm_params?.custom_llm_provider, { + message: `no worker listed ${alias}`, + timeout: 60_000, + }) + .toBe("coralbricks"); + const inputRate = catalogRow?.input_cost_per_token; + const outputRate = catalogRow?.output_cost_per_token; + expect(inputRate, JSON.stringify(catalogRow)).toBeGreaterThan(0); + expect(outputRate, JSON.stringify(catalogRow)).toBeGreaterThan(0); + const expectedSpend = UPSTREAM_PROMPT_TOKENS * inputRate! + UPSTREAM_COMPLETION_TOKENS * outputRate!; + + let row: SpendRow | undefined; + await expect + .poll( + async () => { + row = await spendRow(request, requestId); + return row?.request_id; + }, + { message: `no spend row for ${requestId}`, timeout: 60_000 }, + ) + .toBe(requestId); + expect(row?.custom_llm_provider, JSON.stringify(row)).toBe("coralbricks"); + expect(row?.model, JSON.stringify(row)).toBe(created.litellm_params.model); + expect(row?.model_group, JSON.stringify(row)).toBe(alias); + expect(row?.prompt_tokens, JSON.stringify(row)).toBe(UPSTREAM_PROMPT_TOKENS); + expect(row?.completion_tokens, JSON.stringify(row)).toBe(UPSTREAM_COMPLETION_TOKENS); + expect(Number(row?.spend), JSON.stringify(row)).toBeCloseTo(expectedSpend, 10); + } finally { + await deleteDeployment(request, alias); + } +}); diff --git a/tests/e2e/ui/tests/integrationCritical/expected.json b/tests/e2e/ui/tests/integrationCritical/expected.json index 117df8d4892..e2b3f32645c 100644 --- a/tests/e2e/ui/tests/integrationCritical/expected.json +++ b/tests/e2e/ui/tests/integrationCritical/expected.json @@ -14,5 +14,6 @@ "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::the Add Credential dialog stores each federation identity source with only the values it needs and refuses the invalid ones", "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::the Edit Credential dialog sends only the changed values and deletes the keys the new choice leaves behind", "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::a proxy admin saves a federation credential from the Add Model tab and the model is created against it", - "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::a team admin who is not a proxy admin gets no federation shortcut on the Add Model tab" + "tests/e2e/ui/tests/integrationCritical/anthropicFederationCredential.spec.ts::a team admin who is not a proxy admin gets no federation shortcut on the Add Model tab", + "tests/e2e/ui/tests/integrationCritical/addModelCoralBricks.spec.ts::the Add Model form offers CoralBricks and the deployment it creates serves and bills a chat" ] diff --git a/tests/integration/providers/_coralbricks.py b/tests/integration/providers/_coralbricks.py new file mode 100644 index 00000000000..0cfbcd94d64 --- /dev/null +++ b/tests/integration/providers/_coralbricks.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +import json +import re +import socket +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue + +PROVIDER: Final = "coralbricks" +MODEL: Final = "deepseek-v4.1-flash-fast" +MODELS: Final = (MODEL, "glm-5.3-fast") +API_KEY: Final = "cb_synthetic_integration_key" +ANSWER: Final = "coral ok" +NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +PROMPT_TOKENS: Final = 4357 +CACHED_TOKENS: Final = 4352 +CACHE_WRITE_TOKENS: Final = 5 +COMPLETION_TOKENS: Final = 20 +REASONING_TOKENS: Final = 12 +COST_MAP: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +LEAK_FIELDS: Final = frozenset( + {"litellm_params", "litellm_logging_obj", "litellm_call_id", "litellm_metadata", "proxy_server_request"} +) +READINESS_TEXT: Final = "model readiness" +MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +USER_MESSAGES: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Say hello to the reef"}] +SPEND_COLUMNS: Final = ( + "SELECT request_id, model, custom_llm_provider, model_group, call_type, status, spend, prompt_tokens, " + 'completion_tokens, api_base FROM "LiteLLM_SpendLogs"' +) + + +@dataclass(frozen=True, slots=True) +class Rates: + input: float + output: float + cache_write: float + cache_read: float + + +def number(value: JsonValue) -> float: + assert isinstance(value, (int, float)), value + return float(value) + + +def row_of(model: str) -> Mapping[str, JsonValue]: + row: Final = object_value(JSON_OBJECT.validate_json(COST_MAP.read_bytes())[f"{PROVIDER}/{model}"]) + assert row["litellm_provider"] == PROVIDER, row + return row + + +def catalog_rows() -> frozenset[str]: + rows: Final = frozenset( + key for key in JSON_OBJECT.validate_json(COST_MAP.read_bytes()) if key.startswith(f"{PROVIDER}/") + ) + assert {f"{PROVIDER}/{model}" for model in MODELS} <= rows, sorted(rows) + return rows + + +def rates_of(model: str) -> Rates: + row: Final = row_of(model) + return Rates( + input=number(row["input_cost_per_token"]), + output=number(row["output_cost_per_token"]), + cache_write=number(row["cache_creation_input_token_cost"]), + cache_read=number(row["cache_read_input_token_cost"]), + ) + + +def expected_spend(model: str) -> float: + rates: Final = rates_of(model) + return ( + (PROMPT_TOKENS - CACHED_TOKENS - CACHE_WRITE_TOKENS) * rates.input + + CACHED_TOKENS * rates.cache_read + + CACHE_WRITE_TOKENS * rates.cache_write + + COMPLETION_TOKENS * rates.output + ) + + +def chat_usage() -> dict[str, JsonValue]: + return { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + "prompt_tokens_details": { + "cached_tokens": CACHED_TOKENS, + "cache_write_tokens": CACHE_WRITE_TOKENS, + "billable_cache_write_tokens": CACHE_WRITE_TOKENS, + "cache_write_blocks": 3, + }, + "completion_tokens_details": {"reasoning_tokens": REASONING_TOKENS}, + } + + +def sse(payload: Mapping[str, JsonValue], event: str | None = None) -> bytes: + prefix: Final = f"event: {event}\n" if event is not None else "" + return f"{prefix}data: {json.dumps(payload)}\n\n".encode() + + +def chat_reply(identity: str, model: str, text: str = ANSWER) -> Reply: + body: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": chat_usage(), + } + return Reply(body=json.dumps(body).encode()) + + +def chat_stream_reply(identity: str, model: str, text: str = ANSWER, pause: float = 0) -> Reply: + chunk: Final[dict[str, JsonValue]] = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": model} + frames: Final[tuple[dict[str, JsonValue], ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}]}, + {**chunk, "choices": [{"index": 0, "delta": {"content": text}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**chunk, "choices": [], "usage": chat_usage()}, + ) + return Reply( + content_type="text/event-stream", + chunks=(*(sse(frame) for frame in frames), b"data: [DONE]\n\n"), + pause_between_chunks=pause, + ) + + +def responses_body(identity: str, model: str, text: str = ANSWER) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": model, + "output": [ + { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "usage": { + "input_tokens": PROMPT_TOKENS, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + "input_tokens_details": { + "cached_tokens": CACHED_TOKENS, + "cache_write_tokens": CACHE_WRITE_TOKENS, + "billable_cache_write_tokens": CACHE_WRITE_TOKENS, + }, + "output_tokens_details": {"reasoning_tokens": REASONING_TOKENS}, + }, + } + + +def responses_reply(identity: str, model: str, stream: bool, text: str = ANSWER, pause: float = 0) -> Reply: + body: Final = responses_body(identity, model, text) + if not stream: + return Reply(body=json.dumps(body).encode()) + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "response.created", "sequence_number": 0, "response": {**body, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": text, + }, + {"type": "response.completed", "sequence_number": 2, "response": body}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(sse(event, string_value(event["type"])) for event in events), + pause_between_chunks=pause, + ) + + +def messages_usage(output_tokens: int) -> dict[str, JsonValue]: + return { + "input_tokens": PROMPT_TOKENS - CACHED_TOKENS - CACHE_WRITE_TOKENS, + "cache_read_input_tokens": CACHED_TOKENS, + "cache_creation_input_tokens": CACHE_WRITE_TOKENS, + "output_tokens": output_tokens, + } + + +def messages_reply(identity: str, model: str, text: str = ANSWER) -> Reply: + body: Final[dict[str, JsonValue]] = { + "id": identity, + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": messages_usage(COMPLETION_TOKENS), + } + return Reply(body=json.dumps(body).encode()) + + +def messages_stream_reply(identity: str, model: str, text: str = ANSWER, pause: float = 0) -> Reply: + events: Final[tuple[dict[str, JsonValue], ...]] = ( + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": messages_usage(1), + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": COMPLETION_TOKENS}, + }, + {"type": "message_stop"}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(sse(event, string_value(event["type"])) for event in events), + pause_between_chunks=pause, + ) + + +def error_reply(status: int, message: str) -> Reply: + body: Final[dict[str, JsonValue]] = {"error": {"message": message, "type": "scripted_error", "code": str(status)}} + return Reply(status=status, body=json.dumps(body).encode()) + + +def is_readiness_probe(request: Request) -> bool: + return bool(request.body) and READINESS_TEXT in request.body.decode(errors="replace") + + +def is_streaming(request: Request) -> bool: + return bool(request.body) and JSON_OBJECT.validate_json(request.body).get("stream") is True + + +def marker_of(request: Request) -> str | None: + found: Final = MARKER.search(request.body.decode(errors="replace")) + return found.group(1) if found is not None else None + + +def route(request: Request, identity: str, model: str, text: str = ANSWER, pause: float = 0) -> Reply: + match request.target: + case "/v1/chat/completions": + return ( + chat_stream_reply(identity, model, text, pause) + if is_streaming(request) + else chat_reply(identity, model, text) + ) + case "/v1/responses": + return responses_reply(identity, model, is_streaming(request), text, pause) + case "/v1/messages": + return ( + messages_stream_reply(identity, model, text, pause) + if is_streaming(request) + else messages_reply(identity, model, text) + ) + case _: + return error_reply(404, f"unknown route {request.target}") + + +def provider(identity: str, model: str = MODEL, stream_pause: float = 0) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if is_readiness_probe(request): + return chat_reply(f"probe_{uuid.uuid4().hex}", model) + return route(request, identity, model, pause=stream_pause) + + return respond + + +def marker_identity(request: Request, marker: str) -> str: + match request.target: + case "/v1/responses": + return f"resp_{marker}" + case "/v1/messages": + return f"msg_{marker}" + case _: + return f"req_{marker}" + + +def marker_provider(model: str = MODEL, stream_pause: float = 0) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + marker: Final = marker_of(request) + if marker is None: + return chat_reply(f"probe_{uuid.uuid4().hex}", model) + return route(request, marker_identity(request, marker), model, f"answer marker-{marker}", stream_pause) + + return respond + + +def failing_provider(status: int, message: str) -> Callable[[Request], Reply]: + return lambda _request: error_reply(status, message) + + +def gateway_v1(gateway: Gateway) -> str: + return f"{str(gateway.client.base_url).rstrip('/')}/v1" + + +def openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=gateway_v1(gateway), api_key=gateway.key, max_retries=0) + + +def async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=gateway_v1(gateway), api_key=gateway.key, max_retries=0) + + +def anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url).rstrip("/"), api_key=gateway.key, max_retries=0) + + +def async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url).rstrip("/"), api_key=gateway.key, max_retries=0 + ) + + +def not_yet_routable(response: object) -> bool: + status: Final = getattr(response, "status_code", None) + text: Final = getattr(response, "text", "") + return status == 400 and "Invalid model name" in text + + +def ready(gateway: Gateway, model: str, seconds: float = 60) -> None: + def probe() -> tuple[int, str]: + try: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": READINESS_TEXT}], "cache": NO_CACHE}, + ) + except httpx.TransportError as lost: + return 0, repr(lost) + return response.status_code, response.text + + def routable(status: int, text: str) -> bool: + return status != 0 and not (status == 400 and "Invalid model name" in text) + + eventually( + lambda: tuple(probe() for _ in range(2)), + lambda probes: all(routable(status, text) for status, text in probes), + seconds=seconds, + ) + + +def deployment(scenario: Scenario, wire: Wire, model: str = MODEL, **overrides: JsonValue) -> str: + alias: Final = scenario.model(model=f"{PROVIDER}/{model}", api_base=f"{wire.url}/v1", api_key=API_KEY, **overrides) + ready(scenario.gateway, alias) + return alias + + +def raw_deployment(scenario: Scenario, model_name: str, litellm_params: Mapping[str, JsonValue]) -> str: + created: Final = scenario.gateway.post( + "/model/new", {"model_name": model_name, "litellm_params": dict(litellm_params), "model_info": {}} + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return model_name + + +def closed_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return int(probe.getsockname()[1]) + + +def provider_calls(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if not is_readiness_probe(request)) + + +def only_call(wire: Wire, target: str, model: str = MODEL, key: str = API_KEY) -> dict[str, JsonValue]: + received: Final = provider_calls(wire) + assert [(request.method, request.target) for request in received] == [("POST", target)], [ + (request.method, request.target) for request in received + ] + assert received[0].headers.get("authorization") == f"Bearer {key}", received[0].headers + if target == "/v1/messages": + assert received[0].headers.get("anthropic-version"), received[0].headers + body: Final = JSON_OBJECT.validate_json(received[0].body) + assert body["model"] == model, body + assert not LEAK_FIELDS & body.keys(), sorted(LEAK_FIELDS & body.keys()) + return body + + +def spend_row(identity: str, seconds: float = 70) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE request_id=%s", (identity,)), + lambda found: len(found) == 1, + seconds=seconds, + ) + return rows[0] + + +def group_rows(model_group: str, count: int, seconds: float = 70) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE model_group=%s AND request_id NOT LIKE 'probe_%%'", (model_group,)), + lambda found: len(found) == count, + seconds=seconds, + ) + + +def assert_billed(row: Mapping[str, JsonValue], alias: str, call_type: str, model: str = MODEL) -> None: + assert row["model_group"] == alias, row + assert row["model"] == f"{PROVIDER}/{model}", row + assert row["custom_llm_provider"] == PROVIDER, row + assert row["status"] == "success", row + assert row["call_type"] == call_type, row + assert row["prompt_tokens"] == PROMPT_TOKENS and row["completion_tokens"] == COMPLETION_TOKENS, row + assert number(row["spend"]) == pytest.approx(expected_spend(model)), (row, expected_spend(model)) + + +def assert_failed(row: Mapping[str, JsonValue], alias: str, model: str = MODEL) -> None: + assert row["model_group"] == alias, row + assert row["model"] == f"{PROVIDER}/{model}", row + assert row["custom_llm_provider"] == PROVIDER, row + assert row["status"] == "failure", row + assert number(row["spend"]) == 0, row diff --git a/tests/integration/providers/test_coralbricks_provider_chaos.py b/tests/integration/providers/test_coralbricks_provider_chaos.py new file mode 100644 index 00000000000..dc4014fdb75 --- /dev/null +++ b/tests/integration/providers/test_coralbricks_provider_chaos.py @@ -0,0 +1,396 @@ +from __future__ import annotations + +import asyncio +import re +import threading +import uuid +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._coralbricks import ( + API_KEY, + MARKER, + MODEL, + NO_CACHE, + PROVIDER, + SPEND_COLUMNS, + deployment, + marker_of, + marker_provider, + ready, +) +from pydantic import JsonValue + +Endpoint = Literal["chat", "messages", "responses"] +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") + + +@dataclass(frozen=True, slots=True) +class Call: + endpoint: Endpoint + stream: bool + marker: str + + +@dataclass(frozen=True, slots=True) +class Served: + call: Call + status: int + text: str + call_id: str + + +def path_of(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def body_of(model: str, call: Call) -> dict[str, JsonValue]: + question: Final = f"Say hello marker-{call.marker}" + common: Final[dict[str, JsonValue]] = {"model": model, "stream": call.stream, "cache": NO_CACHE} + match call.endpoint: + case "chat": + return {**common, "messages": [{"role": "user", "content": question}]} + case "messages": + return {**common, "max_tokens": 64, "messages": [{"role": "user", "content": question}]} + case "responses": + return {**common, "input": question} + + +def calls_of(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[Call, ...]: + return tuple( + Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex) + for index in range(count) + ) + + +async def send(client: httpx.AsyncClient, key: str, model: str, call: Call) -> Served: + async with client.stream( + "POST", + path_of(call.endpoint), + json=body_of(model, call), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return Served( + call=call, status=response.status_code, text=raw.decode(), call_id=response.headers.get("x-litellm-call-id", "") + ) + + +async def burst( + base_url: str, key: str, model: str, calls: tuple[Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, Served)) + + +def assert_answered_with_its_own_marker(served: Served) -> None: + assert served.status == 200, served.text + assert set(MARKER.findall(served.text)) == {served.call.marker}, served.text + + +def expected_row_id(call: Call) -> str: + match call.endpoint: + case "chat": + return f"req_{call.marker}" + case "messages": + return f"msg_{call.marker}" + case "responses": + return f"resp_{call.marker}" + + +def landed_once(model_group: str, served: tuple[Served, ...]) -> None: + expected: Final = tuple(expected_row_id(item.call) for item in served) + rows: Final = eventually( + lambda: read_rows( + f"{SPEND_COLUMNS} WHERE model_group=%s AND status='success' AND request_id NOT LIKE 'probe_%%'", + (model_group,), + ), + lambda found: len(found) >= len(expected), + seconds=90, + ) + identities: Final = [string_value(row["request_id"]) for row in rows] + for identity in expected: + matches: Final = [found for found in identities if same_response(found, identity)] + assert len(matches) == 1, (identity, matches) + assert len(rows) == len(expected), (len(rows), len(expected)) + assert {row["model"] for row in rows} == {f"{PROVIDER}/{MODEL}"}, rows + assert {row["custom_llm_provider"] for row in rows} == {PROVIDER}, rows + + +def failed_once(model_group: str, served: tuple[Served, ...]) -> None: + call_ids: Final = tuple(item.call_id for item in served) + assert all(call_ids), [item.call_id for item in served] + rows: Final = eventually( + lambda: read_rows(f"{SPEND_COLUMNS} WHERE model_group=%s AND status='failure'", (model_group,)), + lambda found: {string_value(row["request_id"]) for row in found} >= set(call_ids), + seconds=90, + ) + identities: Final = [string_value(row["request_id"]) for row in rows] + assert sorted(identities) == sorted(call_ids), (sorted(identities), sorted(call_ids)) + + +def assert_no_bleed(received: tuple[Request, ...], markers: frozenset[str]) -> None: + forwarded: Final = [marker_of(request) for request in received if marker_of(request) is not None] + assert sorted(marker for marker in forwarded if marker is not None) == sorted(markers), forwarded + + +@contextmanager +def released_on_exit(release: threading.Event) -> Iterator[threading.Event]: + try: + yield release + finally: + release.set() + + +def held_provider(release: threading.Event, held: SimpleQueue[str]) -> Callable[[Request], Reply]: + respond: Final = marker_provider() + + def hold(request: Request) -> Reply: + marker: Final = marker_of(request) + if marker is None: + return respond(request) + held.put(marker) + assert release.wait(timeout=90), "The burst was never released" + return respond(request) + + return hold + + +def health_of(gateway: Gateway, model: str) -> tuple[int, int, int, str]: + response: Final = gateway.request("GET", "/health", params={"model": model}) + health: Final = JSON_OBJECT.validate_json(response.content) + healthy_count: Final = health["healthy_count"] + unhealthy_count: Final = health["unhealthy_count"] + assert isinstance(healthy_count, int) and isinstance(unhealthy_count, int), response.text + return response.status_code, healthy_count, unhealthy_count, response.text + + +def assert_unhealthy(gateway: Gateway, model: str) -> None: + observed: Final = eventually( + lambda: health_of(gateway, model), lambda found: found[:3] == (503, 0, 1), seconds=30, return_last_on_timeout=True + ) + assert observed[:3] == (503, 0, 1), observed[3] + + +def assert_healthy(gateway: Gateway, model: str) -> None: + observed: Final = eventually( + lambda: health_of(gateway, model), lambda found: found[:3] == (200, 1, 0), seconds=30, return_last_on_timeout=True + ) + assert observed[:3] == (200, 1, 0), observed[3] + + +@pytest.mark.timeout(400) +async def test_provider_outage_mid_burst_fails_fast_reports_unhealthy_and_recovers(gateway: Gateway) -> None: + before: Final = calls_of(15, ("chat", "messages", "responses"), lambda index: index % 2 == 0) + during: Final = calls_of(15, ("chat", "messages", "responses"), lambda index: index % 2 == 1) + after: Final = calls_of(15, ("chat", "messages", "responses"), lambda index: index % 2 == 0) + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + with gateway.scenario() as scenario: + with wire_server(held_provider(release, held)) as wire, released_on_exit(release): + port: Final = int(urlsplit(wire.url).port or 0) + alias: Final = deployment(scenario, wire) + wire.drain() + first: Final = asyncio.create_task(burst(str(gateway.client.base_url), gateway.key, alias, before)) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == 15, 60) + release.set() + served_before: Final = await first + assert len(served_before) == 15 + for item in served_before: + assert_answered_with_its_own_marker(item) + assert_no_bleed(wire.drain(), frozenset(call.marker for call in before)) + served_during: Final = await burst(str(gateway.client.base_url), gateway.key, alias, during) + assert len(served_during) == 15 + for item in served_during: + assert item.status == 500, item.text + assert "CoralbricksException" in item.text, item.text + assert "Connection error" in item.text or "Cannot connect to host" in item.text, item.text + assert "marker-" not in item.text, item.text + assert_unhealthy(gateway, alias) + with wire_server(marker_provider(), port=port) as restarted: + served_after: Final = await burst(str(gateway.client.base_url), gateway.key, alias, after) + assert len(served_after) == 15 + for item in served_after: + assert_answered_with_its_own_marker(item) + assert_no_bleed(restarted.drain(), frozenset(call.marker for call in after)) + assert_healthy(gateway, alias) + landed_once(alias, (*served_before, *served_after)) + failed_once(alias, served_during) + + +@pytest.mark.timeout(300) +async def test_slow_provider_streams_are_forwarded_once_under_concurrency(gateway: Gateway) -> None: + calls: Final = calls_of(20, ("chat", "messages", "responses"), lambda _: True) + with wire_server(marker_provider(stream_pause=0.3)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + wire.drain() + served: Final = await burst(str(gateway.client.base_url), gateway.key, alias, calls) + assert len(served) == 20 + for item in served: + assert_answered_with_its_own_marker(item) + if item.call.endpoint == "chat": + assert item.text.rstrip().endswith("data: [DONE]"), item.text + assert_no_bleed(wire.drain(), frozenset(call.marker for call in calls)) + landed_once(alias, served) + + +def chaos_model() -> str: + return f"coralbricks-chaos-{uuid.uuid4().hex[:12]}" + + +def chaos_config(wire: Wire, directory: Path, model_name: str) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **base, + "model_list": [ + { + "model_name": model_name, + "litellm_params": {"model": f"{PROVIDER}/{MODEL}", "api_base": f"{wire.url}/v1", "api_key": API_KEY}, + } + ], + } + path: Final = directory / "coralbricks-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def worker_pids(log: Path) -> tuple[int, ...]: + return eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(log.read_text())), + lambda pids: len(pids) == 2, + seconds=60, + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_coralbricks(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = calls_of(20, ("chat", "messages", "responses"), lambda _: False) + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + alias: Final = chaos_model() + with wire_server(held_provider(release, held)) as wire, released_on_exit(release): + config: Final = chaos_config(wire, tmp_path, alias) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = worker_pids(owned.log) + ready(candidate, alias, seconds=120) + wire.drain() + task: Final = asyncio.create_task( + burst(str(candidate.client.base_url), candidate.key, alias, calls, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await task + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + assert_answered_with_its_own_marker(item) + follow_up: Final = Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await burst(str(candidate.client.base_url), candidate.key, alias, (follow_up,)) + assert_answered_with_its_own_marker(answered) + assert_no_bleed(wire.drain(), frozenset(call.marker for call in (*calls, follow_up))) + landed_once(alias, (*served, answered)) + eventually( + lambda: owned.log.read_text().count("Application startup complete."), + lambda started: started == 3, + seconds=graceful_stop_seconds(), + ) + + +@pytest.mark.timeout(480) +async def test_proxy_sigterm_mid_burst_then_reboot_serves_coralbricks_again(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = calls_of(20, ("chat", "messages", "responses"), lambda index: index % 2 == 0) + release: Final = threading.Event() + held: Final[SimpleQueue[str]] = SimpleQueue() + alias: Final = chaos_model() + with wire_server(held_provider(release, held)) as wire, released_on_exit(release): + config: Final = chaos_config(wire, tmp_path, alias) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + worker_pids(owned.log) + ready(owned.gateway, alias, seconds=120) + wire.drain() + task: Final = asyncio.create_task( + burst( + str(owned.gateway.client.base_url), + owned.gateway.key, + alias, + calls, + tolerate_transport_errors=True, + ) + ) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == 20, 60) + owned.begin_shutdown() + release.set() + served: Final = await task + answered: Final = tuple(item for item in served if item.status == 200) + for item in answered: + assert_answered_with_its_own_marker(item) + await asyncio.to_thread(eventually, owned.process.poll, lambda code: code is not None, 90) + assert_no_bleed(wire.drain(), frozenset(call.marker for call in calls)) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as rebooted: + worker_pids(rebooted.log) + ready(rebooted.gateway, alias, seconds=120) + wire.drain() + follow_up: Final = Call(endpoint="responses", stream=True, marker=uuid.uuid4().hex) + (recovered,) = await burst(str(rebooted.gateway.client.base_url), rebooted.gateway.key, alias, (follow_up,)) + assert_answered_with_its_own_marker(recovered) + assert_no_bleed(wire.drain(), frozenset({follow_up.marker})) + rows: Final = eventually( + lambda: read_rows( + f"{SPEND_COLUMNS} WHERE model_group=%s AND status='success' AND request_id NOT LIKE 'probe_%%'", + (alias,), + ), + lambda found: any(same_response(string_value(row["request_id"]), expected_row_id(follow_up)) for row in found), + seconds=90, + ) + identities: Final = [string_value(row["request_id"]) for row in rows] + assert len(identities) == len(set(identities)), identities + allowed: Final = (*(expected_row_id(call) for call in calls), expected_row_id(follow_up)) + for identity in identities: + assert any(same_response(identity, candidate) for candidate in allowed), identity + billed: Final = { + item.call.marker: [found for found in identities if same_response(found, expected_row_id(item.call))] + for item in answered + } + assert all(len(found) == 1 for found in billed.values()), billed diff --git a/tests/integration/providers/test_coralbricks_provider_env_wire.py b/tests/integration/providers/test_coralbricks_provider_env_wire.py new file mode 100644 index 00000000000..6d8cc766b62 --- /dev/null +++ b/tests/integration/providers/test_coralbricks_provider_env_wire.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, string_value +from integration._support.process import owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Request, Wire, wire_server +from integration.providers._coralbricks import ( + API_KEY, + MODEL, + NO_CACHE, + PROVIDER, + assert_billed, + is_readiness_probe, + marker_provider, + not_yet_routable, + raw_deployment, + ready, + spend_row, +) +from pydantic import JsonValue + +ENV_KEY: Final = "cb_env_integration_key" +CONFIG_MODEL: Final = "coralbricks-env-config" + + +@dataclass(frozen=True, slots=True) +class EnvRig: + gateway: Gateway + env_wire: Wire + + +def env_config(env_wire: Wire, directory: Path) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **base, + "model_list": [{"model_name": CONFIG_MODEL, "litellm_params": {"model": f"{PROVIDER}/{MODEL}"}}], + "environment_variables": { + **base.get("environment_variables", {}), + "CORALBRICKS_API_KEY": ENV_KEY, + "CORALBRICKS_API_BASE": f"{env_wire.url}/v1", + }, + } + path: Final = directory / "coralbricks-env.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def env_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[EnvRig]: + directory: Final = tmp_path_factory.mktemp("coralbricks-env") + with gateway_from_environment() as rig_gateway, wire_server(marker_provider()) as env_wire: + config: Final = env_config(env_wire, directory) + with owned_proxy_process( + rig_gateway, directory, {}, remove_environment=("CORALBRICKS_API_KEY", "CORALBRICKS_API_BASE"), + config=config, workers=2, + ) as owned: + ready(owned.gateway, CONFIG_MODEL, seconds=120) + yield EnvRig(owned.gateway, env_wire) + + +def calls(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if not is_readiness_probe(request)) + + +def body_for(target: str, model: str, marker: str) -> dict[str, JsonValue]: + question: Final = f"Say hello marker-{marker}" + match target: + case "/v1/responses": + return {"model": model, "input": question, "cache": NO_CACHE} + case "/v1/messages": + return {"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": question}], "cache": NO_CACHE} + case _: + return {"model": model, "messages": [{"role": "user", "content": question}], "cache": NO_CACHE} + + +def only_call_with_key(wire: Wire, target: str, key: str) -> None: + received: Final = calls(wire) + assert [(request.method, request.target) for request in received] == [("POST", target)], [ + (request.method, request.target) for request in received + ] + assert received[0].headers.get("authorization") == f"Bearer {key}", received[0].headers + assert JSON_OBJECT.validate_json(received[0].body)["model"] == MODEL, received[0].body + + +def row_id(target: str, marker: str, served: dict[str, JsonValue]) -> str: + match target: + case "/v1/responses": + assert same_response(string_value(served["id"]), f"resp_{marker}"), served + return string_value(served["id"]) + case "/v1/messages": + assert served["id"] == f"msg_{marker}", served + return f"msg_{marker}" + case _: + assert served["id"] == f"req_{marker}", served + return f"req_{marker}" + + +@pytest.mark.parametrize( + ("target", "call_type"), + (("/v1/chat/completions", "acompletion"), ("/v1/responses", "aresponses"), ("/v1/messages", "anthropic_messages")), +) +def test_config_model_without_credentials_uses_the_env_key_and_base(env_rig: EnvRig, target: str, call_type: str) -> None: + marker: Final = uuid.uuid4().hex + env_rig.env_wire.drain() + served: Final = env_rig.gateway.post(target, body_for(target, CONFIG_MODEL, marker)) + only_call_with_key(env_rig.env_wire, target, ENV_KEY) + assert_billed(spend_row(row_id(target, marker, served)), CONFIG_MODEL, call_type) + + +def test_explicit_credentials_beat_the_env_credentials(env_rig: EnvRig) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(marker_provider()) as explicit_wire, env_rig.gateway.scenario() as scenario: + alias: Final = raw_deployment( + scenario, + f"integration-{uuid.uuid4().hex}", + {"model": f"{PROVIDER}/{MODEL}", "api_base": f"{explicit_wire.url}/v1", "api_key": API_KEY}, + ) + ready(env_rig.gateway, alias) + env_rig.env_wire.drain() + explicit_wire.drain() + served: Final = env_rig.gateway.post("/v1/chat/completions", body_for("/v1/chat/completions", alias, marker)) + only_call_with_key(explicit_wire, "/v1/chat/completions", API_KEY) + assert calls(env_rig.env_wire) == (), "the env base was called although the deployment names its own" + assert_billed(spend_row(row_id("/v1/chat/completions", marker, served)), alias, "acompletion") + + +@pytest.mark.parametrize("api_key", ("missing", None, ""), ids=("missing", "null", "empty")) +def test_deployment_without_a_key_uses_the_env_key(env_rig: EnvRig, api_key: str | None) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(marker_provider()) as explicit_wire, env_rig.gateway.scenario() as scenario: + alias: Final = raw_deployment( + scenario, + f"integration-{uuid.uuid4().hex}", + { + "model": f"{PROVIDER}/{MODEL}", + "api_base": f"{explicit_wire.url}/v1", + **({} if api_key == "missing" else {"api_key": api_key}), + }, + ) + ready(env_rig.gateway, alias) + explicit_wire.drain() + served: Final = env_rig.gateway.post("/v1/chat/completions", body_for("/v1/chat/completions", alias, marker)) + only_call_with_key(explicit_wire, "/v1/chat/completions", ENV_KEY) + assert_billed(spend_row(row_id("/v1/chat/completions", marker, served)), alias, "acompletion") + + +@pytest.mark.parametrize("api_base", ("missing", None, ""), ids=("missing", "null", "empty")) +def test_deployment_without_a_base_uses_the_env_base(env_rig: EnvRig, api_base: str | None) -> None: + marker: Final = uuid.uuid4().hex + with env_rig.gateway.scenario() as scenario: + alias: Final = raw_deployment( + scenario, + f"integration-{uuid.uuid4().hex}", + { + "model": f"{PROVIDER}/{MODEL}", + "api_key": API_KEY, + **({} if api_base == "missing" else {"api_base": api_base}), + }, + ) + ready(env_rig.gateway, alias) + env_rig.env_wire.drain() + served: Final = env_rig.gateway.post("/v1/chat/completions", body_for("/v1/chat/completions", alias, marker)) + only_call_with_key(env_rig.env_wire, "/v1/chat/completions", API_KEY) + assert_billed(spend_row(row_id("/v1/chat/completions", marker, served)), alias, "acompletion") + + +def test_env_credentials_are_not_echoed_by_model_info(env_rig: EnvRig) -> None: + entries: Final = env_rig.gateway.get("/model/info")["data"] + assert isinstance(entries, list) + mine: Final = [JSON_OBJECT.validate_python(entry) for entry in entries] + config_rows: Final = [entry for entry in mine if entry["model_name"] == CONFIG_MODEL] + assert len(config_rows) == 1, [entry["model_name"] for entry in mine] + assert ENV_KEY not in str(config_rows[0]), config_rows[0] + response: Final = eventually( + lambda: env_rig.gateway.request("GET", "/health", params={"model": CONFIG_MODEL}), + lambda candidate: not not_yet_routable(candidate), + seconds=30, + ) + assert response.status_code == 200, response.text + assert ENV_KEY not in response.text, response.text + health: Final = JSON_OBJECT.validate_json(response.content) + assert health["healthy_count"] == 1 and health["unhealthy_count"] == 0, health diff --git a/tests/integration/providers/test_coralbricks_provider_wire.py b/tests/integration/providers/test_coralbricks_provider_wire.py new file mode 100644 index 00000000000..d8231df5151 --- /dev/null +++ b/tests/integration/providers/test_coralbricks_provider_wire.py @@ -0,0 +1,612 @@ +from __future__ import annotations + +import uuid +from itertools import count +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, wire_server +from integration.providers._coralbricks import ( + ANSWER, + API_KEY, + COMPLETION_TOKENS, + MODEL, + MODELS, + NO_CACHE, + PROMPT_TOKENS, + PROVIDER, + USER_MESSAGES, + anthropic_client, + assert_billed, + assert_failed, + catalog_rows, + async_anthropic_client, + async_openai_client, + closed_port, + deployment, + expected_spend, + failing_provider, + group_rows, + is_readiness_probe, + not_yet_routable, + number, + only_call, + openai_client, + provider, + provider_calls, + rates_of, + raw_deployment, + row_of, + ready, + spend_row, +) +from pydantic import JsonValue + +DEPLOYMENT_RETRIES: Final = 2 + + +@pytest.mark.parametrize("model", MODELS) +def test_openai_sdk_chat_reaches_coralbricks_and_bills_the_cost_map_row(gateway: Gateway, model: str) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity, model)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire, model) + response: Final = openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.choices[0].message.content == ANSWER + sent: Final = only_call(wire, "/v1/chat/completions", model) + assert sent["messages"] == USER_MESSAGES, sent + assert_billed(spend_row(identity), alias, "acompletion", model) + + +async def test_async_openai_sdk_chat_reaches_coralbricks(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = await async_openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.choices[0].message.content == ANSWER + assert only_call(wire, "/v1/chat/completions")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "acompletion") + + +def test_openai_sdk_chat_stream_is_forwarded_and_billed_from_the_usage_frame(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + chunks: Final = list( + openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": NO_CACHE}, + ) + ) + assert {chunk.id for chunk in chunks} == {identity} + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == ANSWER + usage: Final = [chunk.usage for chunk in chunks if chunk.usage is not None] + assert len(usage) == 1 and usage[0].prompt_tokens == PROMPT_TOKENS, usage + sent: Final = only_call(wire, "/v1/chat/completions") + assert sent["stream"] is True, sent + assert_billed(spend_row(identity), alias, "acompletion") + + +async def test_async_openai_sdk_chat_stream_is_forwarded_and_billed(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + stream: Final = await async_openai_client(gateway).chat.completions.create( + model=alias, + messages=[{"role": "user", "content": "Say hello to the reef"}], + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": NO_CACHE}, + ) + chunks: Final = [chunk async for chunk in stream] + assert {chunk.id for chunk in chunks} == {identity} + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == ANSWER + assert only_call(wire, "/v1/chat/completions")["stream"] is True + assert_billed(spend_row(identity), alias, "acompletion") + + +def test_raw_chat_request_carries_the_call_id_and_lands_once(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert response.status_code == 200, response.text + assert "x-litellm-call-id" in response.headers, response.headers + body: Final = JSON_OBJECT.validate_json(response.content) + assert body["id"] == identity, response.text + assert only_call(wire, "/v1/chat/completions")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "acompletion") + + +def test_openai_sdk_responses_go_natively_to_coralbricks_responses(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", extra_body={"cache": NO_CACHE} + ) + assert response.output_text == ANSWER + assert same_response(response.id, identity), response.id + sent: Final = only_call(wire, "/v1/responses") + assert sent["input"] == "Say hello to the reef", sent + assert_billed(spend_row(response.id), alias, "aresponses") + + +async def test_async_openai_sdk_responses_go_natively_to_coralbricks_responses(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = await async_openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", extra_body={"cache": NO_CACHE} + ) + assert response.output_text == ANSWER + assert same_response(response.id, identity), response.id + assert only_call(wire, "/v1/responses")["input"] == "Say hello to the reef" + assert_billed(spend_row(response.id), alias, "aresponses") + + +def test_openai_sdk_responses_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + events: Final = list( + openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", stream=True, extra_body={"cache": NO_CACHE} + ) + ) + assert [event.type for event in events] == [ + "response.created", + "response.output_text.delta", + "response.completed", + ], [event.type for event in events] + completed: Final = events[-1] + assert completed.type == "response.completed" + assert completed.response.output_text == ANSWER + sent: Final = only_call(wire, "/v1/responses") + assert sent["stream"] is True and sent["input"] == "Say hello to the reef", sent + rows: Final = group_rows(alias, 1) + assert same_response(completed.response.id, string_value(rows[0]["request_id"])), rows + assert_billed(rows[0], alias, "aresponses") + + +async def test_async_openai_sdk_responses_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + stream: Final = await async_openai_client(gateway).responses.create( + model=alias, input="Say hello to the reef", stream=True, extra_body={"cache": NO_CACHE} + ) + events: Final = [event async for event in stream] + assert [event.type for event in events] == [ + "response.created", + "response.output_text.delta", + "response.completed", + ], [event.type for event in events] + completed: Final = events[-1] + assert completed.type == "response.completed" + assert completed.response.output_text == ANSWER + assert only_call(wire, "/v1/responses")["stream"] is True + rows: Final = group_rows(alias, 1) + assert same_response(completed.response.id, string_value(rows[0]["request_id"])), rows + assert_billed(rows[0], alias, "aresponses") + + +def test_raw_responses_request_lands_once(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = gateway.request( + "POST", "/v1/responses", {"model": alias, "input": "Say hello to the reef", "cache": NO_CACHE} + ) + assert response.status_code == 200, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + assert same_response(string_value(body["id"]), identity), response.text + assert only_call(wire, "/v1/responses")["input"] == "Say hello to the reef" + assert_billed(spend_row(string_value(body["id"])), alias, "aresponses") + + +def test_anthropic_sdk_messages_go_natively_to_coralbricks_messages(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = anthropic_client(gateway).messages.create( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.content[0].type == "text" and response.content[0].text == ANSWER + assert only_call(wire, "/v1/messages")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +async def test_async_anthropic_sdk_messages_go_natively_to_coralbricks_messages(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = await async_anthropic_client(gateway).messages.create( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) + assert response.id == identity + assert response.content[0].type == "text" and response.content[0].text == ANSWER + assert only_call(wire, "/v1/messages")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +def test_anthropic_sdk_messages_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + with anthropic_client(gateway).messages.stream( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) as stream: + text: Final = "".join(stream.text_stream) + final: Final = stream.get_final_message() + assert text == ANSWER + assert final.id == identity and final.usage.output_tokens == COMPLETION_TOKENS, final + assert only_call(wire, "/v1/messages")["stream"] is True + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +async def test_async_anthropic_sdk_messages_stream_is_forwarded_natively_and_billed_once(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + async with async_anthropic_client(gateway).messages.stream( + model=alias, + max_tokens=64, + messages=[{"role": "user", "content": "Say hello to the reef"}], + extra_body={"cache": NO_CACHE}, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert text == ANSWER + assert final.id == identity, final + assert only_call(wire, "/v1/messages")["stream"] is True + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +def test_raw_messages_request_lands_once(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": alias, "max_tokens": 64, "messages": USER_MESSAGES, "cache": NO_CACHE}, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["id"] == identity, response.text + assert only_call(wire, "/v1/messages")["messages"] == USER_MESSAGES + assert_billed(spend_row(identity), alias, "anthropic_messages") + + +def request_body(target: str, alias: str) -> dict[str, JsonValue]: + match target: + case "/v1/responses": + return {"model": alias, "input": "Say hello to the reef", "cache": NO_CACHE} + case "/v1/messages": + return {"model": alias, "max_tokens": 64, "messages": USER_MESSAGES, "cache": NO_CACHE} + case _: + return {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + + +@pytest.mark.parametrize( + ("target", "status"), + (("/v1/chat/completions", 401), ("/v1/chat/completions", 404), ("/v1/responses", 401), ("/v1/messages", 401)), +) +def test_provider_client_error_reaches_the_caller_after_one_attempt(gateway: Gateway, target: str, status: int) -> None: + message: Final = f"scripted coralbricks {status} {uuid.uuid4().hex}" + with wire_server(failing_provider(status, message)) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model(model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY) + probe: Final = eventually( + lambda: gateway.request("POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES}), + lambda response: not not_yet_routable(response), + seconds=60, + ) + assert probe.status_code == status, probe.text + wire.drain() + response: Final = gateway.request("POST", target, request_body(target, alias)) + assert response.status_code == status, response.text + assert message in response.text, response.text + attempts: Final = [(request.method, request.target) for request in wire.drain()] + assert attempts == [("POST", target)], attempts + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize("status", (429, 500)) +def test_provider_server_error_is_retried_then_reaches_the_caller(gateway: Gateway, status: int) -> None: + message: Final = f"scripted coralbricks {status} {uuid.uuid4().hex}" + with wire_server(failing_provider(status, message)) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model( + model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY, num_retries=DEPLOYMENT_RETRIES + ) + eventually( + lambda: gateway.request("POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES}), + lambda response: not not_yet_routable(response), + seconds=60, + ) + wire.drain() + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert response.status_code == status, response.text + assert message in response.text, response.text + attempts: Final = [(request.method, request.target) for request in wire.drain()] + assert attempts == [("POST", "/v1/chat/completions")] * (1 + DEPLOYMENT_RETRIES), attempts + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + + +def test_unreachable_api_base_reports_the_connection_error_and_leaves_the_proxy_serving(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + control: Final = scenario.model() + alias: Final = scenario.model( + model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{closed_port()}/v1", api_key=API_KEY + ) + response: Final = eventually( + lambda: gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 500, response.text + assert "Connection error" in response.text, response.text + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + ready(gateway, control) + assert gateway.chat(control, text="still serving")["object"] == "chat.completion" + + +@pytest.mark.parametrize("api_key", ("missing", None, ""), ids=("missing", "null", "empty")) +def test_deployment_without_a_key_and_no_env_key_fails_before_the_provider( + gateway: Gateway, api_key: str | None +) -> None: + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + litellm_params: Final[dict[str, JsonValue]] = { + "model": f"{PROVIDER}/{MODEL}", + "api_base": f"{wire.url}/v1", + **({} if api_key == "missing" else {"api_key": api_key}), + } + alias: Final = raw_deployment(scenario, f"integration-{uuid.uuid4().hex}", litellm_params) + response: Final = eventually( + lambda: gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 500, response.text + assert "AuthenticationError" in response.text and "CORALBRICKS_API_KEY" in response.text, response.text + assert_failed(spend_row(response.headers["x-litellm-call-id"]), alias) + assert wire.drain() == (), "the provider was called without a key" + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize( + "model", + (7, ["coralbricks/deepseek-v4.1-flash-fast"], f"{PROVIDER}/" + "g" * 5120), + ids=("int", "list", "5kb"), +) +def test_hostile_model_field_is_rejected_without_crashing(gateway: Gateway, model: JsonValue) -> None: + for _ in range(2): + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert response.status_code in (400, 422), response.text + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def test_identical_uncached_chats_are_each_forwarded_and_billed(gateway: Gateway) -> None: + prefix: Final = f"req_{uuid.uuid4().hex}" + calls: Final = count(1) + + def respond(request: Request) -> Reply: + if is_readiness_probe(request): + return provider(f"{prefix}-probe")(request) + return provider(f"{prefix}-{next(calls)}")(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + wire.drain() + first: Final = gateway.post( + "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + second: Final = gateway.post( + "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE} + ) + assert first["id"] != second["id"], (first["id"], second["id"]) + assert [(request.method, request.target) for request in provider_calls(wire)] == [ + ("POST", "/v1/chat/completions") + ] * 2 + rows: Final = group_rows(alias, 2) + assert {string_value(row["request_id"]) for row in rows} == {string_value(first["id"]), string_value(second["id"])} + for row in rows: + assert_billed(row, alias, "acompletion") + + +def test_model_info_reports_the_coralbricks_cost_map_row(gateway: Gateway) -> None: + rates: Final = rates_of(MODEL) + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model(model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY) + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + mine: Final = [object_value(entry) for entry in entries if object_value(entry)["model_name"] == alias] + assert len(mine) == 1, [object_value(entry)["model_name"] for entry in entries] + litellm_params: Final = object_value(mine[0]["litellm_params"]) + info: Final = object_value(mine[0]["model_info"]) + assert litellm_params["model"] == f"{PROVIDER}/{MODEL}", litellm_params + assert info["litellm_provider"] == PROVIDER, info + assert number(info["input_cost_per_token"]) == pytest.approx(rates.input), info + assert number(info["output_cost_per_token"]) == pytest.approx(rates.output), info + assert number(info["cache_creation_input_token_cost"]) == pytest.approx(rates.cache_write), info + assert info["max_input_tokens"] == row_of(MODEL)["max_input_tokens"] and info["mode"] == "chat", info + assert info["supports_function_calling"] is True and info["supports_prompt_caching"] is True, info + + +def test_models_listing_and_model_group_info_carry_the_coralbricks_deployment(gateway: Gateway) -> None: + rates: Final = rates_of(MODEL) + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + alias: Final = scenario.model(model=f"{PROVIDER}/{MODEL}", api_base=f"{wire.url}/v1", api_key=API_KEY) + listed: Final = gateway.get("/v1/models")["data"] + assert isinstance(listed, list) + assert alias in {object_value(entry)["id"] for entry in listed} + groups: Final = gateway.get("/model_group/info")["data"] + assert isinstance(groups, list) + mine: Final = [object_value(group) for group in groups if object_value(group)["model_group"] == alias] + assert len(mine) == 1, [object_value(group)["model_group"] for group in groups] + assert mine[0]["providers"] == [PROVIDER], mine[0] + assert number(mine[0]["input_cost_per_token"]) == pytest.approx(rates.input), mine[0] + assert number(mine[0]["output_cost_per_token"]) == pytest.approx(rates.output), mine[0] + + +def test_wildcard_deployment_lists_every_coralbricks_catalog_row(gateway: Gateway) -> None: + with wire_server(provider("unused")) as wire, gateway.scenario() as scenario: + created: Final = gateway.post( + "/model/new", + { + "model_name": f"{PROVIDER}/*", + "litellm_params": {"model": f"{PROVIDER}/*", "api_base": f"{wire.url}/v1", "api_key": API_KEY}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + catalog: Final = catalog_rows() + listed: Final = gateway.get("/v1/models")["data"] + assert isinstance(listed, list) + listed_names: Final = {object_value(entry)["id"] for entry in listed} + assert catalog <= listed_names, sorted(listed_names) + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + expanded: Final = { + object_value(entry)["model_name"]: object_value(entry) + for entry in entries + if object_value(entry)["model_name"] in catalog + } + assert set(expanded) == catalog, [object_value(entry)["model_name"] for entry in entries] + for name, entry in expanded.items(): + assert object_value(entry["litellm_params"])["model"] == name, entry + assert object_value(entry["model_info"])["litellm_provider"] == PROVIDER, entry + + +def test_health_probe_reaches_coralbricks_with_the_deployment_key(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + wire.drain() + health: Final = gateway.get("/health", {"model": alias}) + assert health["healthy_count"] == 1 and health["unhealthy_count"] == 0, health + healthy: Final = health["healthy_endpoints"] + assert isinstance(healthy, list) and object_value(healthy[0])["model"] == f"{PROVIDER}/{MODEL}", health + probes: Final = provider_calls(wire) + assert [(request.method, request.target) for request in probes] == [("POST", "/v1/chat/completions")], probes + assert probes[0].headers.get("authorization") == f"Bearer {API_KEY}", probes[0].headers + assert JSON_OBJECT.validate_json(probes[0].body)["model"] == MODEL + + +def test_add_model_form_fields_name_the_coralbricks_credentials(gateway: Gateway) -> None: + response: Final = gateway.request("GET", "/public/providers/fields") + assert response.status_code == 200, response.text + providers: Final = response.json() + assert isinstance(providers, list), response.text + mine: Final = [object_value(entry) for entry in providers if object_value(entry)["litellm_provider"] == PROVIDER] + assert len(mine) == 1, [object_value(entry)["litellm_provider"] for entry in providers] + assert mine[0]["provider_display_name"] == "CoralBricks", mine[0] + fields: Final = mine[0]["credential_fields"] + assert isinstance(fields, list), mine[0] + by_key: Final = {string_value(object_value(field)["key"]): object_value(field) for field in fields} + assert set(by_key) == {"api_key", "api_base"}, sorted(by_key) + assert by_key["api_key"]["required"] is True and by_key["api_base"]["required"] is False, by_key + assert string_value(by_key["api_base"]["placeholder"]).endswith("/v1"), by_key["api_base"] + + +def test_key_scoped_to_the_coralbricks_deployment_is_charged_its_spend(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + token: Final = scenario.key(models=[alias], max_budget=1.0) + wire.drain() + response: Final = gateway.post( + "/v1/chat/completions", {"model": alias, "messages": USER_MESSAGES, "cache": NO_CACHE}, key=token + ) + assert response["id"] == identity, response + only_call(wire, "/v1/chat/completions") + assert_billed(spend_row(identity), alias, "acompletion") + info: Final = eventually( + lambda: object_value(gateway.get("/key/info", {"key": token})["info"]), + lambda found: number(found["spend"]) > 0, + seconds=70, + ) + assert number(info["spend"]) == pytest.approx(expected_spend(MODEL)), info + + +def test_wildcard_coralbricks_deployment_routes_the_requested_model(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + prefix: Final = f"integration-{uuid.uuid4().hex}" + raw_deployment( + scenario, + f"{prefix}/*", + {"model": f"{PROVIDER}/*", "api_base": f"{wire.url}/v1", "api_key": API_KEY}, + ) + requested: Final = f"{prefix}/{MODEL}" + response: Final = eventually( + lambda: gateway.request( + "POST", "/v1/chat/completions", {"model": requested, "messages": USER_MESSAGES, "cache": NO_CACHE} + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["id"] == identity, response.text + only_call(wire, "/v1/chat/completions") + row: Final = spend_row(identity) + assert row["model"] == f"{PROVIDER}/{MODEL}" and row["custom_llm_provider"] == PROVIDER, row + assert number(row["spend"]) == pytest.approx(expected_spend(MODEL)), row + + +def test_fallback_from_a_refused_primary_lands_on_coralbricks(gateway: Gateway) -> None: + identity: Final = f"req_{uuid.uuid4().hex}" + with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: + alias: Final = deployment(scenario, wire) + primary: Final = scenario.model( + model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{closed_port()}/v1", api_key=API_KEY + ) + wire.drain() + response: Final = eventually( + lambda: gateway.request( + "POST", + "/v1/chat/completions", + {"model": primary, "messages": USER_MESSAGES, "cache": NO_CACHE, "fallbacks": [alias]}, + ), + lambda candidate: not not_yet_routable(candidate), + seconds=60, + ) + assert response.status_code == 200, response.text + assert response.headers.get("x-litellm-attempted-fallbacks") == "1", response.headers + assert JSON_OBJECT.validate_json(response.content)["id"] == identity, response.text + only_call(wire, "/v1/chat/completions") + assert_billed(spend_row(identity), alias, "acompletion") diff --git a/tests/unit/llms/openai_like/test_coralbricks_provider.py b/tests/unit/llms/openai_like/test_coralbricks_provider.py new file mode 100644 index 00000000000..6187c7aa466 --- /dev/null +++ b/tests/unit/llms/openai_like/test_coralbricks_provider.py @@ -0,0 +1,539 @@ +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import respx + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache + +CORALBRICKS_MODEL: Final = "coralbricks/deepseek-v4.1-flash-fast" +CORALBRICKS_MODEL_ID: Final = "deepseek-v4.1-flash-fast" + + +def test_coralbricks_provider_resolution(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("CORALBRICKS_API_KEY", "coralbricks-test-key") + + model, provider, api_key, api_base = get_llm_provider( + model="coralbricks/glm-5.3-fast", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "glm-5.3-fast" + assert provider == "coralbricks" + assert api_key == "coralbricks-test-key" + assert api_base == "https://inference.coralbricks.ai/v1" + + +def test_coralbricks_provider_keeps_explicit_credentials(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("CORALBRICKS_API_KEY", "coralbricks-env-key") + + _, provider, api_key, api_base = get_llm_provider( + model="coralbricks/glm-5.3-fast", + custom_llm_provider=None, + api_base="https://coralbricks.internal.example/v1", + api_key="coralbricks-explicit-key", + ) + + assert provider == "coralbricks" + assert api_key == "coralbricks-explicit-key" + assert api_base == "https://coralbricks.internal.example/v1" + + +def test_coralbricks_is_available_in_add_model_form(): + fields_path: Final = ( + Path(litellm.__file__).parent / "proxy" / "public_endpoints" / "provider_create_fields.json" + ) + providers: Final = json.loads(fields_path.read_text()) + coralbricks: Final = next( + provider for provider in providers if provider["litellm_provider"] == "coralbricks" + ) + + assert coralbricks["provider"] == "CORALBRICKS" + assert coralbricks["provider_display_name"] == "CoralBricks" + assert coralbricks["default_model_placeholder"] == "coralbricks/glm-5.3-fast" + assert {field["key"]: field["required"] for field in coralbricks["credential_fields"]} == { + "api_base": False, + "api_key": True, + } + + +def test_coralbricks_endpoint_support_matches_providers_json(): + matrix_path: Final = Path(litellm.__file__).parent / "provider_endpoints_support_backup.json" + endpoints: Final = json.loads(matrix_path.read_text())["providers"]["coralbricks"]["endpoints"] + providers_json: Final = json.loads( + (Path(litellm.__file__).parent / "llms" / "openai_like" / "providers.json").read_text() + ) + declared: Final = set(providers_json["coralbricks"]["supported_endpoints"]) + + supported_from_matrix: Final = { + "/v1/chat/completions": endpoints["chat_completions"], + "/v1/responses": endpoints["responses"], + "/v1/messages": endpoints["messages"], + } + assert {endpoint for endpoint, supported in supported_from_matrix.items() if supported} == declared + + +def test_coralbricks_chat_completion_request(): + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_coralbricks", + "object": "chat.completion", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello from CoralBricks"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.completion( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://inference.coralbricks.ai/v1/chat/completions" + assert request.headers["authorization"] == "Bearer coralbricks-test-key" + assert body["model"] == CORALBRICKS_MODEL_ID + assert body["messages"] == [{"role": "user", "content": "Say hello"}] + assert response.choices[0].message.content == "Hello from CoralBricks" + + +def test_coralbricks_responses_request(): + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/responses").respond( + 200, + json={ + "id": "resp_coralbricks", + "object": "response", + "created_at": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "status": "completed", + "output": [ + { + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + {"type": "output_text", "text": "Hello from CoralBricks", "annotations": []} + ], + } + ], + "usage": {"input_tokens": 4, "output_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.responses( + model=CORALBRICKS_MODEL, + input="Say hello", + api_key="coralbricks-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://inference.coralbricks.ai/v1/responses" + assert request.headers["authorization"] == "Bearer coralbricks-test-key" + assert body["model"] == CORALBRICKS_MODEL_ID + assert body["input"] == "Say hello" + assert response.output[0].content[0].text == "Hello from CoralBricks" + + +@pytest.mark.asyncio +async def test_coralbricks_anthropic_messages_request(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/messages").respond( + 200, + json={ + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "model": CORALBRICKS_MODEL_ID, + "content": [{"type": "text", "text": "Hello from CoralBricks"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 4, "output_tokens": 3}, + }, + ) + response: Final = await litellm.anthropic.messages.acreate( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=32, + api_key="coralbricks-test-key", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert str(request.url) == "https://inference.coralbricks.ai/v1/messages" + assert request.headers["authorization"] == "Bearer coralbricks-test-key" + assert request.headers["anthropic-version"] == "2023-06-01" + assert body["model"] == CORALBRICKS_MODEL_ID + assert body["messages"] == [{"role": "user", "content": "Say hello"}] + assert response["content"][0]["text"] == "Hello from CoralBricks" + + +@dataclass(frozen=True) +class CoralBricksRates: + input: float + output: float + cache_write: float + cache_read: float + + +def _coralbricks_rates() -> CoralBricksRates: + row: Final = litellm.model_cost[CORALBRICKS_MODEL] + assert row["litellm_provider"] == "coralbricks" + return CoralBricksRates( + input=float(row["input_cost_per_token"]), + output=float(row["output_cost_per_token"]), + cache_write=float(row["cache_creation_input_token_cost"]), + cache_read=float(row["cache_read_input_token_cost"]), + ) + + +def _expected_cost( + rates: CoralBricksRates, + uncached_input_tokens: int, + cached_read_tokens: int, + cache_write_tokens: int, + billed_output_tokens: int, +) -> float: + return ( + uncached_input_tokens * rates.input + + cached_read_tokens * rates.cache_read + + cache_write_tokens * rates.cache_write + + billed_output_tokens * rates.output + ) + + +def test_coralbricks_chat_completion_cost_includes_cache_write_and_free_cached_read(): + rates: Final = _coralbricks_rates() + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_coralbricks", + "object": "chat.completion", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 4357, + "completion_tokens": 20, + "total_tokens": 4377, + "prompt_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + "cache_write_blocks": 3, + }, + "completion_tokens_details": {"reasoning_tokens": 12}, + }, + }, + ) + response: Final = litellm.completion( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + ) + + cost: Final = litellm.completion_cost(completion_response=response, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=4357 - 4352 - 5, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) + + +def test_coralbricks_responses_cost_includes_cache_write_and_free_cached_read(): + rates: Final = _coralbricks_rates() + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/responses").respond( + 200, + json={ + "id": "resp_coralbricks", + "object": "response", + "created_at": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "status": "completed", + "output": [ + { + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": { + "input_tokens": 4357, + "output_tokens": 20, + "total_tokens": 4377, + "input_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + }, + "output_tokens_details": {"reasoning_tokens": 12}, + }, + }, + ) + response: Final = litellm.responses( + model=CORALBRICKS_MODEL, + input="Say hello", + api_key="coralbricks-test-key", + ) + + cost: Final = litellm.completion_cost(completion_response=response, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=4357 - 4352 - 5, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) + + +@pytest.mark.asyncio +async def test_coralbricks_anthropic_messages_cost_includes_cache_write_and_free_cached_read( + monkeypatch: pytest.MonkeyPatch, +): + rates: Final = _coralbricks_rates() + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/messages").respond( + 200, + json={ + "id": "msg_coralbricks", + "type": "message", + "role": "assistant", + "model": CORALBRICKS_MODEL_ID, + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 0, + "cache_read_input_tokens": 4352, + "cache_creation_input_tokens": 5, + "output_tokens": 20, + }, + }, + ) + response: Final = await litellm.anthropic.messages.acreate( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=32, + api_key="coralbricks-test-key", + ) + + cost: Final = litellm.completion_cost(completion_response=response, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=0, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) + + +def test_coralbricks_rows_are_listed_for_wildcard_expansion(): + catalog: Final = {name for name, info in litellm.model_cost.items() if info.get("litellm_provider") == "coralbricks"} + assert CORALBRICKS_MODEL in catalog + assert set(litellm.get_valid_models(custom_llm_provider="coralbricks")) == catalog + + +def test_coralbricks_cost_map_backup_mirrors_main(): + main_map: Final = json.loads( + (Path(litellm.__file__).parent.parent / "model_prices_and_context_window.json").read_text() + ) + backup_map: Final = json.loads( + (Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json").read_text() + ) + main_keys: Final = {key for key in main_map if key.startswith("coralbricks/")} + backup_keys: Final = {key for key in backup_map if key.startswith("coralbricks/")} + assert main_keys and main_keys == backup_keys + for key in main_keys: + assert backup_map[key] == main_map[key] + + +def test_coralbricks_cost_map_rows_declare_same_endpoints_as_providers_json(): + providers_json: Final = json.loads( + (Path(litellm.__file__).parent / "llms" / "openai_like" / "providers.json").read_text() + ) + declared: Final = providers_json["coralbricks"]["supported_endpoints"] + cost_map_paths: Final = ( + Path(litellm.__file__).parent.parent / "model_prices_and_context_window.json", + Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json", + ) + for cost_map_path in cost_map_paths: + cost_map: Final = json.loads(cost_map_path.read_text()) + coralbricks_rows: Final = { + key: row for key, row in cost_map.items() if key.startswith("coralbricks/") + } + assert coralbricks_rows + for row in coralbricks_rows.values(): + assert row["supported_endpoints"] == declared + + +def _billed_chat_cost(model_id: str) -> float: + with respx.mock() as upstream: + upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl_coralbricks", + "object": "chat.completion", + "created": 1_789_550_000, + "model": model_id, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 4400, + "completion_tokens": 20, + "total_tokens": 4420, + "prompt_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + }, + }, + }, + ) + response: Final = litellm.completion( + model=f"coralbricks/{model_id}", + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + ) + return litellm.completion_cost(completion_response=response, model=f"coralbricks/{model_id}") + + +@pytest.mark.parametrize( + ("deprecated_slug", "current_slug"), + ( + ("glm-5.3-fp4", "glm-5.3-fast"), + ("glm-5.3", "glm-5.3-fast"), + ("deepseek-v4.1-flash-fast-fp4", "deepseek-v4.1-flash-fast"), + ("deepseek-v4.1-flash", "deepseek-v4.1-flash-fast"), + ), +) +def test_coralbricks_deprecated_slug_bills_like_its_current_slug(deprecated_slug: str, current_slug: str): + current_cost: Final = _billed_chat_cost(current_slug) + assert current_cost > 0 + assert _billed_chat_cost(deprecated_slug) == pytest.approx(current_cost) + + +def test_coralbricks_streaming_chat_cost_includes_cache_write_and_free_cached_read(): + rates: Final = _coralbricks_rates() + usage_chunk: Final = { + "id": "chatcmpl_coralbricks", + "object": "chat.completion.chunk", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [], + "usage": { + "prompt_tokens": 4357, + "completion_tokens": 20, + "total_tokens": 4377, + "prompt_tokens_details": { + "cached_tokens": 4352, + "cache_write_tokens": 5, + "billable_cache_write_tokens": 5, + "cache_write_blocks": 3, + }, + "completion_tokens_details": {"reasoning_tokens": 12}, + }, + } + stream_body: Final = ( + "data: " + + json.dumps( + { + "id": "chatcmpl_coralbricks", + "object": "chat.completion.chunk", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}], + } + ) + + "\n\n" + + "data: " + + json.dumps( + { + "id": "chatcmpl_coralbricks", + "object": "chat.completion.chunk", + "created": 1_789_550_000, + "model": CORALBRICKS_MODEL_ID, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + ) + + "\n\n" + + "data: " + + json.dumps(usage_chunk) + + "\n\n" + + "data: [DONE]\n\n" + ) + with respx.mock() as upstream: + route: Final = upstream.post("https://inference.coralbricks.ai/v1/chat/completions").respond( + 200, headers={"content-type": "text/event-stream"}, content=stream_body + ) + response: Final = litellm.completion( + model=CORALBRICKS_MODEL, + messages=[{"role": "user", "content": "Say hello"}], + api_key="coralbricks-test-key", + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = list(response) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert route.call_count == 1 + assert body["stream"] is True + + built: Final = litellm.stream_chunk_builder( + chunks, messages=[{"role": "user", "content": "Say hello"}] + ) + assert built.usage.prompt_tokens_details.cached_tokens == 4352 + cost: Final = litellm.completion_cost(completion_response=built, model=CORALBRICKS_MODEL) + assert cost == pytest.approx( + _expected_cost( + rates, + uncached_input_tokens=0, + cached_read_tokens=4352, + cache_write_tokens=5, + billed_output_tokens=20, + ) + ) diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 52e13642020..b27be2ba966 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -201,6 +201,7 @@ describe("provider_info_helpers", () => { Providers.CLARIFAI, Providers.Cognition, Providers.COMPACTIFAI, + Providers.CORALBRICKS, Providers.DATAROBOT, Providers.DOCKER_MODEL_RUNNER, Providers.DOTPROMPT, @@ -361,6 +362,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Tencent)).toBe("tencent/deepseek-v4-pro"); }); + it("should return a coralbricks/ placeholder for the CoralBricks provider", () => { + expect(getPlaceholder(Providers.CORALBRICKS)).toBe("coralbricks/glm-5.3-fast"); + }); + it("should return decision model placeholders for the TypeSafe and StrandsDecider dropdown keys", () => { expect(getPlaceholder("TypeSafe")).toBe("typesafe/jev-latest"); expect(getPlaceholder("StrandsDecider")).toBe("strands_decider/strands-decider-2B-hobson-v19"); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa774d46f0f..7b6c7d24706 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -99,6 +99,7 @@ export enum Providers { COHERE_CHAT = "Cohere Chat", COMETAPI = "Cometapi", COMPACTIFAI = "Compactifai", + CORALBRICKS = "CoralBricks", Cursor = "Cursor", Dashscope = "Dashscope", Databricks = "Databricks (Qwen API)", @@ -221,6 +222,7 @@ export const provider_map: Record = { COHERE_CHAT: "cohere_chat", COMETAPI: "cometapi", COMPACTIFAI: "compactifai", + CORALBRICKS: "coralbricks", Cursor: "cursor", Dashscope: "dashscope", Databricks: "databricks", @@ -459,6 +461,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Bedrock]: "claude-3-opus", [Providers.CHATGPT]: "chatgpt/gpt-5.4", [Providers.Cognition]: "cognition/swe-1.7", + [Providers.CORALBRICKS]: "coralbricks/glm-5.3-fast", [Providers.Cursor]: "cursor/claude-4-sonnet", [Providers.DeepInfra]: "deepinfra/", [Providers.EDENAI]: "edenai/openai/gpt-mini-latest",