Skip to content

Commit 906483a

Browse files
JavaHelloJavaHello
committed
拆分gpt脚本
1 parent c4707c8 commit 906483a

File tree

6 files changed

+471
-508
lines changed

6 files changed

+471
-508
lines changed

‎lua/kide/gpt/chat.lua‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
local M = {}
2+
local state = {
3+
chatwin = nil,
4+
chatbuf = nil,
5+
codebuf = nil,
6+
chatclosed = true,
7+
cursormoved = false,
8+
chatruning = false,
9+
winleave = false,
10+
}
11+
12+
local chat_request_json = {
13+
messages = {
14+
{
15+
content = "",
16+
role = "system",
17+
},
18+
},
19+
model = "deepseek-chat",
20+
frequency_penalty = 0,
21+
max_tokens = 4096,
22+
presence_penalty = 0,
23+
response_format = {
24+
type = "text",
25+
},
26+
stop = nil,
27+
stream = true,
28+
stream_options = nil,
29+
temperature = 1.3,
30+
top_p = 1,
31+
tools = nil,
32+
tool_choice = "none",
33+
logprobs = false,
34+
top_logprobs = nil,
35+
}
36+
37+
M.chat_config = {
38+
user_title = " :",
39+
system_title = " :",
40+
system_prompt = "You are a general AI assistant.\n\n"
41+
.. "The user provided the additional info about how they would like you to respond:\n\n"
42+
.. "- If you're unsure don't guess and say you don't know instead.\n"
43+
.. "- Ask question if you need clarification to provide better answer.\n"
44+
.. "- Think deeply and carefully from first principles step by step.\n"
45+
.. "- Zoom out first to see the big picture and then zoom in to details.\n"
46+
.. "- Use Socratic method to improve your thinking and coding skills.\n"
47+
.. "- Don't elide any code from your output if the answer requires coding.\n"
48+
.. "- Take a deep breath; You've got this!\n",
49+
}
50+
local close_gpt_win = function()
51+
if state.chatwin then
52+
pcall(vim.api.nvim_win_close, state.chatwin, true)
53+
chat_request_json.messages = {
54+
{
55+
content = "",
56+
role = "system",
57+
},
58+
}
59+
state.chatwin = nil
60+
state.chatbuf = nil
61+
state.codebuf = nil
62+
state.chatclosed = true
63+
state.cursormoved = false
64+
state.chatruning = false
65+
state.winleave = false
66+
end
67+
end
68+
69+
local function create_gpt_win()
70+
state.codebuf = vim.api.nvim_get_current_buf()
71+
vim.cmd("belowright new")
72+
state.chatwin = vim.api.nvim_get_current_win()
73+
state.chatbuf = vim.api.nvim_get_current_buf()
74+
vim.bo[state.chatbuf].buftype = "nofile"
75+
vim.bo[state.chatbuf].bufhidden = "wipe"
76+
vim.bo[state.chatbuf].buflisted = false
77+
vim.bo[state.chatbuf].swapfile = false
78+
vim.bo[state.chatbuf].filetype = "markdown"
79+
vim.api.nvim_put({ M.chat_config.user_title, "" }, "c", true, true)
80+
state.chatclosed = false
81+
82+
vim.keymap.set("n", "q", function()
83+
state.chatclosed = true
84+
close_gpt_win()
85+
end, { noremap = true, silent = true, buffer = state.chatbuf })
86+
vim.keymap.set("n", "<A-k>", function()
87+
M.gpt_chat()
88+
end, { noremap = true, silent = true, buffer = state.chatbuf })
89+
vim.keymap.set("i", "<A-k>", function()
90+
vim.cmd("stopinsert")
91+
M.gpt_chat()
92+
end, { noremap = true, silent = true, buffer = state.chatbuf })
93+
94+
vim.api.nvim_buf_create_user_command(state.chatbuf, "GptSend", function()
95+
M.gpt_chat()
96+
end, { desc = "Gpt Send" })
97+
98+
vim.api.nvim_create_autocmd("BufWipeout", {
99+
buffer = state.chatbuf,
100+
callback = close_gpt_win,
101+
})
102+
103+
vim.api.nvim_create_autocmd("WinClosed", {
104+
buffer = state.chatbuf,
105+
callback = close_gpt_win,
106+
})
107+
vim.api.nvim_create_autocmd("WinLeave", {
108+
buffer = state.chatbuf,
109+
callback = function()
110+
state.winleave = true
111+
end,
112+
})
113+
114+
vim.api.nvim_create_autocmd("CursorMoved", {
115+
buffer = state.chatbuf,
116+
callback = function()
117+
state.cursormoved = true
118+
end,
119+
})
120+
end
121+
122+
local function code_question(selection)
123+
if not selection then
124+
return
125+
end
126+
local qs
127+
---@diagnostic disable-next-line: param-type-mismatch
128+
if vim.api.nvim_buf_is_valid(state.codebuf) then
129+
local filetype = vim.bo[state.codebuf].filetype or "text"
130+
local filename = vim.fn.fnamemodify(vim.fn.bufname(state.codebuf), ":.")
131+
qs = {
132+
"请解释`" .. filename .. "`文件中的这段代码",
133+
"```" .. filetype,
134+
}
135+
else
136+
qs = {
137+
"请解释这段代码",
138+
"```",
139+
}
140+
end
141+
vim.list_extend(qs, selection)
142+
table.insert(qs, "```")
143+
vim.api.nvim_put(qs, "c", true, true)
144+
end
145+
146+
M.toggle_gpt = function(selection)
147+
if state.chatwin then
148+
close_gpt_win()
149+
else
150+
create_gpt_win()
151+
code_question(selection)
152+
end
153+
end
154+
155+
M.gpt_chat = function()
156+
if state.chatwin == nil then
157+
create_gpt_win()
158+
end
159+
if state.chatruning then
160+
vim.api.nvim_put({ "", "", M.chat_config.user_title, "" }, "c", true, true)
161+
state.chatruning = false
162+
return
163+
end
164+
state.chatruning = true
165+
---@diagnostic disable-next-line: param-type-mismatch
166+
local list = vim.api.nvim_buf_get_lines(state.chatbuf, 0, -1, false)
167+
local json = chat_request_json
168+
json.messages[1].content = M.chat_config.system_prompt
169+
-- 1 user, 2 assistant
170+
local flag = 0
171+
local chat_msg = ""
172+
local chat_count = 1
173+
for _, v in ipairs(list) do
174+
if vim.startswith(v, M.chat_config.system_title) then
175+
flag = 2
176+
chat_msg = ""
177+
chat_count = chat_count + 1
178+
elseif vim.startswith(v, M.chat_config.user_title) then
179+
chat_msg = ""
180+
flag = 1
181+
chat_count = chat_count + 1
182+
else
183+
chat_msg = chat_msg .. "\n" .. v
184+
json.messages[chat_count] = {
185+
content = chat_msg,
186+
role = flag == 1 and "user" or "assistant",
187+
}
188+
end
189+
end
190+
-- 跳转到最后一行
191+
vim.cmd("normal! G$")
192+
vim.api.nvim_put({ "", M.chat_config.system_title, "" }, "l", true, true)
193+
194+
local callback = function(opt)
195+
local data = opt.data
196+
local done = opt.done
197+
if state.chatclosed or state.chatruning == false then
198+
vim.fn.jobstop(opt.job)
199+
return
200+
end
201+
if opt.err == 1 then
202+
vim.notify("AI respond Error: " .. opt.data, vim.log.levels.WARN)
203+
return
204+
end
205+
if state.winleave then
206+
-- 防止回答问题时光标已经移动走了
207+
vim.api.nvim_set_current_win(state.chatwin)
208+
state.winleave = false
209+
end
210+
if state.cursormoved then
211+
-- 防止光标移动打乱回答顺序, 总是移动到最后一行
212+
vim.cmd("normal! G$")
213+
state.cursormoved = false
214+
end
215+
if done then
216+
vim.api.nvim_put({ "", "", M.chat_config.user_title, "" }, "c", true, true)
217+
state.chatruning = false
218+
return
219+
end
220+
if state.chatbuf and vim.api.nvim_buf_is_valid(state.chatbuf) then
221+
if data:match("\n") then
222+
local ln = vim.split(data, "\n")
223+
vim.api.nvim_put(ln, "c", true, true)
224+
else
225+
vim.api.nvim_put({ data }, "c", true, true)
226+
end
227+
end
228+
end
229+
230+
require("kide.gpt.sse").request(json, callback)
231+
end
232+
233+
return M

