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://(.*)$")
|
||||
|
||||
Reference in New Issue
Block a user