diff --git a/README.md b/README.md index 534cb9d3..cd7dfe36 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,12 @@ 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. +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. +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. @@ -97,7 +103,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/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 72cf5556..cdbe6cfc 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 @@ -1003,37 +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 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, @@ -1263,43 +1263,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/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..0424fa8f 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 " * 8291, + 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: diff --git a/tests/interfaces/model/openai_compatible/test_http_llm.py b/tests/interfaces/model/openai_compatible/test_http_llm.py index 7eef2786..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 @@ -29,6 +30,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 +471,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 @@ -487,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"), ("|", "|")] ) @@ -694,6 +759,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 +785,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 +797,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 +815,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 +919,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 +1364,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 +1372,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 +1401,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 +1434,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 +1458,74 @@ 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 + + +@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 dcd610ef..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", @@ -185,6 +191,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(