‎lua/kide/gpt/sse.lua‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
local M = {}
2+
3+
M.config = {
4+
API_URL = vim.env["DEEPSEEK_API_ENDPOINT"],
5+
API_KEY = vim.env["DEEPSEEK_API_KEY"],
6+
show_usage = false,
7+
}
8+
9+
local function token()
10+
return "Bearer " .. M.config.API_KEY
11+
end
12+
13+
local function callback_data(job, resp_json, callback)
14+
callback({
15+
data = resp_json.choices[1].delta.content,
16+
job = job,
17+
})
18+
if resp_json.usage ~= nil and M.config.show_usage then
19+
callback({
20+
data = "\n",
21+
job = job,
22+
})
23+
callback({
24+
data = "API[token usage]: " .. vim.inspect(resp_json.usage.prompt_cache_hit_tokens) .. " " .. vim.inspect(
25+
resp_json.usage.prompt_tokens
26+
) .. " + " .. vim.inspect(resp_json.usage.completion_tokens) .. " = " .. vim.inspect(
27+
resp_json.usage.total_tokens
28+
),
29+
job = job,
30+
})
31+
end
32+
end
33+
34+
---@param cmd string[]
35+
---@param callback fun(opt)
36+
local function handle_sse_events(cmd, callback)
37+
local job
38+
local tmp = ""
39+
job = vim.fn.jobstart(cmd, {
40+
on_stdout = function(_, data, _)
41+
for _, value in ipairs(data) do
42+
-- 忽略 SSE 换行输出
43+
if value ~= "" then
44+
if vim.startswith(value, "data: ") then
45+
local text = string.sub(value, 7, -1)
46+
if text == "[DONE]" then
47+
tmp = ""
48+
callback({
49+
data = text,
50+
done = true,
51+
job = job,
52+
})
53+
else
54+
local ok, resp_json = pcall(vim.fn.json_decode, text)
55+
if ok then
56+
tmp = ""
57+
callback_data(job, resp_json, callback)
58+
else
59+
tmp = text
60+
end
61+
end
62+
else
63+
if tmp ~= "" then
64+
tmp = tmp .. value
65+
local ok, resp_json = pcall(vim.fn.json_decode, tmp)
66+
if ok then
67+
callback_data(job, resp_json, callback)
68+
end
69+
else
70+
vim.notify("SSE parse error: " .. value, vim.log.levels.WARN)
71+
end
72+
end
73+
end
74+
end
75+
end,
76+
on_stderr = function(_, _, _) end,
77+
on_exit = function(_, _, _) end,
78+
})
79+
end
80+
81+
function M.request(json, callback)
82+
local body = vim.fn.json_encode(json)
83+
local cmd = {
84+
"curl",
85+
"--no-buffer",
86+
"-s",
87+
"-X",
88+
"POST",
89+
"-H",
90+
"Content-Type: application/json",
91+
"-H",
92+
"Authorization:" .. token(),
93+
"-d",
94+
body,
95+
M.config.API_URL,
96+
}
97+
handle_sse_events(cmd, callback)
98+
end
99+
100+
return M

0 commit comments

Comments
 (0)