Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions docs/en/configuration/config-files.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand All @@ -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`
Expand Down
3 changes: 3 additions & 0 deletions docs/zh/configuration/config-files.md
Original file line number Diff line number Diff line change
Expand Up @@ -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` |

示例:

Expand All @@ -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`
Expand Down
3 changes: 3 additions & 0 deletions src/kimi_cli/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)."""

Expand Down
2 changes: 1 addition & 1 deletion src/kimi_cli/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
22 changes: 22 additions & 0 deletions tests/core/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down
122 changes: 122 additions & 0 deletions tests/core/test_create_llm.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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(
Expand Down
Loading