Skip to content

Commit 3374d07

Browse files
frankroederfrankroeder
committed
Output Claude reasoning to floating window
1 parent b3c8dbd commit 3374d07

File tree

4 files changed

+174
-217
lines changed

4 files changed

+174
-217
lines changed

‎lua/parrot/provider/anthropic.lua

Lines changed: 52 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@ function Anthropic:new(endpoint, api_key)
4141
endpoint = endpoint,
4242
api_key = api_key,
4343
name = "anthropic",
44+
_thinking_buf = nil,
45+
_thinking_win = nil,
46+
_thinking_output = "",
4447
}, self)
4548
end
4649

@@ -60,7 +63,8 @@ function Anthropic:preprocess_payload(payload)
6063
payload.system = payload.messages[1].content
6164
table.remove(payload.messages, 1)
6265
end
63-
return utils.filter_payload_parameters(AVAILABLE_API_PARAMETERS, payload)
66+
local params = utils.filter_payload_parameters(AVAILABLE_API_PARAMETERS, payload)
67+
return params
6468
end
6569

6670
-- Returns the curl parameters for the API request
@@ -97,18 +101,59 @@ function Anthropic:verify()
97101
end
98102
end
99103

100-
-- Processes the stdout from the API response
104+
-- Notification system: displays thinking tokens in a floating window in the top right.
105+
-- The buffer is created with text wrapping enabled, and tokens are accumulated into one coherent string.
106+
function Anthropic:notify_thinking(thinking)
107+
vim.schedule(function()
108+
if not self._thinking_buf or not vim.api.nvim_buf_is_valid(self._thinking_buf) then
109+
self._thinking_buf = vim.api.nvim_create_buf(false, true) -- unlisted scratch buffer
110+
local width = math.floor(vim.o.columns * 0.3)
111+
local height = math.floor(vim.o.lines * 0.3)
112+
local row = 0
113+
local col = vim.o.columns - width
114+
self._thinking_win = vim.api.nvim_open_win(self._thinking_buf, true, {
115+
relative = "editor",
116+
width = width,
117+
height = height,
118+
row = row,
119+
col = col,
120+
style = "minimal",
121+
border = "rounded",
122+
})
123+
vim.api.nvim_buf_set_option(self._thinking_buf, "buftype", "nofile")
124+
vim.api.nvim_win_set_option(self._thinking_win, "wrap", true)
125+
self._thinking_output = ""
126+
end
127+
128+
-- Accumulate tokens into one coherent string.
129+
self._thinking_output = self._thinking_output .. thinking
130+
vim.api.nvim_buf_set_lines(self._thinking_buf, 0, -1, false, { self._thinking_output })
131+
end)
132+
end
133+
134+
-- Processes the stdout from the API response.
135+
-- For "text_delta" responses, returns the text.
136+
-- For "thinking_delta" responses, streams tokens to the floating window.
101137
---@param response string
102138
---@return string|nil
103139
function Anthropic:process_stdout(response)
104-
if response:match("content_block_delta") and response:match("text_delta") then
105-
local success, decoded_line = pcall(vim.json.decode, response)
106-
if success and decoded_line.delta and decoded_line.delta.type == "text_delta" and decoded_line.delta.text then
140+
local success, decoded_line = pcall(vim.json.decode, response)
141+
if not success then
142+
logger.debug("Could not decode response: " .. response)
143+
return nil
144+
end
145+
146+
if decoded_line.delta then
147+
if decoded_line.delta.type == "text_delta" and decoded_line.delta.text then
107148
return decoded_line.delta.text
108-
else
109-
logger.debug("Could not process response: " .. response)
149+
elseif decoded_line.delta.type == "thinking_delta" and decoded_line.delta.thinking then
150+
self:notify_thinking(decoded_line.delta.thinking)
151+
return nil
110152
end
111153
end
154+
155+
logger.debug("Could not process response: " .. response)
156+
return nil
112157
end
113158

114159
-- Processes the onexit event from the API response

