Skip to content
Draft
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions apisix/plugins/ai-drivers/openai-base.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -197,6 +197,7 @@ function _M.read_response(ctx, res)
total_tokens = data.usage.total_tokens or 0,
}
end
::CONTINUEFOR::
end

::CONTINUE::
Expand Down
2 changes: 2 additions & 0 deletions t/APISIX.pm
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -844,6 +845,7 @@ _EOC_
}

log_by_lua_block {
$log_config
apisix.http_log_phase()
}
}
Expand Down
7 changes: 7 additions & 0 deletions t/assets/ai-proxy-stream-response.json
Original file line number Diff line number Diff line change
@@ -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"
]
150 changes: 150 additions & 0 deletions t/plugin/ai-proxy3.t
Original file line number Diff line number Diff line change
@@ -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
Loading