diff --git a/backend/packages/harness/deerflow/models/patched_deepseek.py b/backend/packages/harness/deerflow/models/patched_deepseek.py index 177d0ab98e..a47c355962 100644 --- a/backend/packages/harness/deerflow/models/patched_deepseek.py +++ b/backend/packages/harness/deerflow/models/patched_deepseek.py @@ -9,8 +9,10 @@ on all assistant messages when thinking mode is enabled. from typing import Any +import openai from langchain_core.language_models import LanguageModelInput from langchain_core.messages import AIMessage +from langchain_core.outputs import ChatGenerationChunk, ChatResult from langchain_deepseek import ChatDeepSeek from deerflow.models.assistant_payload_replay import restore_assistant_payloads, restore_reasoning_content @@ -100,3 +102,64 @@ class PatchedChatDeepSeek(ChatDeepSeek): ) return payload + + def _create_chat_result( + self, + response: dict | openai.BaseModel, + generation_info: dict | None = None, + ) -> ChatResult: + """Preserve Moonshot's ``cache_write_tokens`` in the usage metadata. + + The locked ``langchain-openai==1.2.1`` conversion reads + ``prompt_tokens_details.cached_tokens`` (mapped to ``cache_read``) but + not ``cache_write_tokens``, so providers such as Kimi K3.1 lose their + cache-creation count before it reaches the usage journal (#7426). + Restore it as ``input_token_details["cache_creation"]`` from the raw + SDK response. + """ + result = super()._create_chat_result(response, generation_info) + + if not isinstance(response, openai.BaseModel): + return result + prompt_details = getattr(getattr(response, "usage", None), "prompt_tokens_details", None) + cache_write_tokens = getattr(prompt_details, "cache_write_tokens", None) + if cache_write_tokens is None: + return result + for generation in result.generations: + usage = getattr(generation.message, "usage_metadata", None) + if not usage: + continue + input_details = usage.setdefault("input_token_details", {}) + input_details["cache_creation"] = cache_write_tokens + return result + + def _convert_chunk_to_generation_chunk( + self, + chunk: dict, + default_chunk_class: type, + base_generation_info: dict | None, + ) -> ChatGenerationChunk | None: + """Streaming mirror of ``_create_chat_result``'s cache_write restore. + + Streamed runs (the web UI's ``messages-tuple`` mode) build usage + metadata in ``_convert_chunk_to_generation_chunk`` -> + ``_create_usage_metadata``, which ignores ``cache_write_tokens`` + (#7426). Restore it the same way for the final usage chunk. + """ + generation_chunk = super()._convert_chunk_to_generation_chunk( + chunk, + default_chunk_class, + base_generation_info, + ) + if generation_chunk is None: + return None + prompt_details = ((chunk or {}).get("usage") or {}).get("prompt_tokens_details") or {} + cache_write_tokens = prompt_details.get("cache_write_tokens") + if cache_write_tokens is None: + return generation_chunk + usage = getattr(generation_chunk.message, "usage_metadata", None) + if not usage: + return generation_chunk + input_details = usage.setdefault("input_token_details", {}) + input_details["cache_creation"] = cache_write_tokens + return generation_chunk diff --git a/backend/tests/test_patched_deepseek.py b/backend/tests/test_patched_deepseek.py index 532e16c938..c686d99ee8 100644 --- a/backend/tests/test_patched_deepseek.py +++ b/backend/tests/test_patched_deepseek.py @@ -339,3 +339,79 @@ def test_legacy_fallback_assistant_is_removed_before_deepseek_replay(): assert payload["messages"] == base_payload["messages"] assert not any(message.get("role") == "assistant" for message in payload["messages"]) + + +# --------------------------------------------------------------------------- +# cache_write_tokens preservation (#7426) +# --------------------------------------------------------------------------- + + +def _chat_completion_response(): + from openai.types import CompletionUsage + from openai.types.chat import ChatCompletion, ChatCompletionMessage + from openai.types.chat.chat_completion import Choice + + return ChatCompletion( + id="chatcmpl-cache-write", + choices=[ + Choice( + finish_reason="stop", + index=0, + message=ChatCompletionMessage(role="assistant", content="ok"), + ) + ], + created=1761200000, + model="kimi-k3.1-pro", + object="chat.completion", + usage=CompletionUsage( + prompt_tokens=373, + completion_tokens=112, + total_tokens=485, + prompt_tokens_details={"cached_tokens": 100, "cache_write_tokens": 256}, + completion_tokens_details={"reasoning_tokens": 90}, + ), + ) + + +def test_cache_write_tokens_survive_the_conversion(): + """Moonshot reports ``prompt_tokens_details.cache_write_tokens``; the locked + langchain-openai 1.2.1 conversion drops it, so the patch layer must restore + it as ``input_token_details["cache_creation"]`` (#7426).""" + model = _make_model() + result = model._create_chat_result(_chat_completion_response()) + + usage = result.generations[0].message.usage_metadata + assert usage["input_tokens"] == 373 + assert usage["output_tokens"] == 112 + assert usage["input_token_details"]["cache_read"] == 100 + assert usage["input_token_details"]["cache_creation"] == 256 + assert usage["output_token_details"]["reasoning"] == 90 + + +def test_cache_write_tokens_survive_streaming_conversion(): + """Streamed runs (web UI `messages-tuple`) go through + `_convert_chunk_to_generation_chunk`, which must also restore + `cache_creation` (#7426).""" + from langchain_core.messages import ChatMessageChunk + + model = _make_model() + chunk = { + "id": "chatcmpl-cache-write", + "object": "chat.completion.chunk", + "created": 1761200000, + "model": "kimi-k3.1-pro", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 373, + "completion_tokens": 112, + "total_tokens": 485, + "prompt_tokens_details": {"cached_tokens": 100, "cache_write_tokens": 256}, + "completion_tokens_details": {"reasoning_tokens": 90}, + }, + } + generation = model._convert_chunk_to_generation_chunk(chunk, ChatMessageChunk, None) + + usage = generation.message.usage_metadata + assert usage["input_tokens"] == 373 + assert usage["input_token_details"]["cache_read"] == 100 + assert usage["input_token_details"]["cache_creation"] == 256