From 251b58a379e8ede24b8f4c3b3565eea40deb5fc5 Mon Sep 17 00:00:00 2001
From: Michael Yuan <michael@secondstate.io>
Date: Fri, 12 Apr 2024 18:22:33 +0000
Subject: [PATCH] Support multiple qdrant servers

Signed-off-by: Michael Yuan <michael@secondstate.io>
---
 config.json |  1 +
 install.sh  | 46 ++++++++++++++++++++++++++++------------------
 start.sh    | 35 +++++++++++++++++++++--------------
 3 files changed, 50 insertions(+), 32 deletions(-)

diff --git a/config.json b/config.json
index be52c58..bdd9ef0 100644
--- a/config.json
+++ b/config.json
@@ -9,6 +9,7 @@
     "embedding": "https://huggingface.co/second-state/All-MiniLM-L6-v2-Embedding-GGUF/resolve/main/all-MiniLM-L6-v2-ggml-model-f16.gguf",
     "embedding_ctx_size": "384",
     "snapshot": "https://huggingface.co/datasets/gaianet/paris/resolve/main/paris_384_all-minilm-l6-v2_f16.snapshot",
+    "embedding_collection_name": "default",
     "rag_prompt": "Use the following pieces of context to answer the user's question.\nIf you don't know the answer, just say that you don't know, don't try to make up an answer.\n----------------\n",
     "domain": "gaianet.xyz",
     "llamaedge_port": "8080"
diff --git a/install.sh b/install.sh
index a5050ad..607cc3a 100755
--- a/install.sh
+++ b/install.sh
@@ -240,11 +240,13 @@ fi
 # 10. recover from the given qdrant collection snapshot =======================
 printf "[+] Initializing the Qdrant server ...\n\n"
 
-# check 6333 port is in use or not
+qdrant_pid=0
+qdrant_already_running=false
 if [ "$(uname)" == "Darwin" ] || [ "$(expr substr $(uname -s) 1 5)" == "Linux" ]; then
     if lsof -Pi :6333 -sTCP:LISTEN -t >/dev/null ; then
-        printf "It appears that the GaiaNet node is running. Please stop it first.\n\n"
-        exit 1
+        # printf "It appears that the GaiaNet node is running. Please stop it first.\n\n"
+        # exit 1
+	qdrant_already_running=true
     fi
 elif [ "$(expr substr $(uname -s) 1 10)" == "MINGW32_NT" ]; then
     printf "For Windows users, please run this script in WSL.\n"
@@ -254,15 +256,18 @@ else
     exit 1
 fi
 
-# start qdrant
-cd $gaianet_base_dir/qdrant
-nohup $gaianet_base_dir/bin/qdrant > $log_dir/init-qdrant.log 2>&1 &
-sleep 15
-qdrant_pid=$!
+if [ "$qdrant_already_running" = false ]; then
+    # start qdrant
+    cd $gaianet_base_dir/qdrant
+    nohup $gaianet_base_dir/bin/qdrant > $log_dir/init-qdrant.log 2>&1 &
+    sleep 15
+    qdrant_pid=$!
+fi
 
 cd $gaianet_base_dir
 url_snapshot=$(awk -F'"' '/"snapshot":/ {print $4}' config.json)
 url_document=$(awk -F'"' '/"document":/ {print $4}' config.json)
+embedding_collection_name=$(awk -F'"' '/"embedding_collection_name":/ {print $4}' config.json)
 
 if [ -n "$url_snapshot" ]; then
     # 10.1 recover from the given qdrant collection snapshot
@@ -271,18 +276,18 @@ if [ -n "$url_snapshot" ]; then
     curl --progress-bar -L $url_snapshot -o default.snapshot
 
     cd $gaianet_base_dir
