fix(model): validate responses and stabilize token estimates

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

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()

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,

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

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: