diff --git a/docs/en/configuration/config-files.md b/docs/en/configuration/config-files.md index c504bdc3b4..5046cfdea8 100644 --- a/docs/en/configuration/config-files.md +++ b/docs/en/configuration/config-files.md @@ -100,6 +100,7 @@ tool_call_timeout_ms = 60000 | `api_key` | `string` | Yes | API key | | `env` | `table` | No | Environment variables to set before creating provider instance | | `custom_headers` | `table` | No | Custom HTTP headers to attach to requests | +| `prompt_cache_key` | `boolean` | No | Send the session ID as `prompt_cache_key` for `kimi` providers (default: `true`). Set to `false` for compatible APIs that reject this parameter | Example: @@ -109,6 +110,8 @@ type = "kimi" base_url = "https://api.moonshot.cn/v1" api_key = "sk-xxx" custom_headers = { "X-Custom-Header" = "value" } +# For third-party Kimi-compatible APIs that reject prompt_cache_key: +# prompt_cache_key = false ``` ### `models` diff --git a/docs/zh/configuration/config-files.md b/docs/zh/configuration/config-files.md index c5464df402..24e0a998ee 100644 --- a/docs/zh/configuration/config-files.md +++ b/docs/zh/configuration/config-files.md @@ -100,6 +100,7 @@ tool_call_timeout_ms = 60000 | `api_key` | `string` | 是 | API 密钥 | | `env` | `table` | 否 | 创建供应商实例前设置的环境变量 | | `custom_headers` | `table` | 否 | 请求时附加的自定义 HTTP 头 | +| `prompt_cache_key` | `boolean` | 否 | `kimi` 供应商是否将会话 ID 作为 `prompt_cache_key` 发送（默认：`true`）。兼容 API 不支持此参数时设为 `false` | 示例： @@ -109,6 +110,8 @@ type = "kimi" base_url = "https://api.moonshot.cn/v1" api_key = "sk-xxx" custom_headers = { "X-Custom-Header" = "value" } +# 对于不支持 prompt_cache_key 的第三方 Kimi 兼容 API： +# prompt_cache_key = false ``` ### `models` diff --git a/src/kimi_cli/config.py b/src/kimi_cli/config.py index 451f468100..44db1b16a2 100644 --- a/src/kimi_cli/config.py +++ b/src/kimi_cli/config.py @@ -49,6 +49,9 @@ class LLMProvider(BaseModel): """Message field name carrying reasoning content for OpenAI-compatible APIs. Applies to provider type ``openai_legacy``. Defaults to ``reasoning_content`` when unset. Use an empty string to disable reasoning round-tripping.""" + prompt_cache_key: bool = True + """Whether ``kimi`` providers send the session ID as ``prompt_cache_key``. + Disable for third-party compatible APIs that reject this Kimi-specific parameter.""" oauth: OAuthRef | None = None """OAuth credential reference (do not store tokens here).""" diff --git a/src/kimi_cli/llm.py b/src/kimi_cli/llm.py index 4b65fa227c..dfcf697d02 100644 --- a/src/kimi_cli/llm.py +++ b/src/kimi_cli/llm.py @@ -358,7 +358,7 @@ def create_llm( ) gen_kwargs: Kimi.GenerationKwargs = {} - if session_id: + if session_id and provider.prompt_cache_key: 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_config.py b/tests/core/test_config.py index ac8737047d..0036c51032 100644 --- a/tests/core/test_config.py +++ b/tests/core/test_config.py @@ -74,6 +74,28 @@ def test_load_config_text_json(): assert config == get_default_config() +@pytest.mark.parametrize( + ("setting", "expected"), + [ + ("", True), + ("prompt_cache_key = true", True), + ("prompt_cache_key = false", False), + ], +) +def test_load_config_kimi_prompt_cache_key(setting: str, expected: bool): + config = load_config_from_string( + f""" +[providers.third-party] +type = "kimi" +base_url = "https://api.example.com/v1" +api_key = "test-key" +{setting} +""" + ) + + assert config.providers["third-party"].prompt_cache_key is expected + + def test_load_config_sets_source_file(tmp_path): config_file = tmp_path / "custom.toml" diff --git a/tests/core/test_create_llm.py b/tests/core/test_create_llm.py index dbbf0d09e0..253260a047 100644 --- a/tests/core/test_create_llm.py +++ b/tests/core/test_create_llm.py @@ -1,9 +1,17 @@ from __future__ import annotations +import json +import threading +from collections.abc import Iterator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest from inline_snapshot import snapshot from kosong.chat_provider.echo import EchoChatProvider from kosong.chat_provider.kimi import Kimi from kosong.contrib.chat_provider.openai_responses import OpenAIResponses +from kosong.message import Message from pydantic import SecretStr from kimi_cli.config import Config, LLMModel, LLMProvider @@ -107,6 +115,120 @@ def test_create_llm_kimi_prefers_max_completion_tokens_env(monkeypatch): assert llm.chat_provider.model_parameters["max_completion_tokens"] == 5678 +def _chat_completion_response() -> dict[str, object]: + return { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "test-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + +@contextmanager +def _fake_chat_endpoint( + *, accept_prompt_cache_key: bool +) -> Iterator[tuple[str, list[dict[str, object]]]]: + requests: list[dict[str, object]] = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + length = int(self.headers["Content-Length"]) + body = json.loads(self.rfile.read(length)) + requests.append(body) + if not accept_prompt_cache_key and "prompt_cache_key" in body: + status = 400 + response = { + "error": {"message": "Validation: Unsupported parameter(s): `prompt_cache_key`"} + } + else: + status = 200 + response = _chat_completion_response() + encoded = json.dumps(response).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, format: str, *args: object) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + host = server.server_address[0] + port = server.server_address[1] + yield f"http://{host}:{port}/v1", requests + finally: + server.shutdown() + server.server_close() + thread.join() + + +@pytest.mark.asyncio +async def test_create_llm_kimi_can_disable_prompt_cache_key_for_third_party( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost") + with _fake_chat_endpoint(accept_prompt_cache_key=False) as (base_url, requests): + provider = LLMProvider( + type="kimi", + base_url=base_url, + api_key=SecretStr("test-key"), + prompt_cache_key=False, + ) + model = LLMModel( + provider="nvidia", + model="nvidia-model", + max_context_size=4096, + ) + llm = create_llm(provider, model, session_id="session-123") + assert llm is not None + assert isinstance(llm.chat_provider, Kimi) + llm.chat_provider.stream = False + + await llm.chat_provider.generate("", [], [Message(role="user", content="hello")]) + + assert len(requests) == 1 + assert "prompt_cache_key" not in requests[0] + + +@pytest.mark.asyncio +async def test_create_llm_kimi_sends_prompt_cache_key_by_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost") + with _fake_chat_endpoint(accept_prompt_cache_key=True) as (base_url, requests): + provider = LLMProvider( + type="kimi", + base_url=base_url, + api_key=SecretStr("test-key"), + ) + model = LLMModel( + provider="managed:kimi-code", + model="kimi-for-coding", + max_context_size=4096, + ) + llm = create_llm(provider, model, session_id="session-123") + assert llm is not None + assert isinstance(llm.chat_provider, Kimi) + llm.chat_provider.stream = False + + await llm.chat_provider.generate("", [], [Message(role="user", content="hello")]) + + assert len(requests) == 1 + assert requests[0]["prompt_cache_key"] == "session-123" + + def test_compute_max_completion_tokens_uses_response_budget_when_it_fits(): assert ( compute_max_completion_tokens(