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>

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

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://(.*)$")

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

py/rag-service/shell.nix

@@ -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"
+    '';
+}