diff --git a/local/ai-provider/lua/ai-provider/core.lua b/local/ai-provider/lua/ai-provider/core.lua new file mode 100644 index 0000000..67d8b41 --- /dev/null +++ b/local/ai-provider/lua/ai-provider/core.lua @@ -0,0 +1,510 @@ +local M = {} + +local providers = { + ollama = require 'ai-provider.providers.ollama', +} + +local state = { + config = { + providers = {}, + }, + model_cache = {}, +} + +local function is_configured_provider(provider) + return type(state.config.providers[provider]) == 'table' +end + +local function get_provider_config(provider) + return state.config.providers[provider] or {} +end + +local function validate_config(opts) + if type(opts) ~= 'table' then + error('ai-provider setup requires a config table', 3) + end + + if type(opts.default_provider) ~= 'string' or opts.default_provider == '' then + error('ai-provider setup requires default_provider', 3) + end + + if type(opts.providers) ~= 'table' then + error('ai-provider setup requires providers table', 3) + end + + if not opts.providers[opts.default_provider] then + error('ai-provider default_provider must be configured under providers: ' .. opts.default_provider, 3) + end + + for provider, provider_config in pairs(opts.providers) do + if not providers[provider] then + error('ai-provider config references unknown provider: ' .. provider, 3) + end + if type(provider_config.default_model) ~= 'string' or provider_config.default_model == '' then + error('ai-provider setup requires providers.' .. provider .. '.default_model', 3) + end + if provider_config.models ~= nil and type(provider_config.models) ~= 'table' then + error('ai-provider providers.' .. provider .. '.models must be a table when set', 3) + end + end +end + +local function preferences_file() + return vim.fn.stdpath 'data' .. '/ai-provider-preferences.json' +end + +function M.load_preferences() + local file = io.open(preferences_file(), 'r') + if not file then + return {} + end + + local content = file:read '*a' + file:close() + + local ok, prefs = pcall(vim.json.decode, content) + if ok and type(prefs) == 'table' then + return prefs + end + return {} +end + +function M.save_preferences(prefs) + local file = io.open(preferences_file(), 'w') + if not file then + return false + end + + file:write(vim.json.encode(prefs)) + file:close() + return true +end + +function M.get_provider(name) + if not is_configured_provider(name) then + return nil + end + return providers[name] +end + +function M.get_provider_implementation(name) + return providers[name] +end + +function M.get_provider_config(name) + if not is_configured_provider(name) then + return nil + end + return get_provider_config(name) +end + +function M.list_providers() + local names = {} + for name in pairs(state.config.providers) do + table.insert(names, name) + end + table.sort(names) + return names +end + +function M.get_default_provider() + local prefs = M.load_preferences() + if type(prefs.default_provider) == 'string' and is_configured_provider(prefs.default_provider) then + return prefs.default_provider + end + if type(state.config.default_provider) == 'string' and is_configured_provider(state.config.default_provider) then + return state.config.default_provider + end + return nil +end + +function M.set_default_provider(provider) + if not is_configured_provider(provider) then + return false + end + + local prefs = M.load_preferences() + prefs.default_provider = provider + return M.save_preferences(prefs) +end + +function M.get_selected_model(provider) + local prefs = M.load_preferences() + local provider_prefs = prefs[provider] + if type(provider_prefs) == 'table' and type(provider_prefs.model) == 'string' and provider_prefs.model ~= '' then + return provider_prefs.model + end + + local provider_config = state.config.providers[provider] + if type(provider_config) == 'table' and type(provider_config.default_model) == 'string' and provider_config.default_model ~= '' then + return provider_config.default_model + end + + return nil +end + +function M.set_selected_model(provider, model) + if not is_configured_provider(provider) then + return false + end + + local prefs = M.load_preferences() + prefs[provider] = prefs[provider] or {} + prefs[provider].model = model + return M.save_preferences(prefs) +end + +function M.check(provider, callback, opts) + local implementation = M.get_provider(provider) + if not implementation or not implementation.check then + callback(false) + return + end + implementation.check(callback, opts) +end + +function M.auth(provider, callback, opts) + local implementation = M.get_provider(provider) + if not implementation or not implementation.auth then + callback(false) + return + end + implementation.auth(callback, opts) +end + +function M.list_models(provider, callback, opts) + local implementation = M.get_provider(provider) + if not implementation or not implementation.list_models then + callback(nil) + return + end + opts = vim.tbl_extend('force', { provider_config = get_provider_config(provider) }, opts or {}) + implementation.list_models(function(models) + if models then + state.model_cache[provider] = models + end + callback(models) + end, opts) +end + +local function normalize_chat_args(first, second) + if type(first) == 'string' and type(second) == 'table' and is_configured_provider(first) then + local request = vim.tbl_extend('force', {}, second) + request.provider = first + return first, request + end + + local request = first + if type(request) == 'string' then + request = { prompt = request, callback = second } + elseif type(request) == 'table' then + request = vim.tbl_extend('force', {}, request) + else + request = {} + end + + local provider = request.provider or M.get_default_provider() + if not provider then + vim.notify('ai-provider is missing default_provider. Configure it in setup().', vim.log.levels.ERROR) + end + request.provider = provider + return provider, request +end + +function M.chat(first, second) + local provider, request = normalize_chat_args(first, second) + local implementation = M.get_provider(provider) + if not implementation or not implementation.chat then + if request.callback then + request.callback(nil, nil) + end + return nil + end + + request.model = request.model or M.get_selected_model(provider) + request.provider_config = get_provider_config(provider) + if not request.model then + if request.callback then + request.callback(nil, nil) + end + vim.notify('No model selected for ' .. provider .. '. Run :AIProvider ' .. provider .. ' model first.', vim.log.levels.ERROR) + return nil + end + + return implementation.chat(request) +end + +function M.chat_with(provider, request) + return M.chat(provider, request) +end + +local function model_display(provider, model) + return provider .. '/' .. model +end + +local function collect_models(callback) + local all = {} + local provider_names = M.list_providers() + local pending = #provider_names + + if pending == 0 then + callback(all) + return + end + + for _, provider in ipairs(provider_names) do + M.list_models(provider, function(models) + for _, model in ipairs(models or {}) do + table.insert(all, { provider = provider, model = model, label = model_display(provider, model) }) + end + + pending = pending - 1 + if pending == 0 then + table.sort(all, function(a, b) + return a.label < b.label + end) + callback(all) + end + end) + end +end + +function M.select_model(provider) + if not provider then + collect_models(function(models) + if #models == 0 then + vim.notify('No AI provider models are available.', vim.log.levels.WARN) + return + end + + local default_provider = M.get_default_provider() + vim.ui.select(models, { + prompt = 'Select AI model:', + format_item = function(item) + local current = item.provider == default_provider and item.model == M.get_selected_model(item.provider) + local marker = current and '✓ ' or ' ' + return marker .. item.label + end, + }, function(choice) + if not choice then + return + end + M.set_default_provider(choice.provider) + M.set_selected_model(choice.provider, choice.model) + vim.notify('Default AI model set to ' .. choice.label, vim.log.levels.INFO) + end) + end) + return + end + + M.check(provider, function(working) + if not working then + vim.notify(provider .. ' is not reachable', vim.log.levels.ERROR) + return + end + + M.auth(provider, function(authenticated) + if not authenticated then + vim.notify(provider .. ' authentication failed', vim.log.levels.ERROR) + return + end + + M.list_models(provider, function(models) + if not models or #models == 0 then + vim.notify(provider .. ' is reachable but no models are available.', vim.log.levels.WARN) + return + end + + local current = M.get_selected_model(provider) + vim.ui.select(models, { + prompt = 'Select ' .. provider .. ' model:', + format_item = function(model) + local marker = model == current and '✓ ' or ' ' + return marker .. model + end, + }, function(choice) + if not choice then + return + end + M.set_selected_model(provider, choice) + vim.notify(provider .. ' model set to ' .. choice, vim.log.levels.INFO) + end) + end) + end) + end, { force = true }) +end + +local global_actions = { 'default', 'model', 'models' } +local provider_actions = { 'auth', 'check', 'model', 'models' } + +local function starts_with(value, prefix) + return value:sub(1, #prefix) == prefix +end + +local function filter(values, prefix) + local matches = {} + for _, value in ipairs(values) do + if starts_with(value, prefix) then + table.insert(matches, value) + end + end + return matches +end + +function M.command_complete(arglead, cmdline) + local parts = vim.split(cmdline, '%s+', { trimempty = true }) + if cmdline:match '%s$' then + table.insert(parts, '') + end + + local argc = math.max(#parts - 1, 0) + if argc <= 1 then + local first = vim.deepcopy(global_actions) + vim.list_extend(first, M.list_providers()) + return filter(first, arglead) + end + + local first_arg = parts[2] + if first_arg == 'default' then + return argc == 2 and filter(M.list_providers(), arglead) or {} + end + + if first_arg == 'model' then + if argc == 2 then + local cached = {} + for _, provider in ipairs(M.list_providers()) do + for _, model in ipairs(state.model_cache[provider] or {}) do + table.insert(cached, model_display(provider, model)) + end + end + table.sort(cached) + return filter(cached, arglead) + end + return {} + end + + if first_arg == 'models' then + return {} + end + + if not providers[first_arg] then + return {} + end + + if argc == 2 then + return filter(provider_actions, arglead) + end + + if parts[3] == 'model' then + return filter(state.model_cache[first_arg] or {}, arglead) + end + + return {} +end + +local function print_lines(lines) + vim.api.nvim_echo({ { table.concat(lines, '\n'), 'Normal' } }, true, {}) +end + +function M.run_command(args) + if #args == 0 then + M.select_model() + return + end + + if args[1] == 'model' then + if args[2] then + local provider, model = args[2]:match '^([^/]+)/(.+)$' + if not provider or not model or not is_configured_provider(provider) then + vim.notify('Expected model as provider/model, for example ollama/gemma4:e2b', vim.log.levels.ERROR) + return + end + + M.set_default_provider(provider) + M.set_selected_model(provider, model) + vim.notify('Default AI model set to ' .. model_display(provider, model), vim.log.levels.INFO) + else + M.select_model() + end + return + end + + if args[1] == 'models' then + collect_models(function(models) + local lines = vim.tbl_map(function(item) + return item.label + end, models) + print_lines(lines) + end) + return + end + + if args[1] == 'default' then + if args[2] then + if M.set_default_provider(args[2]) then + vim.notify('Default AI provider set to ' .. args[2], vim.log.levels.INFO) + else + vim.notify('Unknown AI provider: ' .. args[2], vim.log.levels.ERROR) + end + return + end + + local provider = M.get_default_provider() + local model = M.get_selected_model(provider) + local suffix = model and (' (' .. model_display(provider, model) .. ')') or '' + vim.notify('Default AI provider: ' .. provider .. suffix, vim.log.levels.INFO) + return + end + + local provider = args[1] + if not is_configured_provider(provider) then + vim.notify('Unknown AI provider: ' .. provider, vim.log.levels.ERROR) + return + end + + local action = args[2] or 'model' + if action == 'check' then + M.check(provider, function(working) + vim.notify(provider .. ' working: ' .. tostring(working), working and vim.log.levels.INFO or vim.log.levels.ERROR) + end, { force = true }) + elseif action == 'auth' then + M.auth(provider, function(authenticated) + vim.notify(provider .. ' authenticated: ' .. tostring(authenticated), authenticated and vim.log.levels.INFO or vim.log.levels.ERROR) + end) + elseif action == 'model' then + if args[3] then + M.set_selected_model(provider, args[3]) + vim.notify(provider .. ' model set to ' .. args[3], vim.log.levels.INFO) + else + M.select_model(provider) + end + elseif action == 'models' then + M.list_models(provider, function(models) + print_lines(models or {}) + end) + else + vim.notify('Unknown AIProvider action: ' .. action, vim.log.levels.ERROR) + end +end + +function M.setup(opts) + validate_config(opts) + state.config = vim.tbl_deep_extend('force', { providers = {} }, opts) + + vim.api.nvim_create_user_command('AIProvider', function(command) + M.run_command(command.fargs) + end, { + nargs = '*', + complete = M.command_complete, + desc = 'Manage AI providers', + force = true, + }) + + vim.defer_fn(function() + for _, provider in ipairs(M.list_providers()) do + M.list_models(provider, function() end) + end + end, 1000) +end + +return M diff --git a/local/ai-provider/lua/ai-provider/curl.lua b/local/ai-provider/lua/ai-provider/curl.lua new file mode 100644 index 0000000..016d9d5 --- /dev/null +++ b/local/ai-provider/lua/ai-provider/curl.lua @@ -0,0 +1,160 @@ +local M = {} + +local DEFAULT_TIMEOUT = 30000 + +local function schedule(callback, ...) + local args = { ... } + vim.schedule(function() + callback(unpack(args)) + end) +end + +local function encode_body(body) + if type(body) == 'string' then + return body + end + if body ~= nil then + return vim.json.encode(body) + end + return nil +end + +local function json_headers(headers) + headers = vim.tbl_extend('force', {}, headers or {}) + headers['Content-Type'] = headers['Content-Type'] or 'application/json' + return headers +end + +---@class AiProviderCurlResponse +---@field status integer HTTP status code. `0` means curl/plenary failed before HTTP response parsing. +---@field body string Raw response body. +---@field json table|nil Decoded JSON body when decoding succeeded. +---@field error string|nil Transport or decode error message. + +---@class AiProviderCurlRequest +---@field method? string HTTP method. Defaults to `GET`. +---@field url string Request URL. +---@field headers? table HTTP headers. +---@field body? table|string Request body. Tables are JSON encoded. +---@field timeout? integer Request timeout in milliseconds. +---@field callback fun(response: AiProviderCurlResponse) Completion callback. + +---Run a non-streaming HTTP request and decode JSON responses when possible. +---@param request AiProviderCurlRequest +function M.json(request) + local curl = require 'plenary.curl' + local method = string.lower(request.method or 'GET') + local body = encode_body(request.body) + + local opts = { + headers = json_headers(request.headers), + body = body, + timeout = request.timeout or DEFAULT_TIMEOUT, + callback = function(response) + local decoded = nil + if type(response.body) == 'string' and response.body ~= '' then + local ok, data = pcall(vim.json.decode, response.body) + if ok then + decoded = data + end + end + + schedule(request.callback, { + status = response.status or 0, + body = response.body or '', + json = decoded, + }) + end, + on_error = function(err) + schedule(request.callback, { + status = 0, + body = '', + error = tostring(err and (err.stderr or err.message) or 'curl request failed'), + }) + end, + } + + if method == 'post' then + return curl.post(request.url, opts) + elseif method == 'put' then + return curl.put(request.url, opts) + elseif method == 'delete' then + return curl.delete(request.url, opts) + end + + return curl.get(request.url, opts) +end + +---@class AiProviderCurlStreamRequest +---@field method? string HTTP method. Defaults to `POST`. +---@field url string Request URL. +---@field headers? table HTTP headers. +---@field body? table|string Request body. Tables are JSON encoded. +---@field timeout? integer Curl max-time in milliseconds. +---@field is_cancelled? fun(): boolean Optional cancellation predicate. +---@field on_json_line? fun(data: table, line: string) Called for each decoded JSON line. +---@field callback fun(code: integer, error: string|nil) Called when curl exits. + +---Run a streaming curl request and decode each stdout line as JSON. +---@param request AiProviderCurlStreamRequest +---@return table job Plenary job handle. +function M.stream_json_lines(request) + local Job = require 'plenary.job' + local body = encode_body(request.body) + local method = request.method or 'POST' + local timeout = request.timeout or DEFAULT_TIMEOUT + local args = { + '--silent', + '--show-error', + '--no-buffer', + '--max-time', + tostring(math.ceil(timeout / 1000)), + '--request', + method, + } + + for name, value in pairs(json_headers(request.headers)) do + table.insert(args, '--header') + table.insert(args, name .. ': ' .. value) + end + + if body then + table.insert(args, '--data') + table.insert(args, body) + end + + table.insert(args, request.url) + + local stderr = {} + local job = Job:new { + command = 'curl', + args = args, + on_stdout = function(_, line) + if not line or line == '' or (request.is_cancelled and request.is_cancelled()) then + return + end + + local ok, data = pcall(vim.json.decode, line) + if not ok or type(data) ~= 'table' then + return + end + + if request.on_json_line then + request.on_json_line(data, line) + end + end, + on_stderr = function(_, line) + if line and line ~= '' then + table.insert(stderr, line) + end + end, + on_exit = function(_, code) + schedule(request.callback, code, #stderr > 0 and table.concat(stderr, '\n') or nil) + end, + } + + job:start() + return job +end + +return M diff --git a/local/ai-provider/lua/ai-provider/init.lua b/local/ai-provider/lua/ai-provider/init.lua new file mode 100644 index 0000000..4a6b021 --- /dev/null +++ b/local/ai-provider/lua/ai-provider/init.lua @@ -0,0 +1,196 @@ +---@meta + +---@class AiProviderPreferences +---@field default_provider? string Default provider used by chat calls without an explicit provider. +---@field [string] table Provider-specific preferences. Provider tables currently support `model`. + +---@class AiProviderProviderConfig +---@field default_model string Default model used when no persisted model exists for this provider. +---@field timeout? integer Optional provider timeout in milliseconds. +---@field context_size? integer Optional provider-wide default context size. +---@field keep_alive? string|integer Optional provider-wide keep-alive/unload timeout. Ollama accepts values like `4h`, `10m`, or `0`. +---@field models? table Optional logical model profiles. Keys are selectable model names. + +---@class AiProviderModelConfig +---@field model string Underlying provider model name. +---@field context_size? integer Optional model-specific context size override. + +---@class AiProviderConfig +---@field default_provider string Default provider used when no persisted default provider exists. +---@field providers table Required provider defaults. + +---@class AiProviderChatMeta +---@field requested_model string Model requested by the caller. +---@field used_model string Model reported by the provider response. +---@field elapsed_ms number Request duration in milliseconds. + +---@class AiProviderChatRequest +---@field provider? string Provider name. Defaults to `get_default_provider()`. +---@field model? string Model name. Defaults to the selected model for the provider. +---@field prompt string User prompt to send to the provider. +---@field stream? boolean Whether the provider should stream chunks. Defaults to true when supported. +---@field max_tokens? integer Maximum generated tokens/provider equivalent. +---@field context_size? integer Per-request context size override. +---@field keep_alive? string|integer Per-request keep-alive/unload timeout override. +---@field on_chunk? fun(chunk: string, raw: table) Called for each streamed text chunk. +---@field callback? fun(message: string|nil, meta: AiProviderChatMeta|nil) Called once the request finishes. +---@field is_cancelled? fun(): boolean Optional cancellation predicate. +---@field register_http_job? fun(job: table) Receives the provider job/process for external cancellation. + +---@class AiProvider +local M = {} + +local core = require 'ai-provider.core' + +---@return AiProviderPreferences prefs Saved provider preferences. +function M.load_preferences() + return core.load_preferences() +end + +---@param prefs AiProviderPreferences +---@return boolean saved Whether preferences were written. +function M.save_preferences(prefs) + return core.save_preferences(prefs) +end + +---@param name string Provider name, for example `ollama`. +---@return table|nil provider Active provider implementation table. Returns nil for known but unconfigured providers. +function M.get_provider(name) + return core.get_provider(name) +end + +---@param name string Provider name, for example `ollama`. +---@return table|nil provider Provider implementation table, even when not enabled by config. +function M.get_provider_implementation(name) + return core.get_provider_implementation(name) +end + +---@param name string Provider name, for example `ollama`. +---@return AiProviderProviderConfig|nil config Active provider config. +function M.get_provider_config(name) + return core.get_provider_config(name) +end + +---@return string[] providers Configured provider names. +function M.list_providers() + return core.list_providers() +end + +---@return string|nil provider Default provider name, or nil before valid setup. +function M.get_default_provider() + return core.get_default_provider() +end + +---@return string|nil provider Current/default provider name, or nil before valid setup. +function M.get_current_provider() + return core.get_default_provider() +end + +---@param provider string Provider name. +---@return boolean saved Whether the provider was saved as default. +function M.set_default_provider(provider) + return core.set_default_provider(provider) +end + +---@param provider string Provider name. +---@return boolean saved Whether the provider was saved as default. +function M.set_provider(provider) + return core.set_default_provider(provider) +end + +---@param provider string Provider name. +---@return string|nil model Selected model for the provider. +function M.get_selected_model(provider) + return core.get_selected_model(provider) +end + +---@param provider string Provider name. +---@return string|nil model Selected model for the provider. +function M.get_model(provider) + return core.get_selected_model(provider) +end + +---@param provider string Provider name. +---@param model string Model name. +---@return boolean saved Whether the model preference was saved. +function M.set_selected_model(provider, model) + return core.set_selected_model(provider, model) +end + +---@param provider string Provider name. +---@param model string Model name. +---@return boolean saved Whether the model preference was saved. +function M.set_model(provider, model) + return core.set_selected_model(provider, model) +end + +---@param provider string Provider name. +---@param callback fun(working: boolean) +---@param opts? table Provider-specific check options. +function M.check(provider, callback, opts) + return core.check(provider, callback, opts) +end + +---@param provider string Provider name. +---@param callback fun(authenticated: boolean) +---@param opts? table Provider-specific auth options. +function M.auth(provider, callback, opts) + return core.auth(provider, callback, opts) +end + +---@param provider string Provider name. +---@param callback fun(models: string[]|nil) +---@param opts? table Provider-specific list options. +function M.list_models(provider, callback, opts) + return core.list_models(provider, callback, opts) +end + +---Chat with the default provider, or with an explicit provider when called as +---`chat('ollama', request)`. Responses are delivered via `request.callback`. +---@param request_or_provider AiProviderChatRequest|string Prompt string, request table, or explicit provider name. +---@param callback_or_request? fun(message: string|nil, meta: AiProviderChatMeta|nil)|AiProviderChatRequest Callback for prompt strings, or request table for explicit providers. +---@return table|nil job Provider job/process handle. +function M.chat(request_or_provider, callback_or_request) + return core.chat(request_or_provider, callback_or_request) +end + +---@param provider string Provider name. +---@param request AiProviderChatRequest Chat request. +---@return table|nil job Provider job/process handle. +function M.chat_with(provider, request) + return core.chat_with(provider, request) +end + +---@param provider? string Provider name. Omit to pick across all providers. +function M.select_model(provider) + return core.select_model(provider) +end + +---@param arglead string Current command-line argument prefix. +---@param cmdline string Full command-line. +---@return string[] completions Matching command completions. +function M.command_complete(arglead, cmdline) + return core.command_complete(arglead, cmdline) +end + +---@param args string[] Parsed `:AIProvider` arguments. +function M.run_command(args) + return core.run_command(args) +end + +---Register user commands for the provider layer. `default_provider` and each +---registered provider's `default_model` are required config values; the plugin +---does not ship provider/model defaults. +--- +---Command semantics: +---- `:AIProvider` opens the all-provider model picker. +---- `:AIProvider model` opens the all-provider model picker. +---- `:AIProvider model provider/model` sets the default provider and that provider's default model. +---- `:AIProvider default [provider]` shows or sets the default provider. +---- `:AIProvider model [model]` picks or sets only that provider's default model. +---@param opts AiProviderConfig +function M.setup(opts) + return core.setup(opts) +end + +return M diff --git a/local/ai-provider/lua/ai-provider/log.lua b/local/ai-provider/lua/ai-provider/log.lua new file mode 100644 index 0000000..43bba19 --- /dev/null +++ b/local/ai-provider/lua/ai-provider/log.lua @@ -0,0 +1,52 @@ +local M = {} + +local state = { + logger = nil, +} + +local function setup_logger() + if state.logger then + return state.logger + end + + local ok, plenary_log = pcall(require, 'plenary.log') + if not ok then + local noop = function() end + state.logger = { debug = noop, info = noop, warn = noop, error = noop } + return state.logger + end + + state.logger = plenary_log.new { + plugin = 'ai-provider', + level = 'debug', + use_console = false, + } + + return state.logger +end + +function M.debug(message) + local logger = setup_logger() + ---@cast logger any + logger.debug(message) +end + +function M.info(message) + local logger = setup_logger() + ---@cast logger any + logger.info(message) +end + +function M.warn(message) + local logger = setup_logger() + ---@cast logger any + logger.warn(message) +end + +function M.error(message) + local logger = setup_logger() + ---@cast logger any + logger.error(message) +end + +return M diff --git a/local/ai-provider/lua/ai-provider/providers/ollama.lua b/local/ai-provider/lua/ai-provider/providers/ollama.lua new file mode 100644 index 0000000..cea9c83 --- /dev/null +++ b/local/ai-provider/lua/ai-provider/providers/ollama.lua @@ -0,0 +1,171 @@ +local M = {} +local curl = require 'ai-provider.curl' + +local DEFAULT_ENDPOINT = 'http://127.0.0.1:11434' +local HEALTH_CACHE_TTL = 30 + +local state = { + health = nil, + health_checked_at = 0, +} + +local function endpoint() + local configured = vim.g.ollama_endpoint or os.getenv 'OLLAMA_HOST' + if configured and configured ~= '' then + configured = configured:gsub('/+$', '') + if not configured:match '^https?://' then + configured = 'http://' .. configured + end + return configured + end + return DEFAULT_ENDPOINT +end + +local function is_cancelled(request) + return request and request.is_cancelled and request.is_cancelled() +end + +local function configured_models(provider_config) + local models = {} + for name in pairs(provider_config.models or {}) do + table.insert(models, name) + end + table.sort(models) + return models +end + +local function model_config(provider_config, model) + local configured = provider_config.models and provider_config.models[model] + if type(configured) == 'string' then + return { model = configured } + end + if type(configured) == 'table' then + return configured + end + return { model = model } +end + +function M.check(callback, opts) + opts = opts or {} + local now = os.time() + if not opts.force and state.health ~= nil and now - state.health_checked_at < HEALTH_CACHE_TTL then + callback(state.health) + return + end + + curl.json { + url = endpoint() .. '/api/tags', + timeout = opts.timeout or 1000, + callback = function(response) + state.health = response.status == 200 + state.health_checked_at = os.time() + callback(state.health) + end, + } +end + +function M.auth(callback, opts) + opts = opts or {} + callback(true) +end + +function M.list_models(callback, opts) + opts = opts or {} + local provider_config = opts.provider_config or {} + curl.json { + url = endpoint() .. '/api/tags', + timeout = opts.timeout or 3000, + callback = function(response) + if response.status ~= 200 or type(response.json) ~= 'table' or type(response.json.models) ~= 'table' then + callback(nil) + return + end + + local models = {} + local seen = {} + for _, model in ipairs(configured_models(provider_config)) do + seen[model] = true + table.insert(models, model) + end + + for _, model in ipairs(response.json.models) do + if type(model) == 'table' and type(model.name) == 'string' and not seen[model.name] then + seen[model.name] = true + table.insert(models, model.name) + end + end + table.sort(models) + callback(models) + end, + } +end + +function M.chat(request) + local started_at = vim.uv.hrtime() + local chunks = {} + local provider_config = request.provider_config or {} + local selected_model = request.model + local selected_config = model_config(provider_config, selected_model) + local raw_model = selected_config.model or selected_model + local context_size = request.context_size or selected_config.context_size or provider_config.context_size + local keep_alive = request.keep_alive or provider_config.keep_alive + local final_model = raw_model + + local body = { + model = raw_model, + messages = { { role = 'user', content = request.prompt } }, + stream = request.stream ~= false, + options = { + num_predict = request.max_tokens, + num_ctx = context_size, + }, + keep_alive = keep_alive, + } + + local job = curl.stream_json_lines { + url = endpoint() .. '/api/chat', + body = body, + timeout = request.timeout or 30000, + is_cancelled = request.is_cancelled, + on_json_line = function(data) + if is_cancelled(request) then + return + end + + final_model = data.model or final_model + local chunk = data.message and data.message.content or '' + if chunk ~= '' then + table.insert(chunks, chunk) + if request.on_chunk then + vim.schedule(function() + request.on_chunk(chunk, data) + end) + end + end + end, + callback = function(code) + if is_cancelled(request) then + return + end + + if code ~= 0 then + if request.callback then + request.callback(nil, nil) + end + return + end + + local elapsed_ms = (vim.uv.hrtime() - started_at) / 1e6 + if request.callback then + request.callback(table.concat(chunks, ''), { requested_model = selected_model, used_model = final_model, elapsed_ms = elapsed_ms }) + end + end, + } + + if request.register_http_job then + request.register_http_job(job) + end + return job +end + +return M diff --git a/local/ai-provider/run_tests.sh b/local/ai-provider/run_tests.sh new file mode 100755 index 0000000..fa8cbf5 --- /dev/null +++ b/local/ai-provider/run_tests.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/../.." + +nvim --headless --noplugin \ + -u local/ai-provider/tests/minimal_init.lua \ + -c "PlenaryBustedDirectory local/ai-provider/tests/ai_provider { minimal_init = 'local/ai-provider/tests/minimal_init.lua' }" diff --git a/local/ai-provider/tests/ai_provider/core_spec.lua b/local/ai-provider/tests/ai_provider/core_spec.lua new file mode 100644 index 0000000..7cbe480 --- /dev/null +++ b/local/ai-provider/tests/ai_provider/core_spec.lua @@ -0,0 +1,29 @@ +local ai_provider = require 'ai-provider' + +describe('ai-provider core', function() + it('exposes only configured providers', function() + ai_provider.setup { + default_provider = 'ollama', + providers = { + ollama = { default_model = 'gemma4:e2b' }, + }, + } + + assert.are.same({ 'ollama' }, ai_provider.list_providers()) + assert.is_truthy(ai_provider.get_provider 'ollama') + assert.is_truthy(ai_provider.get_provider_implementation 'ollama') + assert.is_nil(ai_provider.get_provider 'codex') + end) + + it('rejects missing default provider config', function() + local ok, err = pcall(ai_provider.setup, { + default_provider = 'codex', + providers = { + ollama = { default_model = 'gemma4:e2b' }, + }, + }) + + assert.is_false(ok) + assert.matches('default_provider must be configured', err) + end) +end) diff --git a/local/ai-provider/tests/ai_provider/ollama_options_spec.lua b/local/ai-provider/tests/ai_provider/ollama_options_spec.lua new file mode 100644 index 0000000..eca1248 --- /dev/null +++ b/local/ai-provider/tests/ai_provider/ollama_options_spec.lua @@ -0,0 +1,98 @@ +describe('ollama provider options', function() + local original_curl + + before_each(function() + original_curl = package.loaded['ai-provider.curl'] + package.loaded['ai-provider.providers.ollama'] = nil + end) + + after_each(function() + package.loaded['ai-provider.curl'] = original_curl + package.loaded['ai-provider.providers.ollama'] = nil + end) + + it('sends profile context size and provider keep_alive', function() + ---@type table|nil + local captured = nil + package.loaded['ai-provider.curl'] = { + stream_json_lines = function(request) + captured = request.body + request.on_json_line { model = 'gemma4:e2b', message = { content = 'ok' } } + request.callback(0) + return { shutdown = function() end } + end, + } + + local ollama = require 'ai-provider.providers.ollama' + local done = false + local message = nil + + ollama.chat { + model = 'gemma4:e2b 64k', + prompt = 'Reply with exactly: ok', + max_tokens = 4, + provider_config = { + context_size = 1024 * 8, + keep_alive = '4h', + models = { + ['gemma4:e2b 64k'] = { + model = 'gemma4:e2b', + context_size = 1024 * 64, + }, + }, + }, + callback = function(result) + message = result + done = true + end, + } + + assert.is_true(done) + assert.are.same('ok', message) + assert.is_table(captured) + ---@cast captured table + assert.is_table(captured.options) + assert.are.same('gemma4:e2b', captured.model) + assert.are.same(4, captured.options.num_predict) + assert.are.same(1024 * 64, captured.options.num_ctx) + assert.are.same('4h', captured.keep_alive) + end) + + it('lets request options override profile and provider options', function() + ---@type table|nil + local captured = nil + package.loaded['ai-provider.curl'] = { + stream_json_lines = function(request) + captured = request.body + request.callback(0) + return { shutdown = function() end } + end, + } + + local ollama = require 'ai-provider.providers.ollama' + + ollama.chat { + model = 'gemma4:e2b 32k', + prompt = 'Reply with exactly: ok', + context_size = 1024 * 16, + keep_alive = '10m', + provider_config = { + context_size = 1024 * 8, + keep_alive = '4h', + models = { + ['gemma4:e2b 32k'] = { + model = 'gemma4:e2b', + context_size = 1024 * 32, + }, + }, + }, + callback = function() end, + } + + assert.is_table(captured) + ---@cast captured table + assert.is_table(captured.options) + assert.are.same(1024 * 16, captured.options.num_ctx) + assert.are.same('10m', captured.keep_alive) + end) +end) diff --git a/local/ai-provider/tests/ai_provider/ollama_spec.lua b/local/ai-provider/tests/ai_provider/ollama_spec.lua new file mode 100644 index 0000000..a23cd8d --- /dev/null +++ b/local/ai-provider/tests/ai_provider/ollama_spec.lua @@ -0,0 +1,126 @@ +local ai_provider = require 'ai-provider' + +local config = { + default_provider = 'ollama', + providers = { + ollama = { + default_model = 'gemma4:e2b', + context_size = 1024 * 8, + keep_alive = '4h', + models = { + ['gemma4:e2b 32k'] = { + model = 'gemma4:e2b', + context_size = 1024 * 32, + }, + ['gemma4:e2b 64k'] = { + model = 'gemma4:e2b', + context_size = 1024 * 64, + }, + }, + }, + }, +} + +local function wait_for(done, timeout) + local ok = vim.wait(timeout or 20000, function() + return done() + end, 50) + assert.is_true(ok) +end + +describe('ollama provider integration', function() + before_each(function() + ai_provider.setup(config) + end) + + it('checks provider availability', function() + local done = false + local working = nil + + ai_provider.check('ollama', function(result) + working = result + done = true + end, { force = true, timeout = 1000 }) + + wait_for(function() + return done + end) + assert.is_true(working) + end) + + it('lists raw models and configured logical profiles', function() + local done = false + ---@type string[]|nil + local models = nil + + ai_provider.list_models('ollama', function(result) + models = result + done = true + end) + + wait_for(function() + return done + end) + + assert.is_table(models) + ---@cast models string[] + assert.is_true(vim.tbl_contains(models, 'gemma4:e2b')) + assert.is_true(vim.tbl_contains(models, 'gemma4:e2b 32k')) + assert.is_true(vim.tbl_contains(models, 'gemma4:e2b 64k')) + end) + + it('runs a prompt through a logical model profile', function() + local done = false + local message = nil + ---@type table|nil + local meta = nil + + ai_provider.chat('ollama', { + model = 'gemma4:e2b 32k', + prompt = 'Reply with exactly: ok', + max_tokens = 4, + timeout = 10000, + callback = function(result, result_meta) + message = result + meta = result_meta + done = true + end, + }) + + wait_for(function() + return done + end) + assert.are.same('ok', message) + assert.is_table(meta) + ---@cast meta table + assert.are.same('gemma4:e2b 32k', meta.requested_model) + assert.are.same('gemma4:e2b', meta.used_model) + end) + + it('streams response chunks', function() + local done = false + local chunks = {} + local message = nil + + ai_provider.chat('ollama', { + model = 'gemma4:e2b', + prompt = 'Reply with exactly: ok', + max_tokens = 4, + timeout = 10000, + on_chunk = function(chunk) + table.insert(chunks, chunk) + end, + callback = function(result) + message = result + done = true + end, + }) + + wait_for(function() + return done + end) + assert.are.same('ok', message) + assert.is_true(#chunks > 0) + assert.are.same('ok', table.concat(chunks, '')) + end) +end) diff --git a/local/ai-provider/tests/minimal_init.lua b/local/ai-provider/tests/minimal_init.lua new file mode 100644 index 0000000..6fee827 --- /dev/null +++ b/local/ai-provider/tests/minimal_init.lua @@ -0,0 +1,4 @@ +vim.opt.runtimepath:append '/home/misti/.local/share/nvim/lazy/plenary.nvim' +vim.opt.runtimepath:append '/home/misti/.config/nvim/local/ai-provider' + +vim.cmd 'runtime plugin/plenary.vim' diff --git a/lua/plugins/ai-commit.lua b/lua/plugins/ai-commit.lua index a1effa9..4eae0d7 100644 --- a/lua/plugins/ai-commit.lua +++ b/lua/plugins/ai-commit.lua @@ -1,11 +1,11 @@ ---- AI-powered Git commit message generator using GitHub Copilot +--- AI-powered Git commit message generator using Ollama, GitHub Copilot, or OpenRouter -- Configuration - Change these to customize behavior local CONFIG = { - -- Provider can be 'copilot' or 'openrouter' + -- Fallback provider when Ollama is not reachable. Can be 'copilot' or 'openrouter' provider = 'copilot', - -- Use :AICommitModels to select a model + -- Use :AIProvider to configure the local provider; this is only the remote fallback model. model = nil, -- nil = auto (use Copilot's default) model_name = nil, -- Friendly display name for selected model @@ -91,6 +91,8 @@ local state = { in_flight_buffers = {}, } +local clean_message + -- Get preferences file path local function get_preferences_file() if vim and vim.fn then @@ -194,7 +196,9 @@ end ---@return string local function get_selected_model_name() local provider_prefix = '' - if CONFIG.provider == 'openrouter' then + if CONFIG.provider == 'ollama' then + provider_prefix = '[Ollama] ' + elseif CONFIG.provider == 'openrouter' then provider_prefix = '[OpenRouter] ' elseif CONFIG.provider == 'copilot' and CONFIG.model then provider_prefix = '[Copilot] ' @@ -213,6 +217,55 @@ local function get_selected_model_name() return 'Unknown' end +--- Completion for local Ollama +---@param full_prompt string +---@param callback function(string|nil, table|nil) +---@param status_callback function(string)|nil +---@param request_context table|nil +local function complete_ollama(full_prompt, callback, status_callback, request_context) + local log = setup_logger() + local ai_provider = require 'ai-provider' + local model = ai_provider.get_selected_model 'ollama' + + if not model then + log.error 'Ollama is reachable but no model is selected' + vim.notify('No Ollama model selected. Run :AIProvider ollama model first.', vim.log.levels.ERROR) + callback(nil, nil) + return + end + + if status_callback then + status_callback('Waiting for response from ' .. model) + end + + ai_provider.chat('ollama', { + model = model, + prompt = full_prompt, + stream = true, + max_tokens = CONFIG.max_tokens, + is_cancelled = request_context and request_context.is_cancelled, + register_http_job = request_context and request_context.register_http_job, + callback = function(message, meta) + if request_context and request_context.is_cancelled and request_context.is_cancelled() then + return + end + + if not message then + log.error 'Ollama chat request failed' + vim.notify('Ollama request failed. See ai-commit logs.', vim.log.levels.ERROR) + callback(nil, nil) + return + end + + local cleaned = clean_message(message) + local used_model = meta and meta.used_model or model + local elapsed_ms = meta and meta.elapsed_ms or 0 + log.info(string.format('Generated commit message with Ollama (took %.0fms, model=%s)', elapsed_ms, used_model)) + callback(cleaned, { requested_model = model, used_model = used_model }) + end, + }) +end + -- Try to get headers from Avante's copilot provider (so we don't maintain them) local function get_copilot_headers() local ok, avante_copilot = pcall(require, 'avante.providers.copilot') @@ -505,7 +558,7 @@ end --- Clean up AI-generated message ---@param message string ---@return string -local function clean_message(message) +clean_message = function(message) local result = message:gsub('^%s*```.-\n', ''):gsub('\n```%s*$', ''):gsub('^%s+', ''):gsub('%s+$', '') return result end @@ -745,16 +798,34 @@ end ---@param request_context table|nil local function generate_commit_message_async(branch, recent_commits, diff, callback, status_callback, request_context) local log = setup_logger() - log.info(string.format('Starting commit message generation using %s provider', CONFIG.provider)) + log.info 'Starting commit message generation' local full_prompt = string.format(COMMIT_PROMPT_TEMPLATE, branch, recent_commits, diff) log.debug(string.format('Prompt built (branch=%s, commits=%d chars, diff=%d chars)', branch, #recent_commits, #diff)) - if CONFIG.provider == 'openrouter' then - complete_openrouter(full_prompt, callback, status_callback, request_context) - else - complete_copilot(full_prompt, callback, status_callback, request_context) + local ai_provider = require 'ai-provider' + if status_callback then + status_callback 'Checking Ollama' end + + ai_provider.check('ollama', function(working) + if request_context and request_context.is_cancelled and request_context.is_cancelled() then + return + end + + if working then + log.info 'Routing commit message generation to Ollama' + complete_ollama(full_prompt, callback, status_callback, request_context) + return + end + + log.info(string.format('Ollama unavailable, falling back to %s provider', CONFIG.provider)) + if CONFIG.provider == 'openrouter' then + complete_openrouter(full_prompt, callback, status_callback, request_context) + else + complete_copilot(full_prompt, callback, status_callback, request_context) + end + end) end local function get_copilot_models() @@ -1278,8 +1349,5 @@ return { desc = 'Generate AI commit message', }) - vim.api.nvim_create_user_command('AICommitModels', select_model, { - desc = 'Select Copilot model for commit messages', - }) end, } diff --git a/lua/plugins/ai-provider.lua b/lua/plugins/ai-provider.lua new file mode 100644 index 0000000..4b3825e --- /dev/null +++ b/lua/plugins/ai-provider.lua @@ -0,0 +1,62 @@ +--- @type LazySpec +return { + dir = vim.fn.stdpath 'config' .. '/local/ai-provider', + name = 'ai-provider', + event = 'VeryLazy', + cmd = { 'AIProvider' }, + dependencies = { 'nvim-lua/plenary.nvim' }, + opts = { + default_provider = 'ollama', + providers = { + ollama = { + default_model = 'gemma4:e2b', + context_size = 1024 * 8, + keep_alive = '4h', + models = { + ['gemma4:e2b 32k'] = { + model = 'gemma4:e2b', + context_size = 1024 * 32, + }, + ['gemma4:e2b 64k'] = { + model = 'gemma4:e2b', + context_size = 1024 * 64, + }, + }, + }, + }, + }, + keys = { + { + 'pm', + function() + require('ai-provider').select_model() + end, + desc = 'AI [P]rovider [M]odel', + }, + { + 'pM', + function() + require('ai-provider').select_model 'ollama' + end, + desc = 'AI [P]rovider Ollama [M]odel', + }, + { + 'pd', + 'AIProvider default', + desc = 'AI [P]rovider [D]efault', + }, + { + 'pc', + 'AIProvider ollama check', + desc = 'AI [P]rovider [C]heck Ollama', + }, + { + 'pl', + 'AIProvider models', + desc = 'AI [P]rovider [L]ist Models', + }, + }, + config = function(_, opts) + require('ai-provider').setup(opts) + end, +}