‎lua/parrot/utils.lua

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -367,11 +367,32 @@ M.filter_payload_parameters = function(valid_parameters, payload)
367367
local new_payload = {}
368368
for key, value in pairs(valid_parameters) do
369369
if type(value) == "table" then
370+
-- Initialize table only if needed
370371
if new_payload[key] == nil then
371372
new_payload[key] = {}
372373
end
373-
for tkey, _ in pairs(value) do
374-
new_payload[key][tkey] = payload[tkey]
374+
375+
local found_values = false
376+
if payload[key] then
377+
for tkey, _ in pairs(value) do
378+
if payload[key][tkey] then
379+
new_payload[key][tkey] = payload[key][tkey]
380+
found_values = true
381+
end
382+
end
383+
else
384+
-- Look for the nested keys at top level
385+
for tkey, _ in pairs(value) do
386+
if payload[tkey] then
387+
new_payload[key][tkey] = payload[tkey]
388+
found_values = true
389+
end
390+
end
391+
end
392+
393+
-- Remove empty tables
394+
if not found_values then
395+
new_payload[key] = nil
375396
end
376397
else
377398
new_payload[key] = payload[key]

‎tests/parrot/provider/anthropic_spec.lua

Lines changed: 43 additions & 208 deletions
Original file line numberDiff line numberDiff line change
@@ -15,37 +15,36 @@ describe("Anthropic", function()
1515
assert.are.same(anthropic.name, "anthropic")
1616
end)
1717

18-
-- TODO: preprocess_payload output is nil --
19-
-- describe("preprocess_payload", function()
20-
-- it("should handle payload with system message correctly", function()
21-
-- local input = {
22-
-- max_tokens = 4096,
23-
-- messages = {
24-
-- {
25-
-- content = "You are a versatile AI assistant with capabilities\nextending to general knowledge and coding support. When engaging\nwith users, please adhere to the following guidelines to ensure\nthe highest quality of interaction:\n\n- Admit when unsure by saying 'I don't know.'\n- Ask for clarification when needed.\n- Use first principles thinking to analyze queries.\n- Start with the big picture, then focus on details.\n- Apply the Socratic method to enhance understanding.\n- Include all necessary code in your responses.\n- Stay calm and confident with each task.\n",
26-
-- role = "system",
27-
-- },
28-
-- { content = "Who are you?", role = "user" },
29-
-- },
30-
-- model = "claude-3-haiku-20240307",
31-
-- stream = true,
32-
-- }
33-
--
34-
-- local expected = {
35-
-- max_tokens = 4096,
36-
-- messages = {
37-
-- { content = "Who are you?", role = "user" },
38-
-- },
39-
-- model = "claude-3-haiku-20240307",
40-
-- stream = true,
41-
-- system = "You are a versatile AI assistant with capabilities\nextending to general knowledge and coding support. When engaging\nwith users, please adhere to the following guidelines to ensure\nthe highest quality of interaction:\n\n- Admit when unsure by saying 'I don't know.'\n- Ask for clarification when needed.\n- Use first principles thinking to analyze queries.\n- Start with the big picture, then focus on details.\n- Apply the Socratic method to enhance understanding.\n- Include all necessary code in your responses.\n- Stay calm and confident with each task.",
42-
-- }
43-
--
44-
-- local result = anthropic:preprocess_payload(input)
45-
--
46-
-- assert.are.same(expected, result)
47-
-- end)
48-
-- end)
18+
describe("preprocess_payload", function()
19+
it("should handle payload with system message correctly", function()
20+
local input = {
21+
max_tokens = 4096,
22+
messages = {
23+
{
24+
content = "You are a versatile AI assistant with capabilities\nextending to general knowledge and coding support. When engaging\nwith users, please adhere to the following guidelines to ensure\nthe highest quality of interaction:\n\n- Admit when unsure by saying 'I don't know.'\n- Ask for clarification when needed.\n- Use first principles thinking to analyze queries.\n- Start with the big picture, then focus on details.\n- Apply the Socratic method to enhance understanding.\n- Include all necessary code in your responses.\n- Stay calm and confident with each task.\n",
25+
role = "system",
26+
},
27+
{ content = "Who are you?", role = "user" },
28+
},
29+
model = "claude-3-haiku-20240307",
30+
stream = true,
31+
}
32+
33+
local expected = {
34+
max_tokens = 4096,
35+
messages = {
36+
{ content = "Who are you?", role = "user" },
37+
},
38+
model = "claude-3-haiku-20240307",
39+
stream = true,
40+
system = "You are a versatile AI assistant with capabilities\nextending to general knowledge and coding support. When engaging\nwith users, please adhere to the following guidelines to ensure\nthe highest quality of interaction:\n\n- Admit when unsure by saying 'I don't know.'\n- Ask for clarification when needed.\n- Use first principles thinking to analyze queries.\n- Start with the big picture, then focus on details.\n- Apply the Socratic method to enhance understanding.\n- Include all necessary code in your responses.\n- Stay calm and confident with each task.",
41+
}
42+
43+
local result = anthropic:preprocess_payload(input)
44+
45+
assert.are.same(expected, result)
46+
end)
47+
end)
4948

