diff --git a/src/kimi_cli/llm.py b/src/kimi_cli/llm.py index 4b65fa227c..7c5c2b0fc5 100644 --- a/src/kimi_cli/llm.py +++ b/src/kimi_cli/llm.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, Protocol, Self, cast, get_args +from urllib.parse import urlparse from kosong.chat_provider import ChatProvider, StreamedMessage, ThinkingEffort from kosong.message import ( @@ -47,6 +48,7 @@ DEFAULT_UNKNOWN_CONTEXT_COMPLETION_TOKENS = 32_000 DEFAULT_COMPLETION_TOKEN_SAFETY_MARGIN = 1_024 MEDIA_TOKEN_ESTIMATE = 2_000 +MOONSHOT_API_HOSTS = {"api.kimi.com", "api.moonshot.ai", "api.moonshot.cn"} @dataclass(slots=True) @@ -323,6 +325,16 @@ def _kimi_default_headers(provider: LLMProvider, oauth: OAuthManager | None) -> return headers +def _supports_prompt_cache_key(base_url: str) -> bool: + url = base_url.strip() + if "//" not in url: + # urlparse() only fills in `hostname` when a netloc is present, so a + # scheme-less base_url such as "api.moonshot.ai/v1" parses entirely as + # a path. Give it a netloc before parsing. + url = f"//{url}" + return urlparse(url).hostname in MOONSHOT_API_HOSTS + + def create_llm( provider: LLMProvider, model: LLMModel, @@ -358,7 +370,7 @@ def create_llm( ) gen_kwargs: Kimi.GenerationKwargs = {} - if session_id: + if session_id and _supports_prompt_cache_key(provider.base_url): gen_kwargs["prompt_cache_key"] = session_id if temperature := os.getenv("KIMI_MODEL_TEMPERATURE"): gen_kwargs["temperature"] = float(temperature) diff --git a/tests/core/test_create_llm.py b/tests/core/test_create_llm.py index dbbf0d09e0..ca8958071f 100644 --- a/tests/core/test_create_llm.py +++ b/tests/core/test_create_llm.py @@ -84,6 +84,49 @@ def test_create_llm_kimi_model_parameters(monkeypatch): ) +def test_create_llm_omits_prompt_cache_key_for_third_party_kimi_endpoint(): + provider = LLMProvider( + type="kimi", + base_url="https://integrate.api.nvidia.com/v1", + api_key=SecretStr("test-key"), + ) + model = LLMModel( + provider="nvidia", + model="meta/llama", + max_context_size=4096, + ) + + llm = create_llm(provider, model, session_id="session-123") + + assert llm is not None + assert isinstance(llm.chat_provider, Kimi) + assert "prompt_cache_key" not in llm.chat_provider.model_parameters + + +def test_create_llm_keeps_prompt_cache_key_for_moonshot_endpoints(): + model = LLMModel(provider="kimi", model="kimi-base", max_context_size=4096) + + for base_url in ( + "https://api.kimi.com/coding/v1", + "https://api.moonshot.ai/v1", + "https://api.moonshot.cn/v1", + # scheme-less forms are accepted by the OpenAI-compatible clients too + "api.moonshot.ai/v1", + "api.kimi.com/coding/v1", + ): + provider = LLMProvider( + type="kimi", + base_url=base_url, + api_key=SecretStr("test-key"), + ) + + llm = create_llm(provider, model, session_id="session-123") + + assert llm is not None + assert isinstance(llm.chat_provider, Kimi) + assert llm.chat_provider.model_parameters["prompt_cache_key"] == "session-123" + + def test_create_llm_kimi_prefers_max_completion_tokens_env(monkeypatch): provider = LLMProvider( type="kimi",