diff --git a/lua/avante/config.lua b/lua/avante/config.lua index 34bfa94..68515cf 100644 --- a/lua/avante/config.lua +++ b/lua/avante/config.lua @@ -35,6 +35,7 @@ M._defaults = { tokenizer = "tiktoken", rag_service = { enabled = false, -- Enables the rag service, requires OPENAI_API_KEY to be set + runner = "docker", -- The runner for the rag service, (can use docker, or nix) provider = "openai", -- The provider to use for RAG service. eg: openai or ollama llm_model = "", -- The LLM model to use for RAG service embed_model = "", -- The embedding model to use for RAG service diff --git a/lua/avante/rag_service.lua b/lua/avante/rag_service.lua index ae601fc..e0fadac 100644 --- a/lua/avante/rag_service.lua +++ b/lua/avante/rag_service.lua @@ -1,11 +1,12 @@ local curl = require("plenary.curl") local Path = require("plenary.path") -local Utils = require("avante.utils") local Config = require("avante.config") +local Utils = require("avante.utils") local M = {} local container_name = "avante-rag-service" +local service_path = "/tmp/" .. container_name function M.get_rag_service_image() return "quay.io/yetoneful/avante-rag-service:0.0.6" end @@ -30,6 +31,8 @@ function M.get_current_image() return image end +function M.get_rag_service_runner() return (Config.rag_service and Config.rag_service.runner) or "docker" end + ---@param cb fun() function M.launch_rag_service(cb) local openai_api_key = os.getenv("OPENAI_API_KEY") @@ -40,69 +43,130 @@ function M.launch_rag_service(cb) end end local port = M.get_rag_service_port() - local image = M.get_rag_service_image() - local data_path = M.get_data_path() - local cmd = string.format("docker ps -a | grep '%s'", container_name) - local result = vim.fn.system(cmd) - if result ~= "" then - Utils.debug(string.format("container %s already running", container_name)) - local current_image = M.get_current_image() - if current_image == image then + + if M.get_rag_service_runner() == "docker" then + local image = M.get_rag_service_image() + local data_path = M.get_data_path() + local cmd = string.format("docker ps -a | grep '%s'", container_name) + local result = vim.fn.system(cmd) + if result ~= "" then + Utils.debug(string.format("container %s already running", container_name)) + local current_image = M.get_current_image() + if current_image == image then + cb() + return + end + Utils.debug( + string.format( + "container %s is running with different image: %s != %s, stopping...", + container_name, + current_image, + image + ) + ) + M.stop_rag_service() + else + Utils.debug(string.format("container %s not found, starting...", container_name)) + end + local cmd_ = string.format( + "docker run -d -p %d:8000 --name %s -v %s:/data -v /:/host -e DATA_DIR=/data -e RAG_PROVIDER=%s -e %s_API_KEY=%s -e %s_API_BASE=%s -e RAG_LLM_MODEL=%s -e RAG_EMBED_MODEL=%s %s", + port, + container_name, + data_path, + Config.rag_service.provider, + Config.rag_service.provider:upper(), + openai_api_key, + Config.rag_service.provider:upper(), + Config.rag_service.endpoint, + Config.rag_service.llm_model, + Config.rag_service.embed_model, + image + ) + vim.fn.jobstart(cmd_, { + detach = true, + on_exit = function(_, exit_code) + if exit_code ~= 0 then + Utils.error(string.format("container %s failed to start, exit code: %d", container_name, exit_code)) + else + Utils.debug(string.format("container %s started", container_name)) + cb() + end + end, + }) + elseif M.get_rag_service_runner() == "nix" then + -- Check if service is already running + local check_cmd = string.format("pgrep -f '%s'", service_path) + local check_result = vim.fn.system(check_cmd) + if check_result ~= "" then + Utils.debug(string.format("RAG service already running at %s", service_path)) cb() return end - Utils.debug( - string.format( - "container %s is running with different image: %s != %s, stopping...", - container_name, - current_image, - image - ) + + local dirname = + Utils.trim(string.sub(debug.getinfo(1).source, 2, #"/lua/avante/rag_service.lua" * -1), { suffix = "/" }) + local rag_service_dir = dirname .. "/py/rag-service" + + Utils.debug(string.format("launching %s with nix...", container_name)) + + local cmd = string.format( + "cd %s && PORT=%d DATA_DIR=%s RAG_PROVIDER=%s %s_API_KEY=%s %s_API_BASE=%s RAG_LLM_MODEL=%s RAG_EMBED_MODEL=%s sh run.sh %s", + rag_service_dir, + port, + service_path, + Config.rag_service.provider, + Config.rag_service.provider:upper(), + openai_api_key, + Config.rag_service.provider:upper(), + Config.rag_service.endpoint, + Config.rag_service.llm_model, + Config.rag_service.embed_model, + service_path ) - M.stop_rag_service() - else - Utils.debug(string.format("container %s not found, starting...", container_name)) + + vim.fn.jobstart(cmd, { + detach = true, + on_exit = function(_, exit_code) + if exit_code ~= 0 then + Utils.error(string.format("service %s failed to start, exit code: %d", container_name, exit_code)) + else + Utils.debug(string.format("service %s started", container_name)) + cb() + end + end, + }) end - local cmd_ = string.format( - "docker run -d -p %d:8000 --name %s -v %s:/data -v /:/host -e DATA_DIR=/data -e RAG_PROVIDER=%s -e %s_API_KEY=%s -e %s_API_BASE=%s -e RAG_LLM_MODEL=%s -e RAG_EMBED_MODEL=%s %s", - port, - container_name, - data_path, - Config.rag_service.provider, - Config.rag_service.provider:upper(), - openai_api_key, - Config.rag_service.provider:upper(), - Config.rag_service.endpoint, - Config.rag_service.llm_model, - Config.rag_service.embed_model, - image - ) - vim.fn.jobstart(cmd_, { - detach = true, - on_exit = function(_, exit_code) - if exit_code ~= 0 then - Utils.error(string.format("container %s failed to start, exit code: %d", container_name, exit_code)) - else - Utils.debug(string.format("container %s started", container_name)) - cb() - end - end, - }) end function M.stop_rag_service() - local cmd = string.format("docker ps -a | grep '%s'", container_name) - local result = vim.fn.system(cmd) - if result ~= "" then vim.fn.system(string.format("docker rm -fv %s", container_name)) end + if M.get_rag_service_runner() == "docker" then + local cmd = string.format("docker ps -a | grep '%s'", container_name) + local result = vim.fn.system(cmd) + if result ~= "" then vim.fn.system(string.format("docker rm -fv %s", container_name)) end + else + local cmd = string.format("pgrep -f '%s' | xargs -r kill -9", service_path) + vim.fn.system(cmd) + Utils.debug(string.format("Attempted to kill processes related to %s", service_path)) + end end function M.get_rag_service_status() - local cmd = string.format("docker ps -a | grep '%s'", container_name) - local result = vim.fn.system(cmd) - if result == "" then - return "running" - else - return "stopped" + if M.get_rag_service_runner() == "docker" then + local cmd = string.format("docker ps -a | grep '%s'", container_name) + local result = vim.fn.system(cmd) + if result == "" then + return "stopped" + else + return "running" + end + elseif M.get_rag_service_runner() == "nix" then + local cmd = string.format("pgrep -f '%s'", service_path) + local result = vim.fn.system(cmd) + if result == "" then + return "stopped" + else + return "running" + end end end @@ -113,6 +177,8 @@ function M.get_scheme(uri) end function M.to_container_uri(uri) + local runner = M.get_rag_service_runner() + if runner == "nix" then return uri end local scheme = M.get_scheme(uri) if scheme == "file" then local path = uri:match("^file://(.*)$") diff --git a/py/rag-service/run.sh b/py/rag-service/run.sh new file mode 100755 index 0000000..0848e7e --- /dev/null +++ b/py/rag-service/run.sh @@ -0,0 +1,23 @@ +#!/bin/bash + +# Set the target directory (use the first argument or default to a temporary directory) +TARGET_DIR=$1 +if [ -z "$TARGET_DIR" ]; then + TARGET_DIR="/tmp/avante-rag-service" +fi +# Create the target directory if it doesn't exist +mkdir -p "$TARGET_DIR" + +# Copy the required files to the target directory +cp -r src/ "$TARGET_DIR" +cp requirements.txt "$TARGET_DIR" +cp shell.nix "$TARGET_DIR" + +echo "Files have been copied to $TARGET_DIR" + +# Change to the target directory +cd "$TARGET_DIR" + +# Run the RAG service using nix-shell +# The environment variables (PORT, DATA_DIR, OPENAI_API_KEY, OPENAI_BASE_URL) are passed from the parent process +nix-shell diff --git a/py/rag-service/shell.nix b/py/rag-service/shell.nix new file mode 100755 index 0000000..d193330 --- /dev/null +++ b/py/rag-service/shell.nix @@ -0,0 +1,46 @@ +{ pkgs ? import {} }: +let + logFile = "shell_log.txt"; + python = pkgs.python311; +in pkgs.mkShell { + packages = [ + python + pkgs.uv + pkgs.stdenv.cc.cc.lib + ]; + env = { + PYTHONUNBUFFERED = 1; + PYTHONDONTWRITEBYTECODE = 1; + LD_LIBRARY_PATH = "${pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH"; + }; + shellHook = '' + + # Start with a fresh log file + echo "=== avante.nvim RAG service setup log $(date '+%Y-%m-%d %H:%M:%S') ===" > "${logFile}" + + # Function to run commands and log their output + run_and_log() { + echo "$ $1" >> "${logFile}" + eval "$1" 2>&1 | tee -a "${logFile}" + echo "" >> "${logFile}" + } + + # Log environment info + run_and_log "echo 'Environment: $(uname -a)'" + run_and_log "echo 'Python version: $(python --version)'" + run_and_log "echo 'UV version: $(uv --version)'" + + + if [ ! -d ".venv" ]; then + run_and_log "uv venv" + else + echo "Using existing virtual environment" tee -a "${logFile}" + fi + + run_and_log source ".venv/bin/activate" + + run_and_log "uv pip install -r requirements.txt" + + run_and_log "uv run fastapi run src/main.py --port $PORT --workers 3" + ''; +}