From bcb9d16eaa9860e15e6a2d1edf5370fc6995edba Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 8 Aug 2026 01:47:25 +0000 Subject: [PATCH] feat(cognition): add a distinct provider identity for Cognition Register cognition as a JSON-configured OpenAI-compatible provider so its traffic no longer rides on custom_llm_provider=openai, keeping OpenAI cost discounts and provider-level reporting out of it. Adds the SWE cost map entries, the supported-endpoints matrix row, and the Admin UI provider option. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- README.md | 1 + litellm/constants.py | 2 + .../get_llm_provider_logic.py | 3 + litellm/llms/openai_like/providers.json | 5 + ...odel_prices_and_context_window_backup.json | 31 ++++ .../provider_endpoints_support_backup.json | 17 ++ litellm/types/utils.py | 1 + model_prices_and_context_window.json | 31 ++++ provider_endpoints_support.json | 17 ++ .../openai_like/test_cognition_provider.py | 159 ++++++++++++++++++ .../src/components/provider_info_helpers.tsx | 4 + 11 files changed, 271 insertions(+) create mode 100644 tests/test_litellm/llms/openai_like/test_cognition_provider.py diff --git a/README.md b/README.md index 32b0160dbaa..68aaa09ec98 100644 --- a/README.md +++ b/README.md @@ -292,6 +292,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ | [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | | | [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | | | | | | | | | [Codestral (`codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | | +| [Cognition (`cognition`)](https://docs.litellm.ai/docs/providers/cognition) | ✅ | ✅ | ✅ | | | | | | | | | [Cohere (`cohere`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | | [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | | | [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | | diff --git a/litellm/constants.py b/litellm/constants.py index 6f0e9e7afe2..6a32ead4c7a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -751,6 +751,7 @@ openai_compatible_endpoints: Final[list] = [ "https://api.libertai.io/v1", "https://pinstripes.io/v1", "https://api.meta.ai/v1", + "https://api.cognition.ai/v1", ] @@ -818,6 +819,7 @@ openai_compatible_providers: Final[list] = [ "pinstripes", # Pinstripes - JSON-configured provider "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider + "cognition", # Cognition - JSON-configured provider ] openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions` "together_ai", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index dbb40913e14..fe7b10532d4 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -349,6 +349,9 @@ def get_llm_provider( elif endpoint == "https://api.meta.ai/v1": custom_llm_provider = "meta" dynamic_api_key = get_secret_str("META_API_KEY") + elif endpoint == "https://api.cognition.ai/v1": + custom_llm_provider = "cognition" + dynamic_api_key = get_secret_str("COGNITION_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception(f"api base needs to be a string. api_base={api_base}") diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 164100d4194..5f57aaa78d8 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -175,6 +175,11 @@ "base_class": "openai_gpt", "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] }, + "cognition": { + "base_url": "https://api.cognition.ai/v1", + "api_key_env": "COGNITION_API_KEY", + "api_base_env": "COGNITION_API_BASE" + }, "pinstripes": { "base_url": "https://pinstripes.io/v1", "api_key_env": "PINSTRIPES_API_KEY", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d954e33da9c..6b64e35474f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -46247,6 +46247,37 @@ "supports_reasoning": true, "supports_vision": false }, + "cognition/swe-1.5": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "cache_creation_input_token_cost": 3.75e-07, + "cache_read_input_token_cost": 3e-08, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, + "cognition/swe-1.6": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 3e-08, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, + "cognition/swe-1.7": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, "pinstripes/ps/glm-4.5-air": { "max_tokens": 128000, "max_input_tokens": 128000, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index dd7712aabca..d47d74ead28 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -528,6 +528,23 @@ "interactions": true } }, + "cognition": { + "display_name": "Cognition (`cognition`)", + "url": "https://docs.litellm.ai/docs/providers/cognition", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "cohere": { "display_name": "Cohere (`cohere`)", "url": "https://docs.litellm.ai/docs/providers/cohere", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 35d4250782f..6e48584a8d3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3695,6 +3695,7 @@ class LlmProviders(str, Enum): TENSORMESH = "tensormesh" LIBERTAI = "libertai" PINSTRIPES = "pinstripes" + COGNITION = "cognition" DARKBLOOM = "darkbloom" META = "meta" LITELLM_AGENT = "litellm_agent" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8982b4f2565..7c40efc89c7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -46335,6 +46335,37 @@ "supports_reasoning": true, "supports_vision": false }, + "cognition/swe-1.5": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "cache_creation_input_token_cost": 3.75e-07, + "cache_read_input_token_cost": 3e-08, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, + "cognition/swe-1.6": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 3e-08, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, + "cognition/swe-1.7": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, "pinstripes/ps/glm-4.5-air": { "max_tokens": 128000, "max_input_tokens": 128000, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 65db63dc045..d16484a8acd 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -563,6 +563,23 @@ "interactions": true } }, + "cognition": { + "display_name": "Cognition (`cognition`)", + "url": "https://docs.litellm.ai/docs/providers/cognition", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "cohere": { "display_name": "Cohere (`cohere`)", "url": "https://docs.litellm.ai/docs/providers/cohere", diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/test_litellm/llms/openai_like/test_cognition_provider.py new file mode 100644 index 00000000000..68758b70363 --- /dev/null +++ b/tests/test_litellm/llms/openai_like/test_cognition_provider.py @@ -0,0 +1,159 @@ +""" +Tests for the Cognition provider identity. + +Cognition serves an OpenAI-compatible /v1/chat/completions surface, but it must resolve to its +own `cognition` provider so OpenAI-specific pricing and provider-level reporting never apply to +its traffic. +""" + +import json +from pathlib import Path + +import pytest + +import litellm + + +class TestCognitionProviderIdentity: + def test_cognition_is_a_registered_provider(self): + from litellm import LlmProviders + + assert LlmProviders.COGNITION.value == "cognition" + assert "cognition" in litellm.provider_list + + def test_cognition_json_config(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + cognition = JSONProviderRegistry.get("cognition") + assert cognition is not None + assert cognition.base_url == "https://api.cognition.ai/v1" + assert cognition.api_key_env == "COGNITION_API_KEY" + assert cognition.api_base_env == "COGNITION_API_BASE" + + def test_cognition_in_openai_compatible_providers(self): + from litellm.constants import openai_compatible_providers + + assert "cognition" in openai_compatible_providers + + def test_prefixed_model_resolves_to_cognition_not_openai(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, _, api_base = get_llm_provider( + model="cognition/swe-1.7", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "swe-1.7" + assert provider == "cognition" + assert api_base == "https://api.cognition.ai/v1" + + def test_explicit_api_base_and_key_win(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + _, provider, api_key, api_base = get_llm_provider( + model="cognition/swe-1.7", + custom_llm_provider=None, + api_base="https://cognition.internal.example/v1", + api_key="sk-test", + ) + + assert provider == "cognition" + assert api_base == "https://cognition.internal.example/v1" + assert api_key == "sk-test" + + def test_api_base_autodetects_cognition(self, monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("COGNITION_API_KEY", "sk-cognition-env") + + _, provider, api_key, api_base = get_llm_provider( + model="swe-1.7", + custom_llm_provider=None, + api_base="https://api.cognition.ai/v1", + api_key=None, + ) + + assert provider == "cognition" + assert api_base == "https://api.cognition.ai/v1" + assert api_key == "sk-cognition-env" + + def test_env_api_key_is_read_from_cognition_variable(self, monkeypatch: pytest.MonkeyPatch): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("COGNITION_API_KEY", "sk-cognition-env") + + provider = JSONProviderRegistry.get("cognition") + assert provider is not None + + api_base, api_key = create_config_class(provider)()._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.cognition.ai/v1" + assert api_key == "sk-cognition-env" + + +class TestCognitionCostTracking: + @pytest.mark.parametrize( + "model, input_cost, output_cost", + [ + ("cognition/swe-1.5", 3e-07, 1.5e-06), + ("cognition/swe-1.6", 3e-07, 1.5e-06), + ("cognition/swe-1.7", 2.5e-06, 1.25e-05), + ], + ) + def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float): + info = litellm.get_model_info(model=model) + + assert info["litellm_provider"] == "cognition" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] == input_cost + assert info["output_cost_per_token"] == output_cost + + def test_cost_differs_from_openai_pricing(self): + """A cognition-prefixed model must never be priced off an OpenAI cost entry.""" + from litellm.cost_calculator import cost_per_token + + prompt_cost, completion_cost = cost_per_token( + model="cognition/swe-1.7", + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + custom_llm_provider="cognition", + ) + + assert prompt_cost == pytest.approx(2.5) + assert completion_cost == pytest.approx(12.5) + + def test_supported_endpoints_matrix(self): + matrix = json.loads((Path(litellm.__file__).parent / "provider_endpoints_support_backup.json").read_text()) + + endpoints = matrix["providers"]["cognition"]["endpoints"] + assert endpoints["chat_completions"] is True + assert endpoints["embeddings"] is False + + +class TestCognitionRouting: + @pytest.mark.asyncio + async def test_router_spend_is_attributed_to_cognition_pricing(self): + """Routed traffic is costed off the cognition entry, not an OpenAI one.""" + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "swe", + "litellm_params": {"model": "cognition/swe-1.7", "api_key": "sk-test"}, + } + ] + ) + + response = await router.acompletion( + model="swe", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello from swe", + ) + + usage = response.usage + expected = usage.prompt_tokens * 2.5e-06 + usage.completion_tokens * 1.25e-05 + assert response._hidden_params["response_cost"] == pytest.approx(expected) diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index fa6b3c79230..375b0025dca 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -85,6 +85,7 @@ export enum Providers { CLARIFAI = "Clarifai", CLOUDFLARE = "Cloudflare", CODESTRAL = "Codestral", + Cognition = "Cognition", Cohere = "Cohere", COHERE_CHAT = "Cohere Chat", COMETAPI = "Cometapi", @@ -193,6 +194,7 @@ export const provider_map: Record = { CLARIFAI: "clarifai", CLOUDFLARE: "cloudflare", CODESTRAL: "codestral", + Cognition: "cognition", Cohere: "cohere", COHERE_CHAT: "cohere_chat", COMETAPI: "cometapi", @@ -439,6 +441,8 @@ export const getPlaceholder = (selectedProvider: string): string => { return "cursor/claude-4-sonnet"; } else if (selectedProvider === Providers.ZAI) { return "zai/glm-4.5"; + } else if (selectedProvider === Providers.Cognition) { + return "cognition/swe-1.7"; } else { return "gpt-3.5-turbo"; }