diff --git a/README.md b/README.md index 6f20a32..8574c51 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,12 @@ sudo apt install curl export GEMINI_API_KEY="" ``` +On macOS, the Keychain Services API, which is part of the Security framework, allows you to programmatically add API_KEY to the Keychain using cli app `security` + +``` +security add-generic-password -a ${USER} -s "gemini-cli" -w "YOUR_API_KEY_HERE" +``` + * [lazy.nvim](https://github.com/folke/lazy.nvim) ```lua diff --git a/lua/gemini/api.lua b/lua/gemini/api.lua index 5ee4371..319f676 100644 --- a/lua/gemini/api.lua +++ b/lua/gemini/api.lua @@ -1,120 +1,299 @@ +-- Helper for libuv operations local uv = vim.loop or vim.uv local M = {} -local API = "https://generativelanguage.googleapis.com/v1beta/models/"; +-- Base URL for the Google Gemini API +local API = "https://generativelanguage.googleapis.com/v1beta/models/" +-- A list of available Gemini models M.MODELS = { - GEMINI_2_5_FLASH_PREVIEW = 'gemini-2.5-flash-preview-04-17', - GEMINI_2_5_PRO_PREVIEW = 'gemini-2.5-pro-preview-03-25', - GEMINI_2_0_FLASH = 'gemini-2.0-flash', - GEMINI_2_0_FLASH_LITE = 'gemini-2.0-flash-lite', - GEMINI_2_0_FLASH_EXP = 'gemini-2.0-flash-exp', - GEMINI_2_0_FLASH_THINKING_EXP = 'gemini-2.0-flash-thinking-exp-1219', - GEMINI_1_5_PRO = 'gemini-1.5-pro', - GEMINI_1_5_FLASH = 'gemini-1.5-flash', - GEMINI_1_5_FLASH_8B = 'gemini-1.5-flash-8b', + GEMINI_2_5_FLASH_PREVIEW = "gemini-2.5-flash-preview-04-17", + GEMINI_2_5_PRO_PREVIEW = "gemini-2.5-pro-preview-03-25", + GEMINI_2_5_FLASH = "gemini-2.5-flash", + GEMINI_2_5_PRO = "gemini-2.5-pro", + GEMINI_2_0_FLASH = "gemini-2.0-flash", + GEMINI_2_0_FLASH_LITE = "gemini-2.0-flash-lite", + GEMINI_2_0_FLASH_EXP = "gemini-2.0-flash-exp", + GEMINI_2_0_FLASH_THINKING_EXP = "gemini-2.0-flash-thinking-exp-1219", + GEMINI_1_5_PRO = "gemini-1.5-pro", + GEMINI_1_5_FLASH = "gemini-1.5-flash", + GEMINI_1_5_FLASH_8B = "gemini-1.5-flash-8b", } +local function get_api_key_from_file(callback) + local path = vim.fn.expand("~/.gemini/api.key") + vim.uv.fs_open(path, "r", 438, function(err, fd) + if err then + callback(nil) -- File doesn't exist or can't be opened, fallback. + return + end + vim.uv.fs_fstat(fd, function(err, stat) + if err then + vim.uv.fs_close(fd, function() end) + callback(nil) + return + end + vim.uv.fs_read(fd, stat.size, 0, function(err, data) + vim.uv.fs_close(fd, function() end) + if err then + callback(nil) + return + end + local key = vim.trim(data) + if key == "" then + vim.notify("API key file is empty.", vim.log.levels.WARN) + callback(nil) + else + vim.notify("Gemini: Successfully loaded API key from ~/.gemini/api.key") + callback(key) + end + end) + end) + end) +end + +--- Retrieves the Gemini API key. +-- It first tries to read from ~/.gemini/api.key. +-- If that fails, it falls back to the macOS keychain. +-- @param callback (function) A callback function that receives the API key. +local function get_api_key_async(callback) + get_api_key_from_file(function(key) + if key then + callback(key) + return + end + + -- Fallback to keychain if file method fails. + if vim.fn.has("mac") == 0 then + vim.notify("API key file not found and keychain access is only supported on macOS.", vim.log.levels.ERROR) + callback(nil) + return + end + + local cmd = "security" + local args = { "find-generic-password", "-l", "gemini-cli", "-w" } + local stdout = vim.loop.new_pipe(false) + local stderr = vim.loop.new_pipe(false) + local key_buffer = "" + local err_buffer = "" + + local handle + handle = vim.loop.spawn(cmd, { + args = args, + stdio = { nil, stdout, stderr }, + }, function(code, _) + stdout:close() + stderr:close() + handle:close() + + if code ~= 0 then + vim.notify( + "Gemini: Could not find API key in file or keychain.", + vim.log.levels.ERROR + ) + if err_buffer ~= "" then + vim.notify("Keychain Error: " .. err_buffer, vim.log.levels.INFO) + end + callback(nil) + else + local key = vim.trim(key_buffer) + if key == "" then + vim.notify("Gemini API key from keychain is empty.", vim.log.levels.WARN) + callback(nil) + else + vim.notify("Gemini: Successfully loaded API key from macOS keychain.") + callback(key) + end + end + end) + + vim.loop.read_start(stdout, function(err, data) + assert(not err, err) + if data then + key_buffer = key_buffer .. data + end + end) + + vim.loop.read_start(stderr, function(err, data) + assert(not err, err) + if data then + err_buffer = err_buffer .. data + end + end) + end) +end + +--- +-- Generates content using the Gemini API (non-streaming). +-- @param user_text (string) The user's prompt. +-- @param system_text (string|nil) Optional system instructions. +-- @param model_name (string) The model to use from M.MODELS. +-- @param generation_config (table) Configuration for the generation request. +-- @param callback (function|nil) Optional callback for asynchronous execution. M.gemini_generate_content = function(user_text, system_text, model_name, generation_config, callback) - local api_key = os.getenv("GEMINI_API_KEY") - if not api_key then - return '' - end - - local api = API .. model_name .. ':generateContent?key=' .. api_key - local contents = { - { - role = 'user', - parts = { - { - text = user_text - } - } - } - } - local data = { - contents = contents, - generationConfig = generation_config, - } - if system_text then - data.systemInstruction = { - role = 'user', - parts = { - { - text = system_text, - } - } - } - end - - local json_text = vim.json.encode(data) - local cmd = { 'curl', '-X', 'POST', api, '-H', 'Content-Type: application/json', '--data-binary', '@-' } - local opts = { stdin = json_text } - if callback then - return vim.system(cmd, opts, callback) - else - return vim.system(cmd, opts) - end + get_api_key_async(function(api_key) + if not api_key then + if callback then + callback(nil, "Failed to get API key") + end + return + end + + -- If the user provides a key from the MODELS table (e.g., "GEMINI_2_5_FLASH"), + -- look up its value (e.g., "gemini-2.5-flash"). + local final_model_name = M.MODELS[model_name] or model_name + local api = API .. final_model_name .. ":generateContent?key=" .. api_key + local contents = { + { + role = "user", + parts = { + { + text = user_text, + }, + }, + }, + } + local data = { + contents = contents, + generationConfig = generation_config, + } + if system_text then + data.systemInstruction = { + role = "user", + parts = { + { + text = system_text, + }, + }, + } + end + + local json_text = vim.json.encode(data) + local cmd = { "curl", "-X", "POST", api, "-H", "Content-Type: application/json", "--data-binary", "@-" } + local opts = { stdin = json_text } + + -- Execute synchronously or asynchronously based on whether a callback is provided. + if callback then + vim.system(cmd, opts, callback) + else + return vim.system(cmd, opts) + end + end) end +--- +-- Generates content using the Gemini API (streaming). +-- @param user_text (string) The user's prompt. +-- @param model_name (string) The model to use from M.MODELS. +-- @param generation_config (table) Configuration for the generation request. +-- @param callback (function) Callback to handle each streamed data chunk. M.gemini_generate_content_stream = function(user_text, model_name, generation_config, callback) - local api_key = os.getenv("GEMINI_API_KEY") - if not api_key then - return - end - - if not callback then - return - end - - local api = API .. model_name .. ':streamGenerateContent?alt=sse&key=' .. api_key - local data = { - contents = { - { - role = 'user', - parts = { - { - text = user_text - } - } - } - }, - generationConfig = generation_config, - } - local json_text = vim.json.encode(data) - - local stdin = uv.new_pipe() - local stdout = uv.new_pipe() - local stderr = uv.new_pipe() - local options = { - stdio = { stdin, stdout, stderr }, - args = { api, '-X', 'POST', '-s', '-H', 'Content-Type: application/json', '-d', json_text } - } - - uv.spawn('curl', options, function(code, _) - print("gemini chat finished exit code", code) - end) - - local streamed_data = '' - uv.read_start(stdout, function(err, data) - if not err and data then - streamed_data = streamed_data .. data - - local start_index = string.find(streamed_data, 'data:') - local end_index = string.find(streamed_data, '\r') - local json_text = '' - while start_index and end_index do - if end_index >= start_index then - json_text = string.sub(streamed_data, start_index + 5, end_index - 1) - callback(json_text) - end - streamed_data = string.sub(streamed_data, end_index + 1) - start_index = string.find(streamed_data, 'data:') - end_index = string.find(streamed_data, '\r') - end - end - end) + if not callback then + vim.notify("Streaming requires a callback function.", vim.log.levels.ERROR) + return + end + + get_api_key_async(function(api_key) + if not api_key then + return + end + + -- If the user provides a key from the MODELS table (e.g., "GEMINI_2_5_FLASH"), + -- look up its value (e.g., "gemini-2.5-flash"). + local final_model_name = M.MODELS[model_name] or model_name + local api = API .. final_model_name .. ":streamGenerateContent?alt=sse&key=" .. api_key + local data = { + contents = { + { + role = "user", + parts = { + { + text = user_text, + }, + }, + }, + }, + generationConfig = generation_config, + } + local json_text = vim.json.encode(data) + + local stdout = uv.new_pipe(false) + local stderr = uv.new_pipe(false) + local handle + local stderr_buffer = "" + + -- Arguments for the curl command. Using -N disables output buffering, which is ideal for streams. + -- The API URL is placed at the end, which is the standard position for curl. + local options = { + args = { "curl", "-X", "POST", "-s", "-N", "-H", "Content-Type: application/json", "-d", json_text, api }, + stdio = { nil, stdout, stderr }, + } + + -- Spawn the curl process. + handle = uv.spawn("curl", options, function(code, _) + -- Ensure pipes are closed when the process exits. + stdout:close() + stderr:close() + if code ~= 0 then + vim.schedule(function() + vim.notify( + "Gemini stream finished with non-zero exit code: " .. tostring(code), + vim.log.levels.ERROR + ) + if stderr_buffer ~= "" then + vim.notify("Gemini API Error: " .. stderr_buffer, vim.log.levels.ERROR) + end + end) + end + end) + + if not handle then + vim.notify("Failed to spawn curl for streaming.", vim.log.levels.ERROR) + stdout:close() + stderr:close() + return + end + + local streamed_data = "" + -- Start reading from stdout. + uv.read_start(stdout, function(err, data) + if err then + return + end + if data then + streamed_data = streamed_data .. data + + -- Process Server-Sent Events (SSE). Events are separated by double newlines. + while true do + local s, e = string.find(streamed_data, "\n\n", 1, true) + if not s then + break -- No complete event block found, wait for more data. + end + + -- Extract the complete event block. + local chunk = string.sub(streamed_data, 1, e - 1) + -- Remove the processed block from the buffer. + streamed_data = string.sub(streamed_data, e + 1) + + -- Find the 'data:' line within the event block and pass it to the callback. + local data_line = string.match(chunk, "data: (.*)") + if data_line then + callback(vim.trim(data_line)) + end + end + end + end) + + -- Optionally, read from stderr for debugging. + uv.read_start(stderr, function(err, data) + if not err and data then + stderr_buffer = stderr_buffer .. data + end + end) + end) end + + return M diff --git a/lua/gemini/chat.lua b/lua/gemini/chat.lua index 686fe73..bd090f7 100644 --- a/lua/gemini/chat.lua +++ b/lua/gemini/chat.lua @@ -25,21 +25,38 @@ M.start_chat = function(context) vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, lines) local generation_config = config.get_gemini_generation_config() - local text = '' local model_id = config.get_config({ 'model', 'model_id' }) - api.gemini_generate_content_stream(user_text, model_id, generation_config, function(json_text) - local model_response = vim.json.decode(json_text) - model_response = util.table_get(model_response, { 'candidates', 1, 'content', 'parts', 1, 'text' }) + + api.gemini_generate_content(user_text, nil, model_id, generation_config, function(obj) + if obj.code ~= 0 then + vim.notify("Gemini API Error: " .. vim.inspect(obj), vim.log.levels.ERROR) + return + end + + local data = obj.stdout + if data == "" or data == nil then + vim.notify("Gemini API returned empty response.", vim.log.levels.ERROR) + return + end + + local response_data = vim.json.decode(data) + if not response_data then + vim.notify("Failed to decode Gemini API response: " .. data, vim.log.levels.ERROR) + return + end + + local model_response = util.table_get(response_data, { 'candidates', 1, 'content', 'parts', 1, 'text' }) if not model_response then + vim.notify("Unexpected API response structure: " .. data, vim.log.levels.ERROR) return end - text = text .. model_response vim.schedule(function() - lines = vim.split(text, '\n') + lines = vim.split(model_response, '\n') vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, lines) end) end) end + return M diff --git a/tests/gemini/api_spec.lua b/tests/gemini/api_spec.lua index 12cac15..004903b 100644 --- a/tests/gemini/api_spec.lua +++ b/tests/gemini/api_spec.lua @@ -3,25 +3,38 @@ local util = require('gemini.util') describe('api', function() it('should send message', function() + local completed = false + local result_table + local generation_config = { temperature = 0.9, top_k = 1.0, max_output_tokens = 2048, response_mime_type = 'text/plain', } - local future = api.gemini_generate_content('hello there', nil, api.MODELS.GEMINI_2_0_FLASH, generation_config, nil) - local result = future:wait() - local stdout = result.stdout - print(stdout) - assert(#stdout > 0) - local result = vim.json.decode(stdout) + api.gemini_generate_content('hello there', nil, api.MODELS.GEMINI_2_0_FLASH, generation_config, function(result) + result_table = result + completed = true + end) + + vim.wait(5000, function() return completed end) + + assert.is_not_nil(result_table) + assert.is_equal(result_table.code, 0) + assert.is_not_nil(result_table.stdout) + assert(#result_table.stdout > 0) + + local result = vim.json.decode(result_table.stdout) local model_response = util.table_get(result, { 'candidates', 1, 'content', 'parts', 1, 'text' }) assert(#model_response > 0) end) it('should send long message', function() + local completed = false + local result_table + local generation_config = { temperature = 0.9, top_k = 1.0, @@ -29,13 +42,20 @@ describe('api', function() response_mime_type = 'text/plain', } local long_message = string.rep('this is a very very long message ', 3000) - local future = api.gemini_generate_content(long_message, nil, api.MODELS.GEMINI_2_0_FLASH, generation_config, nil) - local result = future:wait() - local stdout = result.stdout - print(stdout) - assert(#stdout > 0) - local result = vim.json.decode(stdout) + api.gemini_generate_content(long_message, nil, api.MODELS.GEMINI_2_0_FLASH, generation_config, function(result) + result_table = result + completed = true + end) + + vim.wait(20000, function() return completed end) + + assert.is_not_nil(result_table) + assert.is_equal(result_table.code, 0) + assert.is_not_nil(result_table.stdout) + assert(#result_table.stdout > 0) + + local result = vim.json.decode(result_table.stdout) local model_response = util.table_get(result, { 'candidates', 1, 'content', 'parts', 1, 'text' }) assert(#model_response > 0)