From edce9e39c5dcb7381fa28adfe0ba9edacd9d8035 Mon Sep 17 00:00:00 2001 From: Jack Tysoe Date: Thu, 11 Jan 2024 12:02:42 +0000 Subject: [PATCH] fix(ai-proxy): llama2 analytics --- kong/llm/drivers/llama2.lua | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/kong/llm/drivers/llama2.lua b/kong/llm/drivers/llama2.lua index 85a720e6fe12..7bd8d653e906 100644 --- a/kong/llm/drivers/llama2.lua +++ b/kong/llm/drivers/llama2.lua @@ -99,7 +99,7 @@ local transformers_from = { ["llm/v1/chat"] = function(response_string, route_type) local response_table, err = cjson.decode(response_string) if err then - return nil, "failed to decode cohere response" + return nil, "failed to decode llama2 response" end if (not response_table) or (not response_table.data) or (#response_table.data > 1) then @@ -145,13 +145,16 @@ local transformers_from = { end + -- stash analytics for later + if response_table.usage then response_object.usage = response_table.usage end + return cjson.encode(response_object) end, ["llm/v1/completions"] = function(response_string, route_type) local response_table, err = cjson.decode(response_string) if err then - return nil, "failed to decode cohere response" + return nil, "failed to decode llama2 response" end if (not response_table) or (not response_table.data) or (#response_table.data > 1) then @@ -197,6 +200,9 @@ local transformers_from = { end + -- stash analytics for later + if response_table.usage then response_object.usage = response_table.usage end + return cjson.encode(response_object) end, }