-    # remove the 'default' collection if it exists
-    del_response=$(curl -s -X DELETE http://localhost:6333/collections/default \
+    # remove the collection if it exists
+    del_response=$(curl -s -X DELETE http://localhost:6333/collections/$embedding_collection_name \
         -H "Content-Type: application/json")
     status=$(echo "$del_response" | grep -o '"status":"[^"]*"' | cut -d':' -f2 | tr -d '"')
     if [ "$status" != "ok" ]; then
-        printf "    Failed to remove the 'default' collection. $del_response\n\n"
+        printf "    Failed to remove the $embedding_collection_name collection. $del_response\n\n"
         kill $qdrant_pid
         exit 1
     fi
 
     # Import the default.snapshot file
-    response=$(curl -s -X POST 'http://localhost:6333/collections/default/snapshots/upload?priority=snapshot' \
+    response=$(curl -s -X POST http://localhost:6333/collections/$embedding_collection_name/snapshots/upload?priority=snapshot \
         -H 'Content-Type:multipart/form-data' \
         -F 'snapshot=@default.snapshot')
     sleep 5
@@ -301,14 +306,14 @@ elif [ -n "$url_document" ]; then
 
     printf "[+] Creating a Qdrant collection from the given document ...\n\n"
 
-    # Remove the 'default' collection if it exists
-    printf "    * Removing 'default' collection if it exists ...\n\n"
+    # Remove the collection if it exists
+    printf "    * Removing collection if it exists ...\n\n"
     # remove the 'default' collection if it exists
-    del_response=$(curl -s -X DELETE http://localhost:6333/collections/default \
+    del_response=$(curl -s -X DELETE http://localhost:6333/collections/$embedding_collection_name \
         -H "Content-Type: application/json")
     status=$(echo "$del_response" | grep -o '"status":"[^"]*"' | cut -d':' -f2 | tr -d '"')
     if [ "$status" != "ok" ]; then
-        printf "    Failed to remove the 'default' collection. $del_response\n\n"
+        printf "    Failed to remove the collection. $del_response\n\n"
         kill $qdrant_pid
         exit 1
     fi
@@ -337,6 +342,8 @@ elif [ -n "$url_document" ]; then
     embedding_model_stem=$(basename "$embedding_model_name" .gguf)
     # parse context size for embedding model
     embedding_ctx_size=$(awk -F'"' '/"embedding_ctx_size":/ {print $4}' config.json)
+    # parse cli options for embedding vector collection name
+    embedding_collection_name=$(awk -F'"' '/"embedding_collection_name":/ {print $4}' config.json)
     # parse port for LlamaEdge API Server
     llamaedge_port=$(awk -F'"' '/"llamaedge_port":/ {print $4}' config.json)
 
@@ -370,6 +377,7 @@ elif [ -n "$url_document" ]; then
     rag-api-server.wasm -p $prompt_type \
     --model-name $chat_model_stem,$embedding_model_stem \
     --ctx-size $chat_ctx_size,$embedding_ctx_size \
+    --qdrant-collection-name $embedding_collection_name \
     --web-ui ./dashboard \
     --socket-addr 0.0.0.0:$llamaedge_port \
     --log-prompts \
@@ -424,8 +432,10 @@ else
 fi
 printf "\n"
 
-# stop qdrant
-kill $qdrant_pid
+if [ "$qdrant_already_running" = false ]; then
+    # stop qdrant
+    kill $qdrant_pid
+fi
 
 # ======================================================================================
 
diff --git a/start.sh b/start.sh
index 4c68a7f..efc364f 100755
--- a/start.sh
+++ b/start.sh
@@ -44,11 +44,13 @@ fi
 # 1. start a Qdrant instance
 printf "[+] Starting Qdrant instance ...\n"
 
+qdrant_already_running=false
 if [ "$(uname)" == "Darwin" ] || [ "$(expr substr $(uname -s) 1 5)" == "Linux" ]; then
     if lsof -Pi :6333 -sTCP:LISTEN -t >/dev/null ; then
-        printf "    Port 6333 is in use. Stopping the process on 6333 ...\n\n"
-        pid=$(lsof -t -i:6333)
-        kill -9 $pid
+        # printf "    Port 6333 is in use. Stopping the process on 6333 ...\n\n"
+        # pid=$(lsof -t -i:6333)
+        # kill -9 $pid
+        qdrant_already_running=true
     fi
 elif [ "$(expr substr $(uname -s) 1 10)" == "MINGW32_NT" ]; then
     printf "For Windows users, please run this script in WSL.\n"
@@ -58,17 +60,19 @@ else
     exit 1
 fi
 
-qdrant_executable="$gaianet_base_dir/bin/qdrant"
-if [ -f "$qdrant_executable" ]; then
-    cd $gaianet_base_dir/qdrant
-    nohup $qdrant_executable > $log_dir/start-qdrant.log 2>&1 &
-    sleep 2
-    qdrant_pid=$!
-    echo $qdrant_pid > $gaianet_base_dir/qdrant.pid
-    printf "\n    Qdrant instance started with pid: $qdrant_pid\n\n"
-else
-    printf "Qdrant binary not found at $qdrant_executable\n\n"
-    exit 1
+if [ "$qdrant_already_running" = false ]; then
+    qdrant_executable="$gaianet_base_dir/bin/qdrant"
+    if [ -f "$qdrant_executable" ]; then
+        cd $gaianet_base_dir/qdrant
+        nohup $qdrant_executable > $log_dir/start-qdrant.log 2>&1 &
+        sleep 2
+        qdrant_pid=$!
+        echo $qdrant_pid > $gaianet_base_dir/qdrant.pid
+        printf "\n    Qdrant instance started with pid: $qdrant_pid\n\n"
+    else
+        printf "Qdrant binary not found at $qdrant_executable\n\n"
+        exit 1
+    fi
 fi
 
 # 2. start a LlamaEdge instance
@@ -94,6 +98,8 @@ rag_prompt=$(awk -F'"' '/"rag_prompt":/ {print $4}' config.json)
 reverse_prompt=$(awk -F'"' '/"reverse_prompt":/ {print $4}' config.json)
 # parse cli options for embedding model
 url_embedding_model=$(awk -F'"' '/"embedding":/ {print $4}' config.json)
+# parse cli options for embedding vector collection name
+embedding_collection_name=$(awk -F'"' '/"embedding_collection_name":/ {print $4}' config.json)
 # gguf filename
 embedding_model_name=$(basename $url_embedding_model)
 # stem part of the filename
@@ -133,6 +139,7 @@ cmd=(wasmedge --dir .:./dashboard \
   --model-name $chat_model_stem,$embedding_model_stem \
   --ctx-size $chat_ctx_size,$embedding_ctx_size \
   --prompt-template $prompt_type \
+  --qdrant-collection-name $embedding_collection_name \
   --web-ui ./ \
   --socket-addr 0.0.0.0:$llamaedge_port \
   --log-prompts \