diff --git a/README.md b/README.md index 83d769f6..534cb9d3 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,25 @@ Use `message.get_text_content()` when a consumer needs plain text, and retain th `ToolPromptMessage.is_empty()` treats a tool call ID as content, including when the tool result is empty. The invoking Agent/plugin, provider, daemon, and host must all preserve these fields for replay across tool turns. +## OpenAI-Compatible Completion Input + +In SDK 0.11.0 (unreleased), the shared `OAICompatLargeLanguageModel` completion mode accepts text from the first prompt message. +Text content blocks are concatenated in order into one prompt string, including empty lists as an empty string. +Image, audio, video, or document blocks in that prompt raise an SDK invocation error before sending a request, including when mixed with text. +For multimodal input, use chat mode with a model that supports the input modality and message role. + +## OpenAI-Compatible Usage Estimates + +In SDK 0.11.0 (unreleased), streaming and non-streaming responses estimate missing usage fields with the same counters. +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. +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. + +Provider overrides of `_handle_generate_response()`, `_handle_generate_stream_response()`, and `_create_final_llm_result_chunk()` must accept and forward the optional `tools` parameter in SDK 0.11.0. +Overrides of `_num_tokens_from_string()` must follow the SDK signature `(text, tools=None)`; the method does not take a model name. + ## OpenAI-Compatible Streaming In SDK 0.11.0 (unreleased), `OAICompatLargeLanguageModel` preserves SSE data alongside comments and unknown fields. diff --git a/src/dify_plugin/interfaces/model/openai_compatible/llm.py b/src/dify_plugin/interfaces/model/openai_compatible/llm.py index 06fcf6c0..72cf5556 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/llm.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/llm.py @@ -33,6 +33,7 @@ LLMResult, LLMResultChunk, LLMResultChunkDelta, + LLMUsage, ) from dify_plugin.entities.model.message import ( AssistantPromptMessage, @@ -301,6 +302,10 @@ def get_num_tokens( The return value. """ del model + if credentials.get("mode") == LLMMode.COMPLETION.value: + return self._num_tokens_from_string( + prompt_messages[0].get_text_content(), tools + ) return self._num_tokens_from_messages(prompt_messages, tools, credentials) def _request_credentials_validation( @@ -650,7 +655,20 @@ def _generate( ] elif completion_type is LLMMode.COMPLETION: endpoint_url = self._join_endpoint_url(endpoint_url, "completions") - data["prompt"] = prompt_messages[0].content + prompt = prompt_messages[0] + if isinstance(prompt.content, list): + if any( + part.type != PromptMessageContentType.TEXT + for part in prompt.content + ): + msg = ( + "Completion mode only supports text input. " + "Use a compatible chat model for multimodal input." + ) + raise ValueError(msg) + data["prompt"] = prompt.get_text_content() + else: + data["prompt"] = prompt.content else: msg = "Unsupported completion type for model configuration." raise ValueError(msg) @@ -669,7 +687,7 @@ def _generate( for tool in tools ] elif function_calling_type == "tool_call": - data["tool_choice"] = "auto" + data.setdefault("tool_choice", "auto") formatted_tools.extend( PromptMessageFunction(function=tool).model_dump() for tool in tools @@ -683,6 +701,9 @@ def _generate( if user: data["user"] = user + if function_calling_type not in {"tool_call", "function_call"}: + tools = None + response = requests.post( endpoint_url, headers=headers, @@ -713,6 +734,7 @@ def _generate( credentials, response, prompt_messages, + tools=tools, ) with response: @@ -721,54 +743,72 @@ def _generate( credentials, response, prompt_messages, + tools=tools, ) - def _create_final_llm_result_chunk( + def _get_response_usage( self, - index: int, - message: AssistantPromptMessage, - finish_reason: str | None, - usage: dict | None, model: str, - prompt_messages: list[PromptMessage], credentials: dict, - full_content: str, - ) -> LLMResultChunk: - # calculate num tokens + prompt_messages: list[PromptMessage], + message: AssistantPromptMessage, + usage: dict | None, + tools: list[PromptMessageTool] | None = None, + ) -> LLMUsage: + """Estimate only the token counts omitted by the provider.""" prompt_tokens = (usage or {}).get("prompt_tokens") if prompt_tokens is None: if credentials.get("mode") == LLMMode.COMPLETION.value: prompt_tokens = self._num_tokens_from_string( - prompt_messages[0].content or "", + prompt_messages[0].get_text_content(), tools ) else: prompt_tokens = self._num_tokens_from_messages( - prompt_messages, - credentials=credentials, + prompt_messages, tools, credentials ) completion_tokens = (usage or {}).get("completion_tokens") if completion_tokens is None: - completion_tokens = self._num_tokens_from_string(text=full_content) + completion_tokens = self._num_tokens_from_string(message.get_text_content()) completion_tokens += sum( - self._num_tokens_for_tool_call(tool_call.model_dump()) + self._num_tokens_for_function_call(tool_call.function.model_dump()) + if credentials.get("function_calling_type") == "function_call" + else self._num_tokens_for_tool_call(tool_call.model_dump()) for tool_call in message.tool_calls ) - # transform usage - usage = self._calc_response_usage( + return self._calc_response_usage( model, credentials, prompt_tokens, completion_tokens, ) + def _create_final_llm_result_chunk( + self, + index: int, + message: AssistantPromptMessage, + finish_reason: str | None, + usage: dict | None, + model: str, + prompt_messages: list[PromptMessage], + credentials: dict, + full_content: str, + tools: list[PromptMessageTool] | None = None, + ) -> LLMResultChunk: return LLMResultChunk( model=model, delta=LLMResultChunkDelta( index=index, message=message, finish_reason=finish_reason, - usage=usage, + usage=self._get_response_usage( + model, + credentials, + prompt_messages, + message.model_copy(update={"content": full_content}), + usage, + tools, + ), ), ) @@ -778,6 +818,7 @@ def _handle_generate_stream_response( credentials: dict, response: requests.Response, prompt_messages: list[PromptMessage], + tools: list[PromptMessageTool] | None = None, ) -> Generator: """Handle llm stream response @@ -943,6 +984,7 @@ def _handle_generate_stream_response( credentials=credentials, prompt_messages=prompt_messages, full_content=full_assistant_content, + tools=tools, ) def _handle_generate_response( @@ -951,6 +993,7 @@ def _handle_generate_response( credentials: dict, response: requests.Response, prompt_messages: list[PromptMessage], + tools: list[PromptMessageTool] | None = None, ) -> LLMResult: response_json: dict = response.json() @@ -963,11 +1006,15 @@ def _handle_generate_response( tool_calls = None function_calling_type = credentials.get("function_calling_type", "no_call") if completion_type is LLMMode.CHAT: - response_content = output.get("message", {})["content"] + 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 "content" not in response_message and not tool_calls: + msg = "Model response has neither content nor tool calls" + raise InvokeError(msg) elif completion_type is LLMMode.COMPLETION: response_content = output["text"] @@ -987,44 +1034,19 @@ def _handle_generate_response( self._extract_response_function_call(tool_calls), ] - usage = response_json.get("usage") - if usage: - # transform usage - prompt_tokens = usage["prompt_tokens"] - completion_tokens = usage["completion_tokens"] - else: - # calculate num tokens - if prompt_messages[0].content is None: - msg = "Prompt message content is required" - raise ValueError(msg) - prompt_content = cast("str", prompt_messages[0].content) - prompt_tokens = self._num_tokens_from_string( - model, - prompt_content, - ) - if assistant_message.content is None: - msg = "Assistant message content is required" - raise ValueError(msg) - assistant_content = cast("str", assistant_message.content) - completion_tokens = self._num_tokens_from_string( - model, - assistant_content, - ) - - # transform usage - usage = self._calc_response_usage( - model, - credentials, - prompt_tokens, - completion_tokens, - ) - # transform response return LLMResult( id=message_id, model=response_json.get("model", model), message=assistant_message, - usage=usage, + usage=self._get_response_usage( + model, + credentials, + prompt_messages, + assistant_message, + response_json.get("usage"), + tools, + ), ) def _convert_prompt_message_to_dict( diff --git a/tests/interfaces/model/openai_compatible/test_http_llm.py b/tests/interfaces/model/openai_compatible/test_http_llm.py index 55ec0e9b..7eef2786 100644 --- a/tests/interfaces/model/openai_compatible/test_http_llm.py +++ b/tests/interfaces/model/openai_compatible/test_http_llm.py @@ -14,9 +14,14 @@ from dify_plugin.entities.model.llm import LLMResult, LLMResultChunk from dify_plugin.entities.model.message import ( AssistantPromptMessage, + AudioPromptMessageContent, + DocumentPromptMessageContent, + ImagePromptMessageContent, PromptMessage, PromptMessageTool, + TextPromptMessageContent, UserPromptMessage, + VideoPromptMessageContent, ) from dify_plugin.errors.model import CredentialsValidateFailedError, InvokeError from dify_plugin.interfaces.model.openai_compatible import llm @@ -31,8 +36,8 @@ pytestmark = pytest.mark.integration -class LegacyResponseLLM(OAICompatLargeLanguageModel): - """Use the response hooks implemented by existing official plugins.""" +class ResponseHookLLM(OAICompatLargeLanguageModel): + """Forward request tools through provider response hooks.""" def _handle_generate_response( self, @@ -40,9 +45,10 @@ def _handle_generate_response( credentials: dict, response: HTTPResponse, prompt_messages: list[PromptMessage], + tools: list[PromptMessageTool] | None = None, ) -> LLMResult: return super()._handle_generate_response( - model, credentials, response, prompt_messages + model, credentials, response, prompt_messages, tools ) def _handle_generate_stream_response( @@ -51,9 +57,10 @@ def _handle_generate_stream_response( credentials: dict, response: HTTPResponse, prompt_messages: list[PromptMessage], + tools: list[PromptMessageTool] | None = None, ) -> Generator[LLMResultChunk]: yield from super()._handle_generate_stream_response( - model, credentials, response, prompt_messages + model, credentials, response, prompt_messages, tools ) def _create_final_llm_result_chunk( @@ -66,6 +73,7 @@ def _create_final_llm_result_chunk( prompt_messages: list[PromptMessage], credentials: dict, full_content: str, + tools: list[PromptMessageTool] | None = None, ) -> LLMResultChunk: return super()._create_final_llm_result_chunk( index, @@ -76,6 +84,7 @@ def _create_final_llm_result_chunk( prompt_messages, credentials, full_content, + tools, ) @@ -172,6 +181,117 @@ def test_invalid_structured_schema_fails_before_sending_a_request( assert structured_model.started_at == 0 +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("reported_usage", [False, True]) +@pytest.mark.parametrize( + ("content", "expected_prompt"), + [ + ("plain text", "plain text"), + ("", ""), + (None, None), + ([], ""), + ( + [ + TextPromptMessageContent(data="你好😀 "), + TextPromptMessageContent(data=""), + TextPromptMessageContent(data="<|endoftext|>"), + ], + "你好😀 <|endoftext|>", + ), + ], +) +def test_completion_sends_text_content_as_one_prompt( + http_server: HTTPServerFactory, + content: str | list[TextPromptMessageContent] | None, + expected_prompt: str | None, + stream: bool, + reported_usage: bool, +) -> None: + payload = json.dumps({ + "choices": [{"text": "Answer", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 2} + if reported_usage + else None, + }) + body = f"data: {payload}\n\ndata: [DONE]\n\n" if stream else payload + server = http_server([Response(body=body.encode())]) + message = UserPromptMessage(content=content) + original_message = message.model_copy(deep=True) + model = OAICompatLargeLanguageModel([]) + + chunks = list( + model.invoke( + "model", + {"mode": "completion", "endpoint_url": server.url}, + [message], + stream=stream, + ) + ) + + assert server.requests[0].path == "/completions" + assert json.loads(server.requests[0].body)["prompt"] == expected_prompt + assert ( + "".join(chunk.delta.message.get_text_content() for chunk in chunks) == "Answer" + ) + usage = chunks[-1].delta.usage + assert usage is not None + assert usage.completion_tokens == (2 if reported_usage else 1) + if not reported_usage and not expected_prompt: + assert usage.prompt_tokens == 0 + assert message == original_message + assert model.started_at == 0 + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize( + "content_type", + [ + ImagePromptMessageContent, + AudioPromptMessageContent, + VideoPromptMessageContent, + DocumentPromptMessageContent, + ], +) +def test_completion_rejects_media_before_sending_a_request( + http_server: HTTPServerFactory, + content_type: type[ + ImagePromptMessageContent + | AudioPromptMessageContent + | VideoPromptMessageContent + | DocumentPromptMessageContent + ], + mixed: bool, + stream: bool, +) -> None: + media = content_type( + url="https://example.com/attachment", + format="bin", + mime_type="application/octet-stream", + ) + content = ( + [TextPromptMessageContent(data="Describe this."), media] if mixed else [media] + ) + message = UserPromptMessage(content=content) + original_message = message.model_copy(deep=True) + server = http_server([]) + model = OAICompatLargeLanguageModel([]) + + with pytest.raises(InvokeError, match="Completion mode only supports text input"): + list( + model.invoke( + "model", + {"mode": "completion", "endpoint_url": server.url}, + [message], + stream=stream, + ) + ) + + assert server.requests == [] + assert message == original_message + assert model.started_at == 0 + + @pytest.mark.parametrize("mode", ["chat", "completion"]) def test_invoke_sends_json_and_reads_non_stream_response( http_server: HTTPServerFactory, mode: str @@ -237,34 +357,74 @@ def test_invoke_sends_json_and_reads_non_stream_response( @pytest.mark.parametrize("mode", ["chat", "completion"]) -def test_non_stream_usage_keeps_existing_fallback_estimates( +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("text_blocks", [False, True]) +@pytest.mark.parametrize( + "reported_usage", + [ + None, + {}, + {"prompt_tokens": 0}, + {"completion_tokens": 0}, + {"prompt_tokens": None, "completion_tokens": None}, + ], +) +def test_usage_counts_sent_text_and_history( http_server: HTTPServerFactory, - monkeypatch: pytest.MonkeyPatch, mode: str, + stream: bool, + text_blocks: bool, + reported_usage: dict | None, ) -> None: choice = ( - {"message": {"content": "Answer"}} if mode == "chat" else {"text": "Answer"} + {"delta" if stream else "message": {"content": "Answer"}} + if mode == "chat" + else {"text": "Answer"} ) - server = http_server([Response(body=json.dumps({"choices": [choice]}).encode())]) + choice["finish_reason"] = "stop" + payload = json.dumps({"choices": [choice], "usage": reported_usage}) + body = f"data: {payload}\n\ndata: [DONE]\n\n" if stream else payload + server = http_server([Response(body=body.encode())]) model = OAICompatLargeLanguageModel([]) - monkeypatch.setattr(model, "_get_num_tokens_by_gpt2", len) + prompts = [ + UserPromptMessage( + content=[ + TextPromptMessageContent(data="Hello"), + TextPromptMessageContent(data=" world"), + ] + if text_blocks + else "Hello world" + ), + AssistantPromptMessage(content="history"), + ] + credentials = { + "mode": mode, + "endpoint_url": server.url, + "input_price": "2", + "output_price": "4", + "unit": "0.001", + } - chunks = list( - model.invoke( - "model", - {"mode": mode, "endpoint_url": server.url}, - [ - UserPromptMessage(content="Question"), - AssistantPromptMessage(content="history"), - ], - stream=False, - ) - ) + chunks = list(model.invoke("model", credentials, prompts, stream=stream)) usage = chunks[-1].delta.usage assert usage is not None - assert usage.prompt_tokens == 285 - assert usage.completion_tokens == 215 + estimated_prompt_tokens = 15 if mode == "chat" else 2 + prompt_tokens = ( + 0 + if reported_usage and reported_usage.get("prompt_tokens") == 0 + else estimated_prompt_tokens + ) + completion_tokens = ( + 0 if reported_usage and reported_usage.get("completion_tokens") == 0 else 1 + ) + assert usage.prompt_tokens == prompt_tokens + assert usage.completion_tokens == completion_tokens + assert usage.prompt_price == Decimal(prompt_tokens) * Decimal("0.002") + assert usage.completion_price == Decimal(completion_tokens) * Decimal("0.004") + assert ( + model.get_num_tokens("model", credentials, prompts) == estimated_prompt_tokens + ) assert model.started_at == 0 @@ -307,21 +467,18 @@ def test_non_stream_tool_calls_preserve_arguments_and_request_schema( tools=[tool], stream=False, ) - if include_content and reported_usage: - chunks = list(invocation) - call = chunks[-1].delta.message.tool_calls[0] - assert chunks[-1].delta.message.content is None - assert call.function.name == "weather" - assert json.loads(call.function.arguments) == {"city": "北京"} - usage = chunks[-1].delta.usage - assert usage is not None + chunks = list(invocation) + call = chunks[-1].delta.message.tool_calls[0] + assert chunks[-1].delta.message.content is None + assert call.function.name == "weather" + assert json.loads(call.function.arguments) == {"city": "北京"} + usage = chunks[-1].delta.usage + assert usage is not None + if reported_usage: assert (usage.prompt_tokens, usage.completion_tokens) == (0, 0) else: - with pytest.raises(InvokeError) as exc_info: - list(invocation) - assert isinstance( - exc_info.value.__cause__, KeyError if not include_content else ValueError - ) + assert usage.prompt_tokens > len("Weather?Get weather") + assert usage.completion_tokens >= len(function["arguments"]) assert model.started_at == 0 payload = json.loads(server.requests[0].body) if calling_type == "tool_call": @@ -649,7 +806,7 @@ def test_stream_rejects_eof_without_completion_signal( @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("calling_type", ["tool_call", "function_call", "no_call"]) -def test_missing_usage_keeps_estimates_and_provider_response_overrides( +def test_missing_usage_counts_sent_tools_through_provider_hooks( http_server: HTTPServerFactory, monkeypatch: pytest.MonkeyPatch, stream: bool, @@ -668,7 +825,7 @@ def test_missing_usage_keeps_estimates_and_provider_response_overrides( else Response(body=b'{"choices":[{"message":{"content":"Answer"}}]}') ) server = http_server([response]) - model = LegacyResponseLLM([]) + model = ResponseHookLLM([]) monkeypatch.setattr(model, "_get_num_tokens_by_gpt2", len) prompts = [UserPromptMessage(content="Question")] tools = [ @@ -685,7 +842,7 @@ def test_missing_usage_keeps_estimates_and_provider_response_overrides( "input_price": "1", "unit": "0.001", } - expected = 18 if stream else 285 + expected = 18 if calling_type == "no_call" else 79 + len(tools[0].description) chunks = list( model.invoke("model", credentials, prompts, tools=tools, stream=stream) @@ -705,6 +862,39 @@ def test_missing_usage_keeps_estimates_and_provider_response_overrides( assert "functions" not in payload +def test_tool_definitions_preserve_an_explicit_tool_choice( + http_server: HTTPServerFactory, +) -> None: + server = http_server([ + Response( + body=b'{"choices":[{"message":{"content":"Answer"}}],' + b'"usage":{"prompt_tokens":0,"completion_tokens":0}}' + ) + ]) + model = OAICompatLargeLanguageModel([]) + with model.timing_context(): + model._generate( + "model", + { + "mode": "chat", + "endpoint_url": server.url, + "function_calling_type": "tool_call", + }, + [UserPromptMessage(content="Search")], + {"tool_choice": "required"}, + tools=[ + PromptMessageTool( + name="search", description="Search", parameters={"type": "object"} + ) + ], + stream=False, + ) + + payload = json.loads(server.requests[0].body) + assert payload["tool_choice"] == "required" + assert payload["tools"][0]["function"]["name"] == "search" + + @pytest.mark.parametrize( ("mode", "stream"), [("chat", False), ("completion", False), ("chat", True)] )