feat: add Nix as Rag service runner option (#1480)
* feat: add nix as option for RAG runner * fix: remove default embedding model * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * stylua format --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Jorge Luis Suarez
committed by
GitHub
GitHub
parent
e1d2d825e2
commit
e408b820c8
@@ -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
|
||||
|
||||
@@ -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://(.*)$")
|
||||
|
||||
23
py/rag-service/run.sh
Executable file
23
py/rag-service/run.sh
Executable file
@@ -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
|
||||
46
py/rag-service/shell.nix
Executable file
46
py/rag-service/shell.nix
Executable file
@@ -0,0 +1,46 @@
|
||||
{ pkgs ? import <nixpkgs> {} }:
|
||||
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"
|
||||
'';
|
||||
}
|
||||
Reference in New Issue
Block a user