fix: pass context to provider for stream data parsing (#1475)

* fix: pass context to provider for stream data parsing * fix: luatype --------- Co-authored-by: yetone <yetoneful@gmail.com>
yetone · Mar 3, 2025 · 6bd966e · 6bd966e
1 parent c6d5527
commit 6bd966e
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 6 deletions.
diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua
@@ -154,6 +154,8 @@ end
 function M._stream(opts)
   local provider = opts.provider or Providers[Config.provider]
 
+  ---@cast provider AvanteProviderFunctor
+
   local prompt_opts = M.generate_prompts(opts)
 
   ---@type string
@@ -285,10 +287,10 @@ function M._stream(opts)
               { once = true }
             )
           end
-          provider.parse_stream_data(data, handler_opts)
+          provider.parse_stream_data(resp_ctx, data, handler_opts)
         else
           if provider.parse_stream_data ~= nil then
-            provider.parse_stream_data(data, handler_opts)
+            provider.parse_stream_data(resp_ctx, data, handler_opts)
           else
             parse_stream_data(data)
           end

diff --git a/lua/avante/providers/bedrock.lua b/lua/avante/providers/bedrock.lua
@@ -28,7 +28,7 @@ function M.build_bedrock_payload(prompt_opts, body_opts)
   return model_handler.build_bedrock_payload(prompt_opts, body_opts)
 end
 
-function M.parse_stream_data(data, opts)
+function M.parse_stream_data(ctx, data, opts)
   -- @NOTE: Decode and process Bedrock response
   -- Each response contains a Base64-encoded `bytes` field, which is decoded into JSON.
   -- The `type` field in the decoded JSON determines how the response is handled.
@@ -37,7 +37,7 @@ function M.parse_stream_data(data, opts)
     local jsn = vim.json.decode(bedrock_data_match)
     local data_stream = vim.base64.decode(jsn.bytes)
     local json = vim.json.decode(data_stream)
-    M.parse_response({}, data_stream, json.type, opts)
+    M.parse_response(ctx, data_stream, json.type, opts)
   end
 end
 

diff --git a/lua/avante/providers/cohere.lua b/lua/avante/providers/cohere.lua
@@ -57,7 +57,7 @@ function M.parse_messages(opts)
   return { messages = messages }
 end
 
-function M.parse_stream_data(data, opts)
+function M.parse_stream_data(ctx, data, opts)
   ---@type CohereChatResponse
   local json = vim.json.decode(data)
   if json.type ~= nil then

diff --git a/lua/avante/types.lua b/lua/avante/types.lua
@@ -239,7 +239,7 @@ vim.g.avante_login = vim.g.avante_login
 ---@field tool_use_list? AvanteLLMToolUse[]
 ---@field retry_after? integer
 ---
----@alias AvanteStreamParser fun(line: string, handler_opts: AvanteHandlerOptions): nil
+---@alias AvanteStreamParser fun(ctx: any, line: string, handler_opts: AvanteHandlerOptions): nil
 ---@alias AvanteLLMStartCallback fun(opts: AvanteLLMStartCallbackOptions): nil
 ---@alias AvanteLLMChunkCallback fun(chunk: string): any
 ---@alias AvanteLLMStopCallback fun(opts: AvanteLLMStopCallbackOptions): nil