From fe21d7704ad7e6256db774807850efa179af5557 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Mon, 2 Nov 2026 01:39:26 +0000 Subject: [PATCH 1/4] fix(model): correct OpenAI-compatible response handling --- README.md | 4 +- .../interfaces/model/openai_compatible/llm.py | 46 ++---------- .../model/openai_compatible/test_http_llm.py | 73 ++++++++++++++----- .../test_model_configuration.py | 16 ++++ 4 files changed, 80 insertions(+), 59 deletions(-) diff --git a/README.md b/README.md index 534cb9d3..81710d59 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,8 @@ In SDK 0.11.0 (unreleased), streaming and non-streaming responses estimate missi Provider counts, including zero, take precedence; a missing or null count is estimated independently. Completion counts the text sent from the first prompt message, while chat counts the complete message history and the tool definitions sent with the request. Output estimates include text and tool calls, including responses with no text content. +Tool parameter estimates include the complete JSON schema, including optional and nested fields. +Tool-only non-streaming replies use an empty string for absent or null text so the daemon can decode the result. These local estimates use the SDK tokenizer and pricing configuration; they do not represent the provider's billing record. The public `get_num_tokens()` budget for completion also counts only the first prompt message and any supplied tool definitions. @@ -97,7 +99,7 @@ The parser handles LF, CRLF, and CR line endings within events, including with c A single leading UTF-8 BOM is ignored for both default and custom delimiters. Raw JSON objects, single-prefix multiline JSON, and final payloads without a blank-line terminator remain supported. Explicit response charsets are respected; responses without a charset default to UTF-8. -Streams must include a non-null `finish_reason` for the selected choice or a `[DONE]` event. +Streams must include a non-empty `finish_reason` for the selected choice or a `[DONE]` event. EOF without either raises an invocation error instead of returning a successful final chunk with partial text or tool calls. ## Version Management diff --git a/src/dify_plugin/interfaces/model/openai_compatible/llm.py b/src/dify_plugin/interfaces/model/openai_compatible/llm.py index 72cf5556..6d130f37 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/llm.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/llm.py @@ -130,7 +130,7 @@ def _iter_stream_json( raise InvokeError(msg) from data_error yield chunk choices = chunk.get("choices") - if choices and choices[0].get("finish_reason") is not None: + if choices and choices[0].get("finish_reason"): has_finish_reason = True if not has_finish_reason: @@ -856,7 +856,7 @@ def _handle_generate_stream_response( continue choice = chunk_json["choices"][0] - if choice.get("finish_reason") is not None: + if choice.get("finish_reason"): finish_reason = choice["finish_reason"] chunk_index += 1 @@ -1008,6 +1008,8 @@ def _handle_generate_response( if completion_type is LLMMode.CHAT: response_message = output.get("message", {}) response_content = response_message.get("content") + if response_content is None: + response_content = "" if function_calling_type == "tool_call": tool_calls = output.get("message", {}).get("tool_calls") elif function_calling_type == "function_call": @@ -1263,43 +1265,9 @@ def _num_tokens_for_tools(self, tools: list[PromptMessageTool]) -> int: return num_tokens def _num_tokens_for_tool_parameters(self, parameters: dict) -> int: - num_tokens = 0 - if "title" in parameters: - num_tokens += self._get_num_tokens_by_gpt2("title") - num_tokens += self._get_num_tokens_by_gpt2(parameters.get("title")) - num_tokens += self._get_num_tokens_by_gpt2("type") - num_tokens += self._get_num_tokens_by_gpt2(parameters.get("type")) - if "properties" in parameters: - num_tokens += self._get_num_tokens_by_gpt2("properties") - for key, value in parameters.get("properties", {}).items(): - num_tokens += self._num_tokens_for_tool_property(key, value) - if "required" in parameters: - num_tokens += self._get_num_tokens_by_gpt2("required") - for required_field in parameters["required"]: - num_tokens += 3 - num_tokens += self._get_num_tokens_by_gpt2(required_field) - return num_tokens - - def _num_tokens_for_tool_property(self, key: str, value: dict) -> int: - num_tokens = self._get_num_tokens_by_gpt2(key) - for field_key, field_value in value.items(): - num_tokens += self._get_num_tokens_by_gpt2(field_key) - if field_key == "enum": - num_tokens += self._num_tokens_for_enum_field(field_value) - else: - num_tokens += self._get_num_tokens_by_gpt2(field_key) - num_tokens += self._get_num_tokens_by_gpt2(str(field_value)) - return num_tokens - - def _num_tokens_for_enum_field(self, field_value: list[Any]) -> int: - num_tokens = 0 - for enum_field in field_value: - num_tokens += 3 - value = ( - enum_field if isinstance(enum_field, str) else json.dumps(enum_field) - ) - num_tokens += self._get_num_tokens_by_gpt2(value) - return num_tokens + return self._get_num_tokens_by_gpt2( + json.dumps(parameters, ensure_ascii=False, allow_nan=False) + ) def _extract_response_tool_calls( self, diff --git a/tests/interfaces/model/openai_compatible/test_http_llm.py b/tests/interfaces/model/openai_compatible/test_http_llm.py index 7eef2786..d266f8dc 100644 --- a/tests/interfaces/model/openai_compatible/test_http_llm.py +++ b/tests/interfaces/model/openai_compatible/test_http_llm.py @@ -29,6 +29,7 @@ OAICompatLargeLanguageModel, ) from tests.http_server import HTTPServerFactory, Response +from tests.servers.helpers import RecordingResponseWriter if TYPE_CHECKING: from requests import Response as HTTPResponse @@ -469,7 +470,12 @@ def test_non_stream_tool_calls_preserve_arguments_and_request_schema( ) chunks = list(invocation) call = chunks[-1].delta.message.tool_calls[0] - assert chunks[-1].delta.message.content is None + writer = RecordingResponseWriter() + writer.session_message("session-1", writer.stream_object(chunks[-1])) + wire_message = writer.messages[0]["data"]["data"]["delta"]["message"] + assert isinstance(wire_message["content"], str) + assert not wire_message["content"] + assert wire_message["tool_calls"][0]["function"] == function assert call.function.name == "weather" assert json.loads(call.function.arguments) == {"city": "εŒ—δΊ¬"} usage = chunks[-1].delta.usage @@ -694,6 +700,7 @@ def test_stream_errors_after_content_do_not_report_success( ("stop", ""), ("length", "\n\n"), (None, "\n\ndata: [DONE]\n\n"), + ("", "\n\ndata: [DONE]\n\n"), (None, "\n\nevent: message\ndata: [DONE]\n\n"), ("stop", "\n\ndata: [DONE]\n\n"), ], @@ -719,7 +726,7 @@ def test_stream_accepts_finish_reason_or_done( assert next(stream).delta.message.content == "Answer" terminal = next(stream) assert terminal.delta.message == AssistantPromptMessage(content="") - assert terminal.delta.finish_reason == finish_reason + assert terminal.delta.finish_reason == (finish_reason or None) assert terminal.delta.usage is not None assert list(stream) == [] assert model.started_at == 0 @@ -731,6 +738,7 @@ def test_stream_accepts_finish_reason_or_done( (None, ""), ({"delta": {"role": "assistant"}}, ""), ({"delta": {"content": "partial"}, "finish_reason": None}, "partial"), + ({"delta": {"content": "partial"}, "finish_reason": ""}, "partial"), ({"delta": {"reasoning_content": "partial"}}, "\npartial"), ( { @@ -748,8 +756,18 @@ def test_stream_accepts_finish_reason_or_done( "", ), ({"text": "partial"}, "partial"), + ({"text": "partial", "finish_reason": ""}, "partial"), + ], + ids=[ + "empty", + "role", + "content", + "empty-chat-finish", + "reasoning", + "tool", + "completion", + "empty-completion-finish", ], - ids=["empty", "role", "content", "reasoning", "tool", "completion"], ) @pytest.mark.parametrize( "ending", @@ -842,7 +860,7 @@ def test_missing_usage_counts_sent_tools_through_provider_hooks( "input_price": "1", "unit": "0.001", } - expected = 18 if calling_type == "no_call" else 79 + len(tools[0].description) + expected = 18 if calling_type == "no_call" else 87 + len(tools[0].description) chunks = list( model.invoke("model", credentials, prompts, tools=tools, stream=stream) @@ -1287,7 +1305,7 @@ def test_completion_stream_handles_empty_text_and_partial_usage( http_server: HTTPServerFactory, ) -> None: frames = [ - {"choices": [{"text": ""}]}, + {"choices": [{"text": "", "finish_reason": ""}]}, {"choices": [{"text": "An"}]}, {"choices": [{"text": ""}]}, {"choices": [{"text": "swer"}]}, @@ -1295,6 +1313,7 @@ def test_completion_stream_handles_empty_text_and_partial_usage( "choices": [{"text": "", "finish_reason": "length"}], "usage": {"prompt_tokens": 0}, }, + {"choices": [{"text": "", "finish_reason": ""}]}, ] body = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) server = http_server([Response(body=(body + "data: [DONE]\n\n").encode())]) @@ -1323,8 +1342,24 @@ def test_completion_stream_handles_empty_text_and_partial_usage( @pytest.mark.parametrize("stream", [False, True]) -def test_missing_usage_can_estimate_tools_with_integer_enum_parameters( - http_server: HTTPServerFactory, stream: bool +@pytest.mark.parametrize( + "parameters", + [ + {}, + {"properties": {}}, + {"type": ["object"]}, + {"type": "object", "properties": {"allowed": True, "denied": False}}, + {"allOf": [{"type": "object"}]}, + { + "type": "object", + "properties": {"attempts": {"type": "integer", "enum": [1, 2]}}, + "required": ["attempts"], + }, + ], + ids=["empty", "implicit-type", "type-array", "boolean-schema", "all-of", "enum"], +) +def test_missing_usage_can_estimate_tool_schemas( + http_server: HTTPServerFactory, stream: bool, parameters: dict ) -> None: response = ( Response( @@ -1340,22 +1375,21 @@ def test_missing_usage_can_estimate_tools_with_integer_enum_parameters( tool = PromptMessageTool( name="configure", description="Choose the number of attempts", - parameters={ - "type": "object", - "properties": {"attempts": {"type": "integer", "enum": [1, 2]}}, - "required": ["attempts"], - }, + parameters=parameters, ) + credentials = { + "mode": "chat", + "endpoint_url": server.url, + "function_calling_type": "tool_call", + } + prompts = [UserPromptMessage(content="Configure two attempts")] + token_budget = model.get_num_tokens("model", credentials, prompts, [tool]) chunks = list( model.invoke( "model", - { - "mode": "chat", - "endpoint_url": server.url, - "function_calling_type": "tool_call", - }, - [UserPromptMessage(content="Configure two attempts")], + credentials, + prompts, tools=[tool], stream=stream, ) @@ -1365,7 +1399,8 @@ def test_missing_usage_can_estimate_tools_with_integer_enum_parameters( "".join(chunk.delta.message.get_text_content() for chunk in chunks) == "Answer" ) assert chunks[-1].delta.usage is not None - assert chunks[-1].delta.usage.prompt_tokens > 0 + assert chunks[-1].delta.usage.prompt_tokens == token_budget > 0 + assert chunks[-1].delta.usage.completion_tokens == 1 payload = json.loads(server.requests[0].body) assert payload["tools"] == [{"type": "function", "function": tool.model_dump()}] assert model.started_at == 0 diff --git a/tests/interfaces/model/openai_compatible/test_model_configuration.py b/tests/interfaces/model/openai_compatible/test_model_configuration.py index dcd610ef..828de2db 100644 --- a/tests/interfaces/model/openai_compatible/test_model_configuration.py +++ b/tests/interfaces/model/openai_compatible/test_model_configuration.py @@ -185,6 +185,22 @@ def test_token_budget_accepts_json_schema_enum_values(parameter: dict) -> None: ) +def test_token_budget_counts_nested_schema_content() -> None: + model = OAICompatLargeLanguageModel([]) + prompts = [UserPromptMessage(content="Configure the input")] + tool = PromptMessageTool( + name="configure", + description="Configure the input", + parameters={"allOf": [{"type": "object"}]}, + ) + baseline = model.get_num_tokens("model", {}, prompts, [tool]) + tool.parameters["allOf"][0]["description"] = ( + "A detailed explanation of the input. " * 100 + ) + + assert model.get_num_tokens("model", {}, prompts, [tool]) > baseline + 100 + + def test_invalid_custom_model_mode_reports_a_configuration_error() -> None: with pytest.raises(ValueError, match="invalid-mode"): OAICompatLargeLanguageModel([]).get_model_schema( From ef9b911f3f30749c5b96ddf5f71715f5cddc1daa Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Mon, 2 Nov 2026 02:08:49 +0000 Subject: [PATCH 2/4] fix(model): validate parsed replies and tokenize long inputs --- README.md | 3 + src/dify_plugin/interfaces/model/ai_model.py | 26 ++-- .../interfaces/model/openai_compatible/llm.py | 36 +++-- .../model/openai_compatible/test_http_llm.py | 125 ++++++++++++++++++ .../model/openai_compatible/test_llm.py | 14 +- .../test_model_configuration.py | 20 ++- 6 files changed, 183 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 81710d59..99b2d901 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,10 @@ Provider counts, including zero, take precedence; a missing or null count is est Completion counts the text sent from the first prompt message, while chat counts the complete message history and the tool definitions sent with the request. Output estimates include text and tool calls, including responses with no text content. Tool parameter estimates include the complete JSON schema, including optional and nested fields. +Long inputs use tokenizer chunks of at most 100,000 characters instead of switching to character counts; this shared estimator also serves embedding input budgets. +Token merges across chunk boundaries remain approximate. Tool-only non-streaming replies use an empty string for absent or null text so the daemon can decode the result. +Replies without a content field must contain a parsed tool call; an unrecognized tool entry does not make an empty reply valid. These local estimates use the SDK tokenizer and pricing configuration; they do not represent the provider's billing record. The public `get_num_tokens()` budget for completion also counts only the first prompt message and any supplied tool definitions. diff --git a/src/dify_plugin/interfaces/model/ai_model.py b/src/dify_plugin/interfaces/model/ai_model.py index 2baaf358..966a1dda 100644 --- a/src/dify_plugin/interfaces/model/ai_model.py +++ b/src/dify_plugin/interfaces/model/ai_model.py @@ -386,21 +386,23 @@ def _get_num_tokens_by_gpt2(self, text: str) -> int: The return value. """ - # ENHANCEMENT: - # to avoid performance issue, do not calculate the number of tokens - # for too long text - # only to promise text length is less than TOKEN_ESTIMATION_TEXT_LIMIT - if len(text) >= TOKEN_ESTIMATION_TEXT_LIMIT: - return len(text) - - # check if gevent is patched to main thread import tiktoken # ruff:ignore[import-outside-top-level] + def count_tokens() -> int: + encoding = tiktoken.encoding_for_model("gpt2") + # shortcut: chunk boundaries affect token merges; revisit for exact counts. + return sum( + len( + encoding.encode_ordinary( + text[offset : offset + TOKEN_ESTIMATION_TEXT_LIMIT] + ) + ) + for offset in range(0, len(text), TOKEN_ESTIMATION_TEXT_LIMIT) + ) + if socket.socket is gevent.socket.socket: # using gevent real thread to avoid blocking main thread - result = threadpool.spawn( - lambda: len(tiktoken.encoding_for_model("gpt2").encode_ordinary(text)) - ) + result = threadpool.spawn(count_tokens) return result.get(block=True) or 0 - return len(tiktoken.encoding_for_model("gpt2").encode_ordinary(text)) + return count_tokens() diff --git a/src/dify_plugin/interfaces/model/openai_compatible/llm.py b/src/dify_plugin/interfaces/model/openai_compatible/llm.py index 6d130f37..cdbe6cfc 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/llm.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/llm.py @@ -1003,39 +1003,37 @@ def _handle_generate_response( message_id = response_json.get("id") response_content = "" - tool_calls = None + tool_calls = [] function_calling_type = credentials.get("function_calling_type", "no_call") if completion_type is LLMMode.CHAT: response_message = output.get("message", {}) - response_content = response_message.get("content") - if response_content is None: - response_content = "" - if function_calling_type == "tool_call": - tool_calls = output.get("message", {}).get("tool_calls") - elif function_calling_type == "function_call": - tool_calls = output.get("message", {}).get("function_call") + if function_calling_type == "tool_call" and ( + response_tool_calls := response_message.get("tool_calls") + ): + tool_calls = self._extract_response_tool_calls(response_tool_calls) + elif function_calling_type == "function_call" and ( + response_function_call := response_message.get("function_call") + ): + tool_call = self._extract_response_function_call(response_function_call) + if tool_call is not None: + tool_calls = [tool_call] + if "content" not in response_message and not tool_calls: msg = "Model response has neither content nor tool calls" raise InvokeError(msg) + response_content = response_message.get("content") + if response_content is None: + response_content = "" + elif completion_type is LLMMode.COMPLETION: response_content = output["text"] assistant_message = AssistantPromptMessage( content=response_content, - tool_calls=[], + tool_calls=tool_calls, ) - if tool_calls: - if function_calling_type == "tool_call": - assistant_message.tool_calls = self._extract_response_tool_calls( - tool_calls, - ) - elif function_calling_type == "function_call": - assistant_message.tool_calls = [ - self._extract_response_function_call(tool_calls), - ] - # transform response return LLMResult( id=message_id, diff --git a/tests/interfaces/model/openai_compatible/test_http_llm.py b/tests/interfaces/model/openai_compatible/test_http_llm.py index d266f8dc..fb607b74 100644 --- a/tests/interfaces/model/openai_compatible/test_http_llm.py +++ b/tests/interfaces/model/openai_compatible/test_http_llm.py @@ -7,6 +7,7 @@ import gevent import pytest +import tiktoken from gevent.event import Event from dify_plugin.entities import I18nObject @@ -493,6 +494,64 @@ def test_non_stream_tool_calls_preserve_arguments_and_request_schema( assert payload["functions"] == [tool.model_dump()] +@pytest.mark.parametrize( + "invalid_call", + [ + {}, + {"function": None}, + {"function": {}}, + {"function": []}, + {"function": ""}, + {"function": False}, + {"function": 0}, + ], +) +@pytest.mark.parametrize("include_valid_call", [False, True]) +def test_non_stream_missing_content_requires_a_parsed_tool_call( + http_server: HTTPServerFactory, + http_responses: tuple[list[HTTPResponse], Event], + invalid_call: dict, + include_valid_call: bool, +) -> None: + responses, _ready = http_responses + valid_call = { + "id": "call-1", + "type": "function", + "function": {"name": "weather", "arguments": '{"city":"Paris"}'}, + } + calls = [invalid_call, valid_call] if include_valid_call else [invalid_call] + server = http_server([ + Response( + body=json.dumps({"choices": [{"message": {"tool_calls": calls}}]}).encode() + ) + ]) + model = OAICompatLargeLanguageModel([]) + invocation = model.invoke( + "model", + { + "mode": "chat", + "endpoint_url": server.url, + "function_calling_type": "tool_call", + }, + [UserPromptMessage(content="Weather?")], + stream=False, + ) + + if include_valid_call: + chunks = list(invocation) + assert isinstance(chunks[-1].delta.message.content, str) + assert not chunks[-1].delta.message.content + assert [call.model_dump() for call in chunks[-1].delta.message.tool_calls] == [ + valid_call + ] + else: + with pytest.raises(InvokeError, match="neither content nor tool calls"): + list(invocation) + + assert responses[0].raw.closed + assert model.started_at == 0 + + @pytest.mark.parametrize( ("delimiter", "configured"), [("\n\n", "\n\n"), ("\n", r"\n"), ("|", "|")] ) @@ -1404,3 +1463,69 @@ def test_missing_usage_can_estimate_tool_schemas( payload = json.loads(server.requests[0].body) assert payload["tools"] == [{"type": "function", "function": tool.model_dump()}] assert model.started_at == 0 + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("calling_type", ["tool_call", "function_call"]) +def test_large_tool_schemas_keep_token_and_price_estimates_consistent( + http_server: HTTPServerFactory, stream: bool, calling_type: str +) -> None: + choice = ( + {"delta": {"content": "Answer"}, "finish_reason": "stop"} + if stream + else {"message": {"content": "Answer"}} + ) + body = json.dumps({"choices": [choice], "usage": {"completion_tokens": 0}}) + if stream: + body = f"data: {body}\n\ndata: [DONE]\n\n" + server = http_server([Response(body=body.encode()) for _ in range(2)]) + model = OAICompatLargeLanguageModel([]) + credentials = { + "mode": "chat", + "endpoint_url": server.url, + "function_calling_type": calling_type, + "input_price": "2", + "unit": "0.001", + } + prompts = [UserPromptMessage(content="Configure the input")] + tool = PromptMessageTool( + name="configure", description="Configure the input", parameters={} + ) + encoding = tiktoken.encoding_for_model("gpt2") + overhead = model.get_num_tokens("model", credentials, prompts, [tool]) - len( + encoding.encode_ordinary("{}") + ) + + for field_count in (310, 311): + tool.parameters = { + "type": "object", + "properties": { + f"field_{index}": { + "type": "string", + "description": "A detailed optional input. " * 10, + } + for index in range(field_count) + }, + } + serialized_schema = json.dumps(tool.parameters, ensure_ascii=False) + assert (len(serialized_schema) < 100_000) == (field_count == 310) + expected = overhead + len(encoding.encode_ordinary(serialized_schema)) + budget = model.get_num_tokens("model", credentials, prompts, [tool]) + assert budget == pytest.approx(expected, rel=0.001) + + chunks = list( + model.invoke("model", credentials, prompts, tools=[tool], stream=stream) + ) + usage = chunks[-1].delta.usage + assert usage is not None + assert usage.prompt_tokens == budget + assert usage.completion_tokens == 0 + assert usage.prompt_price == Decimal(budget) * Decimal("0.002") + payload = json.loads(server.requests[-1].body) + request_tool = ( + payload["tools"][0]["function"] + if calling_type == "tool_call" + else payload["functions"][0] + ) + assert request_tool["parameters"] == tool.parameters + assert model.started_at == 0 diff --git a/tests/interfaces/model/openai_compatible/test_llm.py b/tests/interfaces/model/openai_compatible/test_llm.py index 41edbe4a..f3de047b 100644 --- a/tests/interfaces/model/openai_compatible/test_llm.py +++ b/tests/interfaces/model/openai_compatible/test_llm.py @@ -7,6 +7,7 @@ import gevent.socket import pytest import requests +import tiktoken from gevent.threadpool import ThreadPool from dify_plugin.entities.model.message import ( @@ -319,8 +320,9 @@ def test_num_tokens_counts_non_user_text_parts_like_plain_text( @pytest.mark.parametrize("use_threadpool", [False, True]) +@pytest.mark.parametrize("repetitions", [1, 5000]) def test_num_tokens_treats_special_token_spelling_as_text( - use_threadpool: bool, monkeypatch: pytest.MonkeyPatch + use_threadpool: bool, repetitions: int, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr( ai_model, @@ -329,7 +331,8 @@ def test_num_tokens_treats_special_token_spelling_as_text( ) pool = ThreadPool(1) monkeypatch.setattr(ai_model, "threadpool", pool, raising=False) - content = [TextPromptMessageContent(data="Explain the <|endoftext|> token.")] + text = "Explain the <|endoftext|> token. " * repetitions + content = [TextPromptMessageContent(data=text)] messages = [ SystemPromptMessage(content=content), AssistantPromptMessage(content=content), @@ -340,7 +343,12 @@ def test_num_tokens_treats_special_token_spelling_as_text( credentials = {"function_calling_type": "tool_call"} count = model.get_num_tokens("model", credentials, messages) empty = [message.model_copy(update={"content": ""}) for message in messages] - assert count > model.get_num_tokens("model", credentials, empty) + expected = len(tiktoken.encoding_for_model("gpt2").encode_ordinary(text)) * len( + messages + ) + assert count - model.get_num_tokens( + "model", credentials, empty + ) == pytest.approx(expected, rel=0.001) finally: pool.kill() diff --git a/tests/interfaces/model/openai_compatible/test_model_configuration.py b/tests/interfaces/model/openai_compatible/test_model_configuration.py index 828de2db..ffbcdefc 100644 --- a/tests/interfaces/model/openai_compatible/test_model_configuration.py +++ b/tests/interfaces/model/openai_compatible/test_model_configuration.py @@ -1,6 +1,7 @@ from decimal import Decimal import pytest +import tiktoken from dify_plugin.entities.model import ModelFeature, ModelPropertyKey from dify_plugin.entities.model.message import ( @@ -145,17 +146,22 @@ def test_embedding_token_estimates_preserve_input_order_and_empty_texts() -> Non assert model.get_num_tokens("model", {}, ["Hello", "", "Hello world"]) == [1, 0, 2] -def test_large_prompts_have_positive_nondecreasing_token_estimates() -> None: +@pytest.mark.parametrize("pattern", ["hello world ", "δ½ ε₯½πŸ˜€ <|endoftext|> "]) +@pytest.mark.parametrize("size", [99_999, 100_001, 210_003]) +def test_large_text_estimates_track_the_tokenizer(pattern: str, size: int) -> None: + text = (pattern * (size // len(pattern) + 1))[:size] + expected = len(tiktoken.encoding_for_model("gpt2").encode_ordinary(text)) model = OAICompatLargeLanguageModel([]) + count = model.get_num_tokens( + "model", {"mode": "completion"}, [UserPromptMessage(content=text)] + ) - counts = [ - model.get_num_tokens("model", {}, [UserPromptMessage(content="a" * size)]) - for size in (100_000, 200_000) + assert count == pytest.approx(expected, rel=0.001) + assert OAICompatEmbeddingModel([]).get_num_tokens("model", {}, [text, ""]) == [ + count, + 0, ] - assert all(type(count) is int and count > 0 for count in counts) - assert counts[1] >= counts[0] - @pytest.mark.parametrize( "parameter", From b55713bb07c104cd20b6cd27f1ed8accf1d9252c Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Mon, 2 Nov 2026 02:18:36 +0000 Subject: [PATCH 3/4] fix(embedding): measure truncated inputs against token budgets --- README.md | 3 +- .../model/openai_compatible/text_embedding.py | 31 ++++--- .../test_http_embedding_rerank.py | 89 +++++++++++++++++++ 3 files changed, 112 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 99b2d901..cd7dfe36 100644 --- a/README.md +++ b/README.md @@ -86,7 +86,8 @@ Output estimates include text and tool calls, including responses with no text c Tool parameter estimates include the complete JSON schema, including optional and nested fields. Long inputs use tokenizer chunks of at most 100,000 characters instead of switching to character counts; this shared estimator also serves embedding input budgets. Token merges across chunk boundaries remain approximate. -Tool-only non-streaming replies use an empty string for absent or null text so the daemon can decode the result. +Embedding inputs that exceed the configured context are shortened to a measured text prefix within the estimated token budget. +Non-streaming chat replies use an empty string for absent or null text so the daemon can decode the result. Replies without a content field must contain a parsed tool call; an unrecognized tool entry does not make an empty reply valid. These local estimates use the SDK tokenizer and pricing configuration; they do not represent the provider's billing record. The public `get_num_tokens()` budget for completion also counts only the first prompt message and any supplied tool definitions. diff --git a/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py b/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py index 1f6db3f4..bd59fbe9 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/text_embedding.py @@ -140,21 +140,32 @@ def _invoke( max_chunks = self._get_max_chunks(model, credentials) inputs = [] - indices = [] used_tokens = 0 - for i, text in enumerate(texts): - # Here token count is only an approximation based on the GPT2 tokenizer - # Better token estimation would make chunking more precise. + for text in texts: num_tokens = self._get_num_tokens_by_gpt2(text) - - if num_tokens >= context_size: - cutoff = int((len(text) * context_size) // num_tokens) - # if num tokens is larger than context length, only use the start - inputs.append(text[0:cutoff]) + if num_tokens > context_size: + lower, upper = 0, len(text) + # Keep fitting estimates: token merges can make prefix counts decrease. + for estimate in ( + min(context_size, len(text)), + len(text) * context_size // num_tokens, + ): + if ( + lower < estimate < upper + and self._get_num_tokens_by_gpt2(text[:estimate]) + <= context_size + ): + lower = estimate + while lower + 1 < upper: + middle = (lower + upper) // 2 + if self._get_num_tokens_by_gpt2(text[:middle]) <= context_size: + lower = middle + else: + upper = middle + inputs.append(text[:lower]) else: inputs.append(text) - indices += [i] batched_embeddings = [] dimension = None diff --git a/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py b/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py index 79aa0ed2..d283baec 100644 --- a/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py +++ b/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py @@ -4,6 +4,7 @@ import pytest import requests +import tiktoken from pytest_mock import MockerFixture from dify_plugin.entities import I18nObject @@ -86,6 +87,94 @@ def test_embedding_batches_truncates_and_accumulates_usage( assert all(payload["user"] == "user-1" for payload in payloads) +@pytest.mark.parametrize( + ("text", "context_size"), + [ + ("a!" * 10_000 + "hello world " * 10_000, 8192), + ( + "| a | b | c | d | e | f |\n" * 1000 + + ( + "The document describes the current configuration " + "and the expected output.\n" + ) + * 2000, + 8192, + ), + ("δ½ ε₯½πŸ˜€ " * 30 + "hello world " * 100, 8), + ("hello world", 2), + ("", 2), + ("πŸ˜€", 1), + ], + ids=["mixed-long", "table-long", "unicode", "already-fits", "empty", "no-fit"], +) +def test_embedding_sends_a_prefix_within_the_estimated_context( + http_server: HTTPServerFactory, text: str, context_size: int +) -> None: + server = http_server([ + Response(body=b'{"data":[{"embedding":[1,2]}],"usage":{"total_tokens":0}}') + ]) + model = OAICompatEmbeddingModel([]) + result = model.invoke( + "model", + {"endpoint_url": server.url, "context_size": str(context_size)}, + [text], + ) + + sent = json.loads(server.requests[0].body)["input"][0] + encoding = tiktoken.encoding_for_model("gpt2") + assert text.startswith(sent) + assert len(encoding.encode_ordinary(sent)) <= context_size + if len(encoding.encode_ordinary(text)) <= context_size: + assert sent == text + else: + assert len(encoding.encode_ordinary(text[: len(sent) + 1])) > context_size + assert result.embeddings == [[1, 2]] + assert result.usage.total_tokens == 0 + assert model.started_at == 0 + + +@pytest.mark.parametrize( + ("text", "context_size", "minimum_prefix"), + [ + ( + "a\n" * 255 + "internationalization!!", + 512, + "a\n" * 255 + "internationalization", + ), + ("θ¦šι†’foo!", 1, "θ¦šι†’"), + ( + "a\n" * 254 + "\nπŸ˜€θ¦šι†’" + "hello world " * 8300, + 512, + "a\n" * 254 + "\nπŸ˜€θ¦šι†’", + ), + ], + ids=["word-merge", "unicode-merge", "long-input-merge"], +) +def test_embedding_preserves_fitting_estimates_across_token_merges( + http_server: HTTPServerFactory, + text: str, + context_size: int, + minimum_prefix: str, +) -> None: + server = http_server([ + Response(body=b'{"data":[{"embedding":[1,2]}],"usage":{"total_tokens":0}}') + ]) + model = OAICompatEmbeddingModel([]) + model.invoke( + "model", + {"endpoint_url": server.url, "context_size": str(context_size)}, + [text], + ) + + sent = json.loads(server.requests[0].body)["input"][0] + assert sent.startswith(minimum_prefix) + assert text.startswith(sent) + assert ( + len(tiktoken.encoding_for_model("gpt2").encode_ordinary(sent)) <= context_size + ) + assert model.started_at == 0 + + def test_embedding_empty_input_does_not_request_the_server( http_server: HTTPServerFactory, embedding: OAICompatEmbeddingModel ) -> None: From f3492782ab33bb23e9d612b06cb9b693ec8e5483 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Mon, 2 Nov 2026 02:39:23 +0000 Subject: [PATCH 4/4] test(embedding): catch loss of fitting character prefixes --- .../model/openai_compatible/test_http_embedding_rerank.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py b/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py index d283baec..0424fa8f 100644 --- a/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py +++ b/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py @@ -143,7 +143,7 @@ def test_embedding_sends_a_prefix_within_the_estimated_context( ), ("θ¦šι†’foo!", 1, "θ¦šι†’"), ( - "a\n" * 254 + "\nπŸ˜€θ¦šι†’" + "hello world " * 8300, + "a\n" * 254 + "\nπŸ˜€θ¦šι†’" + "hello world " * 8291, 512, "a\n" * 254 + "\nπŸ˜€θ¦šι†’", ),