5049
describe("verify", function()
5150
it("should return true for a valid API key", function()
@@ -74,7 +73,7 @@ describe("Anthropic", function()
7473

7574
describe("process_stdout", function()
7675
it("should extract text from content_block_delta with text_delta", function()
77-
local input = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello, world!"}}'
76+
local input = '{"delta":{"type":"text_delta","text":"Hello, world!"}}'
7877

7978
local result = anthropic:process_stdout(input)
8079

@@ -83,7 +82,7 @@ describe("Anthropic", function()
8382

8483
it("should return nil for non-text_delta messages", function()
8584
local input =
86-
'{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":8}}'
85+
'{"delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":8}}'
8786

8887
local result = anthropic:process_stdout(input)
8988

@@ -113,180 +112,16 @@ describe("Anthropic", function()
113112

114113
assert.is_nil(result)
115114
end)
115+
116+
it("should accumulate thinking tokens in the floating window for reasoning", function()
117+
-- Simulate two consecutive streaming events with reasoning tokens.
118+
local stream1 = '{"delta": {"type": "thinking_delta", "thinking": "To calculate 27"}}'
119+
local stream2 = '{"delta": {"type": "thinking_delta", "thinking": " * 453, I\'ll multiply these"}}'
120+
anthropic:process_stdout(stream1)
121+
anthropic:process_stdout(stream2)
122+
-- Wait a short time to allow the scheduled callback to run.
123+
vim.wait(50)
124+
assert.equals("To calculate 27 * 453, I'll multiply these", anthropic._thinking_output)
125+
end)
116126
end)
117127
end)
118-
119-
-- Thinking process
120-
-- curl https://api.anthropic.com/v1/messages \ 22:53:45
121-
-- --header "x-api-key: $ANTHROPIC_API_KEY" \
122-
-- --header "anthropic-version: 2023-06-01" \
123-
-- --header "content-type: application/json" \
124-
-- --data \
125-
-- '{
126-
-- "model": "claude-3-7-sonnet-20250219",
127-
-- "max_tokens": 20000,
128-
-- "stream": true,
129-
-- "thinking": {
130-
-- "type": "enabled",
131-
-- "budget_tokens": 16000
132-
-- },
133-
-- "messages": [
134-
-- {
135-
-- "role": "user",
136-
-- "content": "What is 27 * 453?"
137-
-- }
138-
-- ]
139-
-- }'
140-
141-
-- event: message_start
142-
-- data: {"type":"message_start","message":{"id":"msg_017bnwLbFy7uzWgDGkS4FpK5","type":"message","role":"assistant","model":"claude-3-7-sonnet-20250219","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":44,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":5}} }
143-
144-
-- event: content_block_start
145-
-- data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""} }
146-
147-
-- event: ping
148-
-- data: {"type": "ping"}
149-
150-
-- event: content_block_delta
151-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"To calculate 27"} }
152-
153-
-- event: content_block_delta
154-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" * 453, I'll multiply these"} }
155-
156-
-- event: content_block_delta
157-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" numbers step by step."} }
158-
159-
-- event: content_block_delta
160-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"\n\nFirst, let me break this down:\n27 "} }
161-
162-
-- event: content_block_delta
163-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"* 453 = (20"} }
164-
165-
-- event: content_block_delta
166-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" + 7) * "} }
167-
168-
-- event: content_block_delta
169-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"453\n = 20 * 453"} }
170-
171-
-- event: content_block_delta
172-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" + 7 * 453\n "} }
173-
174-
-- event: content_block_delta
175-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"= 9060 + 7"} }
176-
177-
-- event: content_block_delta
178-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" * 453\n\nNow I need to calculate"} }
179-
180-
-- event: content_block_delta
181-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" 7 * 453."}}
182-
183-
-- event: content_block_delta
184-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"\n7 * 453 = 7"} }
185-
186-
-- event: content_block_delta
187-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" * 400 "}}
188-
189-
-- event: content_block_delta
190-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"+ 7 *"} }
191-
192-
-- event: content_block_delta
193-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" 50 + 7 "} }
194-
195-
-- event: content_block_delta
196-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"* 3\n "} }
197-
198-
-- event: content_block_delta
199-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"= 2800 + 350 +"}}
200-
201-
-- event: content_block_delta
202-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" 21\n ="} }
203-
204-
-- event: content_block_delta
205-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" 3171\n\nSo,"} }
206-
207-
-- event: content_block_delta
208-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" 27 * 453 = 9"} }
209-
210-
-- event: content_block_delta
211-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"060 + 3"} }
212-
213-
-- event: content_block_delta
214-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"171 = 12231.\n\nLet me"} }
215-
216-
-- event: content_block_delta
217-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" double-check using the standar"} }
218-
219-
-- event: content_block_delta
220-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"d multiplication algorithm:\n\n 453"} }
221-
222-
-- event: content_block_delta
223-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"\n × 27\n "} }
224-
225-
-- event: content_block_delta
226-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"-----\n 3171"} }
227-
228-
-- event: content_block_delta
229-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" (7 × 453)"} }
230-
231-
-- event: content_block_delta
232-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"\n 9060 (20 "} }
233-
234-
-- event: content_block_delta
235-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"× 453)\n -----\n "} }
236-
237-
-- event: content_block_delta
238-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"12231\n\nSo 27 * 453"} }
239-
240-
-- event: content_block_delta
241-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":" = 12231."} }
242-
243-
-- event: content_block_delta
244-
-- data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"ErUBCkYIARgCIkBWljGprkAg3jqP0QHtn6lZ+uduvsVrxQsZhpM26RFq+lmLbwPbv6Ow9qvUmnU5T7HlLD47T0vL6RcgyYOh77qmEgz1esz0owfDCYDrga0aDGviVWnJRQv6cOR4MSIwYU+28tQAdNJ3m74ryq86qEdZt8d/tXfdV67t9DNc5FXaP0T4ZelOAB6XbVcj3IONKh1UFyd1cHJrPfpbv1wQcW3lPCYbjkPUJUZP/XQH5g=="} }
245-
246-
-- event: content_block_stop
247-
-- data: {"type":"content_block_stop","index":0 }
248-
249-
-- event: content_block_start
250-
-- data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""} }
251-
252-
-- event: content_block_delta
253-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"To multiply 27 × "} }
254-
255-
-- event: content_block_delta
256-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"453, I'll work"} }
257-
258-
-- event: content_block_delta
259-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":" through this step by step"}}
260-
261-
-- event: content_block_delta
262-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":":\n\n 453\n × 27"} }
263-
264-
-- event: content_block_delta
265-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"\n -----\n "} }
266-
267-
-- event: content_block_delta
268-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"3171 (7 × 453)"} }
269-
270-
-- event: content_block_delta
271-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"\n 9060 (20 × "} }
272-
273-
-- event: content_block_delta
274-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"453)\n -----\n 12231"} }
275-
276-
-- event: content_block_delta
277-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"\n\nTherefore, 27 ×"} }
278-
279-
-- event: content_block_delta
280-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":" 453 = 12,"} }
281-
282-
-- event: content_block_delta
283-
-- data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"231"} }
284-
285-
-- event: content_block_stop
286-
-- data: {"type":"content_block_stop","index":1 }
287-
288-
-- event: message_delta
289-
-- data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":324} }
290-
291-
-- event: message_stop
292-
-- data: {"type":"message_stop" }

0 commit comments

Comments
 (0)