scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-02 02:11:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix: address cometapi review comments

Browse source
This commit is contained in:
TensorNull 2026-06-03 21:03:39 +08:00
parent bc229725f9
commit bdecd5497c
3 changed files with 13 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/images/main.py
View file

@ -442,8 +442,8 @@ def image_generation( # noqa: PLR0915
if custom_llm_provider == litellm.LlmProviders.COMETAPI:
_api_base = api_base
api_key = api_key or dynamic_api_key or litellm.cometapi_key
litellm_params_dict.pop("api_key", None)
litellm_params_dict["api_base"] = _api_base
litellm_params_dict["api_key"] = api_key
return llm_http_handler.image_generation_handler(
api_key=api_key,

3
litellm/llms/cometapi/common_utils.py
View file

@ -25,6 +25,9 @@ def get_cometapi_api_base(api_base: Optional[str] = None) -> str:
def get_cometapi_complete_url(api_base: Optional[str], endpoint: str) -> str:
base_url = get_cometapi_api_base(api_base).rstrip("/")
normalized_endpoint = endpoint.strip("/")
parsed_endpoint = urlsplit(normalized_endpoint)
if not normalized_endpoint or parsed_endpoint.query or parsed_endpoint.fragment:
raise ValueError("CometAPI endpoint must be a non-empty path")
parsed_base_url = urlsplit(base_url)
path_segments = [segment for segment in parsed_base_url.path.split("/") if segment]
endpoint_segments = normalized_endpoint.split("/")

11
tests/test_litellm/llms/cometapi/test_cometapi_endpoints.py
View file

@ -121,6 +121,14 @@ def test_cometapi_complete_url_rejects_non_v1_paths(api_base, monkeypatch):
get_cometapi_complete_url(api_base, "embeddings")
@pytest.mark.parametrize("endpoint", ["", "/", "embeddings?foo=bar", "embeddings#frag"])
def test_cometapi_complete_url_rejects_invalid_endpoints(endpoint, monkeypatch):
_clear_cometapi_env(monkeypatch)
with pytest.raises(ValueError, match="CometAPI endpoint must be a non-empty path"):
get_cometapi_complete_url("https://api.cometapi.com/v1", endpoint)
def test_cometapi_embedding_uses_api_key_alias(monkeypatch):
_clear_cometapi_env(monkeypatch)
monkeypatch.setenv("COMETAPI_API_KEY", "comet-env-key")
@ -253,6 +261,7 @@ def test_cometapi_image_generation_maps_new_openai_params(monkeypatch):
assert call_kwargs["api_key"] == "comet-explicit-key"
assert call_kwargs["custom_llm_provider"] == "cometapi"
assert call_kwargs["litellm_params"]["api_base"] is None
assert "api_key" not in call_kwargs["litellm_params"]
assert (
call_kwargs["image_generation_optional_request_params"]["output_compression"]
== 70
@ -510,8 +519,6 @@ def test_provider_endpoint_matrix_only_updates_cometapi():
"moderations",
):
assert providers["cometapi"]["endpoints"][endpoint] is True
assert providers["a2a"]["endpoints"][endpoint] is False
assert providers["bedrock"]["endpoints"][endpoint] is False
class _ModerationResponse: