From e07fc30b6b504dc9b5b68fbaa09cb15e8a38433c Mon Sep 17 00:00:00 2001 From: vislee Date: Thu, 17 Apr 2025 17:31:57 +0800 Subject: [PATCH 1/3] bugfix: The events array after splitting may contain empty strings, which can cause subsequent messages to be skipped. --- apisix/plugins/ai-drivers/openai-base.lua | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/apisix/plugins/ai-drivers/openai-base.lua b/apisix/plugins/ai-drivers/openai-base.lua index a4f061fe4c13..24243557bd55 100644 --- a/apisix/plugins/ai-drivers/openai-base.lua +++ b/apisix/plugins/ai-drivers/openai-base.lua @@ -167,24 +167,24 @@ function _M.read_response(ctx, res) for _, event in ipairs(events) do if not core.string.find(event, "data:") or core.string.find(event, "[DONE]") then - goto CONTINUE + goto CONTINUEFOR end local parts, err = ngx_re.split(event, ":", nil, nil, 2) if err then core.log.warn("failed to split data event [", event, "] to parts: ", err) - goto CONTINUE + goto CONTINUEFOR end if #parts ~= 2 then core.log.warn("malformed data event: ", event) - goto CONTINUE + goto CONTINUEFOR end local data, err = core.json.decode(parts[2]) if err then core.log.warn("failed to decode data event [", parts[2], "] to json: ", err) - goto CONTINUE + goto CONTINUEFOR end -- usage field is null for non-last events, null is parsed as userdata type @@ -197,6 +197,7 @@ function _M.read_response(ctx, res) total_tokens = data.usage.total_tokens or 0, } end + ::CONTINUEFOR:: end ::CONTINUE:: From c130ecf4823854931f9f1905f3890bbf85e1daf6 Mon Sep 17 00:00:00 2001 From: vislee Date: Mon, 21 Apr 2025 11:21:12 +0800 Subject: [PATCH 2/3] case: add test case --- t/APISIX.pm | 2 + t/assets/ai-proxy-stream-response.json | 7 ++ t/plugin/ai-proxy3.t | 150 +++++++++++++++++++++++++ 3 files changed, 159 insertions(+) create mode 100644 t/assets/ai-proxy-stream-response.json create mode 100644 t/plugin/ai-proxy3.t diff --git a/t/APISIX.pm b/t/APISIX.pm index f4b9b8055d4a..50d35558c631 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -731,6 +731,7 @@ _EOC_ $ipv6_listen_conf = "listen \[::1\]:1984;" } + my $log_config = $block->log_config // ''; my $config = $block->config // ''; $config .= <<_EOC_; $ipv6_listen_conf @@ -844,6 +845,7 @@ _EOC_ } log_by_lua_block { + $log_config apisix.http_log_phase() } } diff --git a/t/assets/ai-proxy-stream-response.json b/t/assets/ai-proxy-stream-response.json new file mode 100644 index 000000000000..3c5c50f667cc --- /dev/null +++ b/t/assets/ai-proxy-stream-response.json @@ -0,0 +1,7 @@ +[ +"data: {\"id\":\"cmpl-130\",\"object\":\"chat.completion.chunk\",\"created\":1698999575,\"model\":\"gpt-4o-2024-05-13\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\" 1\"},\"finish_reason\":null}]}\n", +"data: {\"id\":\"cmpl-130\",\"object\":\"chat.completion.chunk\",\"created\":1698999575,\"model\":\"gpt-4o-2024-05-13\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" + 1\"},\"finish_reason\":null}]}\n", +"data: {\"id\":\"cmpl-130\",\"object\":\"chat.completion.chunk\",\"created\":1698999575,\"model\":\"gpt-4o-2024-05-13\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" = 2\"},\"finish_reason\":null}]}\n", +"data: {\"id\":\"cmpl-130\",\"object\":\"chat.completion.chunk\",\"created\":1698999575,\"model\":\"gpt-4o-2024-05-13\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\",\"usage\":{\"prompt_tokens\":19,\"completion_tokens\":13,\"total_tokens\":32}}]}\n", +"data: [DONE]\n" +] \ No newline at end of file diff --git a/t/plugin/ai-proxy3.t b/t/plugin/ai-proxy3.t new file mode 100644 index 000000000000..b5ee02e7d836 --- /dev/null +++ b/t/plugin/ai-proxy3.t @@ -0,0 +1,150 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +use t::APISIX 'no_plan'; + +log_level("info"); +repeat_each(1); +no_long_string(); +no_root_location(); + + +my $resp_file = 't/assets/ai-proxy-stream-response.json'; +open(my $fh, '<', $resp_file) or die "Could not open file '$resp_file' $!"; +my $resp = do { local $/; <$fh> }; +close($fh); + +print "Hello, World!\n"; +print $resp; + + +add_block_preprocessor(sub { + my ($block) = @_; + + if (!defined $block->request) { + $block->set_value("request", "GET /t"); + } + + my $http_config = $block->http_config // <<_EOC_; + server { + server_name openai; + listen 6724; + + default_type 'application/json'; + + location /v1/chat/completions { + content_by_lua_block { + local json = require("cjson.safe") + + if ngx.req.get_method() ~= "POST" then + ngx.status = 400 + ngx.say("Unsupported request method: ", ngx.req.get_method()) + end + ngx.req.read_body() + local body, err = ngx.req.get_body_data() + body, err = json.decode(body) + + local query_auth = ngx.req.get_uri_args()["api_key"] + + if query_auth ~= "apikey" then + ngx.status = 401 + ngx.say("Unauthorized") + return + end + + local data, err = json.decode([[$resp]]) + if not err then + ngx.status = 500 + ngx.say(err) + return + end + + ngx.status = 200 + for _, val in ipairs(data) do + ngx.say(val) + ngx.flush(true) + end + return + } + } + } +_EOC_ + + $block->set_value("http_config", $http_config); +}); + +run_tests(); + +__DATA__ + + +=== TEST 1: send request +--- log_config +local api_ctx = ngx.ctx.api_ctx +if api_ctx then + ngx.log(ngx.INFO, "prompt_tokens: ", api_ctx.ai_token_usage and api_ctx.ai_token_usage.prompt_tokens or 0) + ngx.log(ngx.INFO, "completion_tokens: ", api_ctx.ai_token_usage and api_ctx.ai_token_usage.completion_tokens or 0) + ngx.log(ngx.INFO, "total_tokens: ", api_ctx.ai_token_usage and api_ctx.ai_token_usage.total_tokens or 0) +end + +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { + "header": { + "Authorization": "Bearer token" + } + }, + "options": { + "model": "gpt-35-turbo-instruct", + "max_tokens": 1024, + "temperature": 1.0 + }, + "override": { + "endpoint": "http://localhost:6724" + }, + "ssl_verify": false + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } + +--- request +POST /anything +{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } +--- error_code: 200 +--- no_error_log +[error] +--- error_log +prompt_tokens: 19 +completion_tokens: 13 +total_tokens: 32 From 60ff4f674acbeb1311ea82813426405780202713 Mon Sep 17 00:00:00 2001 From: vislee Date: Sat, 24 May 2025 12:31:02 +0800 Subject: [PATCH 3/3] fix failed lint CI --- t/assets/ai-proxy-stream-response.json | 2 +- t/plugin/ai-proxy3.t | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/t/assets/ai-proxy-stream-response.json b/t/assets/ai-proxy-stream-response.json index 3c5c50f667cc..a8408e22f8a9 100644 --- a/t/assets/ai-proxy-stream-response.json +++ b/t/assets/ai-proxy-stream-response.json @@ -4,4 +4,4 @@ "data: {\"id\":\"cmpl-130\",\"object\":\"chat.completion.chunk\",\"created\":1698999575,\"model\":\"gpt-4o-2024-05-13\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" = 2\"},\"finish_reason\":null}]}\n", "data: {\"id\":\"cmpl-130\",\"object\":\"chat.completion.chunk\",\"created\":1698999575,\"model\":\"gpt-4o-2024-05-13\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\",\"usage\":{\"prompt_tokens\":19,\"completion_tokens\":13,\"total_tokens\":32}}]}\n", "data: [DONE]\n" -] \ No newline at end of file +] diff --git a/t/plugin/ai-proxy3.t b/t/plugin/ai-proxy3.t index b5ee02e7d836..5aa712e36a02 100644 --- a/t/plugin/ai-proxy3.t +++ b/t/plugin/ai-proxy3.t @@ -91,7 +91,6 @@ run_tests(); __DATA__ - === TEST 1: send request --- log_config local api_ctx = ngx.ctx.api_ctx