feat: add CometAPI provider support with chat completions and streaming (#13458)

* feat: add CometAPI support with config, error handling and tests * fix: specify type for extra_body in CometAPIConfig --------- Signed-off-by: NULL <129579691+TensorNull@users.noreply.github.com>
2026-08-02 06:22:48 +00:00 · 2025-08-11 18:06:37 -07:00
parent 7484a19edf
commit f3dcae2528
9 changed files with 585 additions and 1 deletions
@@ -233,6 +233,7 @@ novita_api_key: Optional[str] = None
 snowflake_key: Optional[str] = None
 gradient_ai_api_key: Optional[str] = None
 nebius_key: Optional[str] = None
+cometapi_key: Optional[str] = None
 common_cloud_provider_auth_params: dict = {
    "params": ["project", "region_name", "token"],
    "providers": ["vertex_ai", "bedrock", "watsonx", "azure", "vertex_ai_beta"],
@@ -535,6 +536,7 @@ morph_models: List = []
 lambda_ai_models: List = []
 hyperbolic_models: List = []
 recraft_models: List = []
+cometapi_models: List = []
 oci_models: List = []


@@ -727,6 +729,8 @@ def add_known_models():
            hyperbolic_models.append(key)
        elif value.get("litellm_provider") == "recraft":
            recraft_models.append(key)
+        elif value.get("litellm_provider") == "cometapi":
+            cometapi_models.append(key)
        elif value.get("litellm_provider") == "oci":
            oci_models.append(key)

@@ -818,6 +822,7 @@ model_list = (
    + morph_models
    + lambda_ai_models
    + recraft_models
+    + cometapi_models
    + oci_models
 )

@@ -893,6 +898,7 @@ models_by_provider: dict = {
    "lambda_ai": lambda_ai_models,
    "hyperbolic": hyperbolic_models,
    "recraft": recraft_models,
+    "cometapi": cometapi_models,
    "oci": oci_models,
 }

@@ -1198,7 +1204,7 @@ from .llms.azure.azure import (
    AzureOpenAIError,
    AzureOpenAIAssistantsAPIConfig,
 )
-
+from .llms.cometapi.chat.transformation import CometAPIConfig
 from .llms.azure.chat.gpt_transformation import AzureOpenAIConfig
 from .llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config
 from .llms.azure.completion.transformation import AzureOpenAITextConfig
@@ -224,6 +224,7 @@ LITELLM_CHAT_PROVIDERS = [
    "together_ai",
    "datarobot",
    "openrouter",
+    "cometapi",
    "vertex_ai",
    "vertex_ai_beta",
    "gemini",
@@ -878,6 +879,7 @@ SENTRY_DENYLIST = [
    "CLOUDFLARE_API_KEY",
    "BASETEN_KEY",
    "OPENROUTER_KEY",
+    "COMETAPI_KEY",
    "DATAROBOT_API_TOKEN",
    "FIREWORKS_API_KEY",
    "FIREWORKS_AI_API_KEY",
@@ -358,6 +358,9 @@ def get_llm_provider(  # noqa: PLR0915
        # bytez models
        elif model.startswith("bytez/"):
            custom_llm_provider = "bytez"
+        # cometapi models
+        elif model.startswith("cometapi/"):
+            custom_llm_provider = "cometapi"
        elif model.startswith("oci/"):
            custom_llm_provider = "oci"
        if not custom_llm_provider:
@@ -0,0 +1,207 @@
+"""
+Support for CometAPI's `/v1/chat/completions` endpoint.
+
+Based on OpenAI-compatible API interface implementation
+Documentation: [CometAPI Documentation Link]
+"""
+
+from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union
+
+import httpx
+
+from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
+from litellm.types.utils import ModelResponse, ModelResponseStream
+
+from ...openai.chat.gpt_transformation import OpenAIGPTConfig
+from ..common_utils import CometAPIException
+
+
+class CometAPIConfig(OpenAIGPTConfig):
+    """
+    CometAPI configuration class, inherits from OpenAIGPTConfig
+    
+    Since CometAPI is OpenAI-compatible API, we inherit from OpenAIGPTConfig
+    and only need to override necessary methods to handle CometAPI-specific features
+    """
+    
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        """
+        Map OpenAI format parameters to CometAPI format
+        """
+        mapped_openai_params = super().map_openai_params(
+            non_default_params, optional_params, model, drop_params
+        )
+
+        # CometAPI-specific parameters (if any)
+        extra_body: dict[str, Any] = {}
+        # TODO: Add CometAPI-specific parameter handling here
+        # Example:
+        # custom_param = non_default_params.pop("custom_param", None)
+        # if custom_param is not None:
+        #     extra_body["custom_param"] = custom_param
+        
+        if extra_body:
+            mapped_openai_params["extra_body"] = extra_body
+        
+        return mapped_openai_params
+
+    def remove_cache_control_flag_from_messages_and_tools(
+        self,
+        model: str,
+        messages: List[AllMessageValues],
+        tools: Optional[List["ChatCompletionToolParam"]] = None,
+    ) -> Tuple[List[AllMessageValues], Optional[List["ChatCompletionToolParam"]]]:
+        """
+        Remove cache control flags from messages and tools if not supported
+        """
+        # For CometAPI, use default behavior (remove cache control)
+        return super().remove_cache_control_flag_from_messages_and_tools(
+            model, messages, tools
+        )
+
+    def transform_request(
+        self,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        headers: dict,
+    ) -> dict:
+        """
+        Transform the overall request to be sent to the API.
+
+        Returns:
+            dict: The transformed request. Sent as the body of the API call.
+        """
+        extra_body = optional_params.pop("extra_body", {})
+        response = super().transform_request(
+            model, messages, optional_params, litellm_params, headers
+        )
+        response.update(extra_body)
+        return response
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        """
+        Get the complete URL for the CometAPI call.
+
+        Returns:
+            str: The complete URL for the API call.
+        """
+        # Default base
+        if api_base is None:
+            api_base = "https://api.cometapi.com/v1"
+        endpoint = "chat/completions"
+
+        # Normalize
+        api_base = api_base.rstrip("/")
+
+        # If endpoint already present, return as-is
+        if endpoint in api_base:
+            return api_base
+
+        # Ensure we include /v1 prefix when missing
+        if api_base.endswith("/v1"):
+            return f"{api_base}/{endpoint}"
+        if api_base.endswith("/v1/"):
+            return f"{api_base}{endpoint}"
+        # If user provided https://api.cometapi.com, add /v1
+        if api_base == "https://api.cometapi.com":
+            return f"{api_base}/v1/{endpoint}"
+        # Generic fallback: if '/v1' not in path, add it
+        if "/v1" not in api_base.split("//", 1)[-1]:
+            return f"{api_base}/v1/{endpoint}"
+        return f"{api_base}/{endpoint}"
+
+    def get_error_class(
+        self, 
+        error_message: str, 
+        status_code: int, 
+        headers: Union[dict, httpx.Headers]
+    ) -> BaseLLMException:
+        """
+        Return CometAPI-specific error class
+        """
+        return CometAPIException(
+            message=error_message,
+            status_code=status_code,
+            headers=headers,
+        )
+
+    def get_model_response_iterator(
+        self,
+        streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse],
+        sync_stream: bool,
+        json_mode: Optional[bool] = False,
+    ) -> Any:
+        """
+        Get model response iterator for streaming responses
+        """
+        return CometAPIChatCompletionStreamingHandler(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
+
+
+class CometAPIChatCompletionStreamingHandler(BaseModelResponseIterator):
+    """
+    Handler for CometAPI streaming chat completion responses
+    """
+    
+    def chunk_parser(self, chunk: dict) -> ModelResponseStream:
+        """
+        Parse individual chunks from streaming response
+        """
+        try:
+            # Handle error in chunk
+            if "error" in chunk:
+                error_chunk = chunk["error"]
+                error_message = "CometAPI Error: {}".format(
+                    error_chunk.get("message", "Unknown error")
+                )
+                raise CometAPIException(
+                    message=error_message,
+                    status_code=error_chunk.get("code", 400),
+                    headers={"Content-Type": "application/json"},
+                )
+
+            # Process choices
+            new_choices = []
+            for choice in chunk["choices"]:
+                # Handle reasoning content if present
+                if "delta" in choice and "reasoning" in choice["delta"]:
+                    choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning")
+                new_choices.append(choice)
+            
+            return ModelResponseStream(
+                id=chunk["id"],
+                object="chat.completion.chunk",
+                created=chunk["created"],
+                usage=chunk.get("usage"),
+                model=chunk["model"],
+                choices=new_choices,
+            )
+        except KeyError as e:
+            raise CometAPIException(
+                message=f"KeyError: {e}, Got unexpected response from CometAPI: {chunk}",
+                status_code=400,
+                headers={"Content-Type": "application/json"},
+            )
+        except Exception as e:
+            raise e
@@ -0,0 +1,6 @@
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+
+class CometAPIException(BaseLLMException):
+    """CometAPI exception handling class"""
+    pass
@@ -1883,6 +1883,45 @@ def completion(  # type: ignore # noqa: PLR0915
                encoding=encoding,
                stream=stream,
            )
+        elif custom_llm_provider == "cometapi":
+            api_key = (
+                api_key
+                or litellm.cometapi_key
+                or get_secret_str("COMETAPI_KEY")
+                or litellm.api_key
+            )
+            
+            api_base = (
+                api_base
+                or litellm.api_base
+                or get_secret_str("COMETAPI_API_BASE")
+                or "https://api.cometapi.com/v1"
+            )
+
+            ## COMPLETION CALL
+            response = base_llm_http_handler.completion(
+                model=model,
+                messages=messages,
+                headers=headers,
+                model_response=model_response,
+                api_key=api_key,
+                api_base=api_base,
+                acompletion=acompletion,
+                logging_obj=logging,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                timeout=timeout,
+                client=client,
+                custom_llm_provider=custom_llm_provider,
+                encoding=encoding,
+                stream=stream,
+                provider_config=provider_config,
+            )
+            
+            ## LOGGING
+            logging.post_call(
+                input=messages, api_key=api_key, original_response=response
+            )
        elif (
            model in litellm.open_ai_chat_completion_models
            or custom_llm_provider == "custom_openai"
@@ -2330,6 +2330,7 @@ class LlmProviders(str, Enum):
    PG_VECTOR = "pg_vector"
    HYPERBOLIC = "hyperbolic"
    RECRAFT = "recraft"
+    COMETAPI = "cometapi"
    OCI = "oci"
    AUTO_ROUTER = "auto_router"
    DOTPROMPT = "dotprompt"
@@ -6859,6 +6859,8 @@ class ProviderConfigManager:
            return litellm.TogetherAIConfig()
        elif litellm.LlmProviders.OPENROUTER == provider:
            return litellm.OpenrouterConfig()
+        elif litellm.LlmProviders.COMETAPI == provider:
+            return litellm.CometAPIConfig()
        elif litellm.LlmProviders.DATAROBOT == provider:
            return litellm.DataRobotConfig()
        elif litellm.LlmProviders.GEMINI == provider:
@@ -0,0 +1,318 @@
+"""
+Unit tests for CometAPI Chat Configuration
+
+Tests the CometAPIChatConfig class methods using mocks
+"""
+
+import os
+import sys
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.llms.cometapi.chat.transformation import (
+    CometAPIChatCompletionStreamingHandler,
+    CometAPIConfig,
+)
+from litellm.llms.cometapi.common_utils import CometAPIException
+
+
+class TestCometAPIChatCompletionStreamingHandler:
+    def test_chunk_parser_successful(self):
+        handler = CometAPIChatCompletionStreamingHandler(
+            streaming_response=None, sync_stream=True
+        )
+
+        # Test input chunk
+        chunk = {
+            "id": "test_id",
+            "created": 1234567890,
+            "model": "gpt-3.5-turbo",
+            "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
+            "choices": [
+                {"delta": {"content": "test content", "reasoning": "test reasoning"}}
+            ],
+        }
+
+        # Parse chunk
+        result = handler.chunk_parser(chunk)
+
+        # Verify response
+        assert result.id == "test_id"
+        assert result.object == "chat.completion.chunk"
+        assert result.created == 1234567890
+        assert result.model == "gpt-3.5-turbo"
+        assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"]
+        assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"]
+        assert result.usage.total_tokens == chunk["usage"]["total_tokens"]
+        assert len(result.choices) == 1
+        assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning"
+
+    def test_chunk_parser_error_response(self):
+        handler = CometAPIChatCompletionStreamingHandler(
+            streaming_response=None, sync_stream=True
+        )
+
+        # Test error chunk
+        error_chunk = {
+            "error": {
+                "message": "test error",
+                "code": 400,
+            }
+        }
+
+        # Verify error handling
+        with pytest.raises(CometAPIException) as exc_info:
+            handler.chunk_parser(error_chunk)
+
+        assert "CometAPI Error: test error" in str(exc_info.value)
+        assert exc_info.value.status_code == 400
+
+    def test_chunk_parser_key_error(self):
+        handler = CometAPIChatCompletionStreamingHandler(
+            streaming_response=None, sync_stream=True
+        )
+
+        # Test invalid chunk missing required fields
+        invalid_chunk = {"incomplete": "data"}
+
+        # Verify KeyError handling
+        with pytest.raises(CometAPIException) as exc_info:
+            handler.chunk_parser(invalid_chunk)
+
+        assert "KeyError" in str(exc_info.value)
+        assert exc_info.value.status_code == 400
+
+
+class TestCometAPIConfig:
+    def test_transform_request_basic(self):
+        """Test basic request transformation"""
+        config = CometAPIConfig()
+        
+        transformed_request = config.transform_request(
+            model="cometapi/gpt-3.5-turbo",
+            messages=[
+                {"role": "user", "content": "Hello, world!"}
+            ],
+            optional_params={},
+            litellm_params={},
+            headers={},
+        )
+
+        assert transformed_request["model"] == "cometapi/gpt-3.5-turbo"
+        assert transformed_request["messages"] == [
+            {"role": "user", "content": "Hello, world!"}
+        ]
+
+    def test_transform_request_with_extra_body(self):
+        """Test request transformation with extra_body parameters"""
+        config = CometAPIConfig()
+        
+        transformed_request = config.transform_request(
+            model="cometapi/gpt-4",
+            messages=[{"role": "user", "content": "Hello, world!"}],
+            optional_params={"extra_body": {"custom_param": "custom_value"}},
+            litellm_params={},
+            headers={},
+        )
+
+        # Validate that extra_body parameters are merged into the request
+        assert transformed_request["custom_param"] == "custom_value"
+        assert transformed_request["messages"] == [
+            {"role": "user", "content": "Hello, world!"}
+        ]
+
+    def test_cache_control_flag_removal(self):
+        """Test cache control flag removal from messages"""
+        config = CometAPIConfig()
+        
+        transformed_request = config.transform_request(
+            model="cometapi/gpt-3.5-turbo",
+            messages=[
+                {
+                    "role": "user",
+                    "content": "Hello, world!",
+                    "cache_control": {"type": "ephemeral"},
+                }
+            ],
+            optional_params={},
+            litellm_params={},
+            headers={},
+        )
+        
+        # CometAPI should remove cache_control flags by default
+        assert transformed_request["messages"][0].get("cache_control") is None
+
+    def test_map_openai_params(self):
+        """Test OpenAI parameter mapping"""
+        config = CometAPIConfig()
+        
+        non_default_params = {
+            "temperature": 0.7,
+            "max_tokens": 100,
+            "top_p": 0.9,
+        }
+        
+        mapped_params = config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params={},
+            model="cometapi/gpt-3.5-turbo",
+            drop_params=False,
+        )
+        
+        assert mapped_params["temperature"] == 0.7
+        assert mapped_params["max_tokens"] == 100
+        assert mapped_params["top_p"] == 0.9
+
+    def test_get_error_class(self):
+        """Test error class creation"""
+        config = CometAPIConfig()
+        
+        error = config.get_error_class(
+            error_message="Test error",
+            status_code=400,
+            headers={"Content-Type": "application/json"}
+        )
+        
+        assert isinstance(error, CometAPIException)
+        assert error.message == "Test error"
+        assert error.status_code == 400
+
+
+# Integration test example (requires real API key)
+@pytest.mark.skip(reason="Skipping integration test")
+def test_cometapi_integration():
+    """
+    Integration test - requires real API key
+    Run with: pytest -k test_cometapi_integration -s
+    """
+    import os
+    from litellm import completion
+    
+    # Try to get API key from multiple environment variables
+    api_key = (
+        os.getenv("COMETAPI_API_KEY") 
+        or os.getenv("COMETAPI_KEY")
+        or os.getenv("COMET_API_KEY")
+    )
+    
+    if not api_key:
+        pytest.skip("COMETAPI_API_KEY not set - skipping integration test")
+    
+    response = completion(
+        model="cometapi/gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "Say hello in one word"}],
+        api_key=api_key,
+        max_tokens=10,
+        temperature=0.7
+    )
+    
+    # Verify response structure
+    assert response.choices[0].message.content
+    assert len(response.choices[0].message.content.strip()) > 0
+    assert response.model
+    assert response.usage
+    assert response.usage.total_tokens > 0
+
+
+def test_cometapi_streaming_integration():
+    """
+    Integration test for streaming - requires real API key
+    Run with: pytest -k test_cometapi_streaming_integration -s
+    """
+    import os
+    from litellm import completion
+    
+    # Try to get API key from multiple environment variables
+    api_key = (
+        os.getenv("COMETAPI_API_KEY") 
+        or os.getenv("COMETAPI_KEY")
+        or os.getenv("COMET_API_KEY")
+    )
+    
+    if not api_key:
+        pytest.skip("COMETAPI_API_KEY not set - skipping streaming integration test")
+    
+    try:
+        print(f"🔍 Testing streaming with API key: {api_key[:6]}...{api_key[-4:]} (length: {len(api_key)})")
+        print(f"🔍 API base URL: {os.getenv('COMETAPI_API_BASE', 'default')}")
+        
+        # test streaming API call
+        response = completion(
+            model="cometapi/gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "Count from 1 to 5"}],
+            api_key=api_key,
+            max_tokens=50,
+            stream=True
+        )
+
+        # collect streaming response
+        chunks = []
+        content_parts = []
+
+        for chunk in response:
+            chunks.append(chunk)
+            if chunk.choices[0].delta.content:
+                content_parts.append(chunk.choices[0].delta.content)
+
+        # Verify we received at least one chunk and content
+        assert len(chunks) > 0, "Should receive at least one chunk"
+        assert len(content_parts) > 0, "Should receive content in chunks"
+
+        full_content = "".join(content_parts)
+        assert len(full_content.strip()) > 0, "Should have non-empty content"
+
+        print(f"✅ Received {len(chunks)} chunks")
+        print(f"✅ Full content: {full_content}")
+
+    except Exception as e:
+        print(f"❌ Streaming integration test error details:")
+        print(f"   Error type: {type(e).__name__}")
+        print(f"   Error message: {str(e)}")
+        if hasattr(e, 'status_code'):
+            print(f"   Status code: {e.status_code}")
+        if hasattr(e, 'response'):
+            print(f"   Response: {e.response}")
+            
+        # Re-raise with more context for pytest
+        pytest.fail(f"Streaming integration test failed: {type(e).__name__}: {str(e)}")
+def test_cometapi_with_custom_base_url():
+    """
+    Test CometAPI with custom base URL
+    """
+    import os
+    from litellm import completion
+    
+    api_key = (
+        os.getenv("COMETAPI_API_KEY") 
+        or os.getenv("COMETAPI_KEY")
+        or os.getenv("COMET_API_KEY")
+    )
+    
+    custom_base_url = os.getenv("COMETAPI_API_BASE", "https://api.cometapi.com/v1")
+    
+    if not api_key:
+        pytest.skip("COMETAPI_API_KEY not set - skipping custom base URL test")
+    
+    try:
+        response = completion(
+            model="cometapi/gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "Hello"}],
+            api_key=api_key,
+            api_base=custom_base_url,
+            max_tokens=5
+        )
+        
+        assert response.choices[0].message.content
+        print(f"✅ Custom base URL test passed: {response.choices[0].message.content}")
+        
+    except Exception as e:
+        pytest.fail(f"Custom base URL test failed: {str(e)}")
+
+
+if __name__ == "__main__":
+    # Quick test runner
+    pytest.main([__file__, "-v"])