Repository: kaito-project/kaito

Issue:
Why do we need to label nodes when we already provide the node name in preferredNode?
**Is your feature request related to a problem? Please describe.**
This feature request addresses a user experience issue. While I understand this relates to Kubernetes' native orchestration, the current behavior creates a confusing user experience.

The current implementation also leads to poor error handling: if a customer mislabels a node, the controller still creates a nodeclaim but eventually times out because no controller handles the nodeclaim. This timeout error is disconnected from the actual root cause (the labeling error), making debugging difficult.
**Describe the solution you'd like**
1. Have the workspace controller automatically label the node when a preferredNode is specified.
2. If option 1 is not feasible, at minimum return a "no node found" error immediately when a preferredNode is set but no matching labeled node exists.

**Describe alternatives you've considered**

**Additional context**

Discussion:
@chewong: @qfai would you be interested in raising a PR to fix this?@Jont828: /assign@Fei-Guo: I think the preferred nodes still need to be properly labeled. 
1) User can find the nodes used by a particular workspace using label selector easily. 
2) This label selector is used as node affinity label when creating inference deployment to ensure the workload will be run in preferred nodes.

If the mechanism of falling back to NAP is not desired, i.e., making the preferred nodes as "required nodes", we can either 
1) change the workspace webhook validation logic to verify the number of preferred nodes is equal to the count field and make sure all nodes exist and in running state. But this brings another problem such that if any of the preferred node is crashing or down during the runtime, the controller will be stuck. User needs to rely on other workflow to fix the broken nodes and update the workspace with "new" preferred nodes (complicating the day 2 operations a lot).

2) Add an annotation for specific user scenario to change the logic if the users know exactly what they are doing. e.g., if the special annotation is set, we add another check in validation webhook, if the node name is wrong or the node label is missing, fail the workspace creation.

@qfai, will option 2 work for you?

@qfai: I can see your point on the labelled node.
But what's the desired user experience if the user-managed node is crashed without NAP after creation? The workspace will fail and the user needs to bring the node back?
I think the problem is that the controller needs to know whether there is NAP or not.
I propose to add a new field in helm value like Compatitable_NAP: true. 
helm installed customer needs to set this value to false, while AKS will set it to true.
When it is set to true, it means it will fall back to NAP(original logic).
When it is set to false, the preferred node means required nodes and will have checks in the validation webhook. 

What do you think? Option 2 can solve the problem, I am ok with that. My proposal gives a better BYO node UX. The field can be used in other places as well.
@Fei-Guo: If this is set via helm, we will need to add a controller boot argument to specify the behavior or adding a feature gate instead of using an annotation. This also sounds reasonable to me. Thanks for the suggestion.@chewong: I think having the ability to turn off NAP is a great addition. This would also make instance type an optional parameter to Workspace resources when NAP is off, which fixes https://github.com/kaito-project/kaito/issues/1280.

@Jont828 could you update https://github.com/kaito-project/kaito/pull/1306 (or open a separate PR) to address the asks?

Solution patch:
diff --git a/.github/actions/build-image-action/action.yaml b/.github/actions/build-image-action/action.yaml
index 0ab6550..c073968 100644
--- a/.github/actions/build-image-action/action.yaml
+++ b/.github/actions/build-image-action/action.yaml
@@ -82,7 +82,9 @@ runs:
     - name: Install Hugging Face CLI
       run: |
         echo "Installing Hugging Face CLI..."
-        pip install huggingface-hub
+        if ! pip install huggingface-hub; then
+          pip install huggingface-hub --break-system-packages
+        fi
       shell: bash
 
     - name: Install ORAS CLI
diff --git a/.github/determine_models.py b/.github/determine_models.py
index 7c2bc08..db57b45 100644
--- a/.github/determine_models.py
+++ b/.github/determine_models.py
@@ -23,30 +23,30 @@ import yaml
 
 def read_yaml(file_path):
     try:
-        with open(file_path, 'r') as file:
+        with open(file_path) as file:
             data = yaml.safe_load(file)
             return data
-    except (IOError, yaml.YAMLError) as e:
+    except (OSError, yaml.YAMLError) as e:
         print(f"Error reading {file_path}: {e}")
         return None
 
 
-supp_models_yaml = 'presets/workspace/models/supported_models.yaml'
+supp_models_yaml = "presets/workspace/models/supported_models.yaml"
 YAML_PR = read_yaml(supp_models_yaml)
 # Format: {falcon-7b : {model_name:falcon-7b, type:text-generation, version: #, tag: #}}
-MODELS = {model['name']: model for model in YAML_PR['models']}
+MODELS = {model["name"]: model for model in YAML_PR["models"]}
 KAITO_REPO_URL = "https://github.com/kaito-project/kaito.git"
 GITREMOTE_TARGET = "_ciupstream"
 
 
 def set_multiline_output(name, value):
-    if not os.getenv('GITHUB_OUTPUT'):
-        print(f"Not in github env, skip writing to $GITHUB_OUTPUT .")
+    if not os.getenv("GITHUB_OUTPUT"):
+        print("Not in github env, skip writing to $GITHUB_OUTPUT .")
         return
 
-    with open(os.getenv('GITHUB_OUTPUT'), 'a') as fh:
+    with open(os.getenv("GITHUB_OUTPUT"), "a") as fh:
         delimiter = uuid.uuid1()
-        print(f'{name}<<{delimiter}', file=fh)
+        print(f"{name}<<{delimiter}", file=fh)
         print(value, file=fh)
         print(delimiter, file=fh)
 
@@ -60,13 +60,14 @@ def create_matrix(models_list):
 def run_command(command):
     """Execute a shell command and return the output."""
     try:
-        process = subprocess.Popen(command, stdout=subprocess.PIPE,
-                                   stderr=subprocess.PIPE, shell=True)
+        process = subprocess.Popen(
+            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
+        )
         output, error = process.communicate()
         if process.returncode != 0:
             print(f"Error: {error.decode('utf-8')}")
             return None
-        return output.decode('utf-8').strip()
+        return output.decode("utf-8").strip()
     except Exception as e:
         print(f"An error occurred: {e}")
         return None
@@ -74,33 +75,33 @@ def run_command(command):
 
 def get_yaml_from_branch(branch, file_path):
     """Read YAML from a branch"""
-    subprocess.run(['git', 'fetch', GITREMOTE_TARGET, branch], check=True)
-    subprocess.run(['git', 'checkout', f"{GITREMOTE_TARGET}/" + branch], check=True)
+    subprocess.run(["git", "fetch", GITREMOTE_TARGET, branch], check=True)
+    subprocess.run(["git", "checkout", f"{GITREMOTE_TARGET}/" + branch], check=True)
     content = read_yaml(file_path)
-    subprocess.run(['git', 'checkout', '-'], check=True)
+    subprocess.run(["git", "checkout", "-"], check=True)
     return content
 
 
 def detect_changes_in_yaml(yaml_main, yaml_pr):
     """Detecting relevant changes in support_models.yaml"""
-    yaml_main, yaml_pr = yaml_main['models'], yaml_pr['models']
+    yaml_main, yaml_pr = yaml_main["models"], yaml_pr["models"]
 
     models_to_build = []
     for model_pr in yaml_pr:
         # Searches for matching models
-        model_main = next((m for m in yaml_main if m['name'] == model_pr['name']), None)
+        model_main = next((m for m in yaml_main if m["name"] == model_pr["name"]), None)
         # New Model
         if not model_main:
-            print("New Model: ", model_pr['name'])
-            models_to_build.append(model_pr['name'])
+            print("New Model: ", model_pr["name"])
+            models_to_build.append(model_pr["name"])
         # Model Version Update
-        elif model_pr.get('version') != model_main.get('version'):
-            print("Updated Version of Model: ", model_pr['name'])
-            models_to_build.append(model_pr['name'])
+        elif model_pr.get("version") != model_main.get("version"):
+            print("Updated Version of Model: ", model_pr["name"])
+            models_to_build.append(model_pr["name"])
         # Model Tag Update
-        elif model_pr.get('tag') != model_main.get('tag'):
-            print("Update Tag of Model: ", model_pr['name'])
-            models_to_build.append(model_pr['name'])
+        elif model_pr.get("tag") != model_main.get("tag"):
+            print("Update Tag of Model: ", model_pr["name"])
+            models_to_build.append(model_pr["name"])
     return models_to_build
 
 
@@ -108,13 +109,15 @@ def models_to_build(files_changed):
     """Models to build based on changed files."""
     models, seen_model_types = set(), set()
     if supp_models_yaml in files_changed:
-        yaml_main = get_yaml_from_branch('main', supp_models_yaml)
+        yaml_main = get_yaml_from_branch("main", supp_models_yaml)
         models.update(detect_changes_in_yaml(yaml_main, YAML_PR))
     for model, model_info in MODELS.items():
-        if model_info["type"] not in seen_model_types:
-            if any(file.startswith(f'presets/workspace/inference/{model_info["type"]}') for file in files_changed):
-                models.add(model)
-                seen_model_types.add(model_info["type"])
+        if model_info["type"] not in seen_model_types and any(
+            file.startswith(f"presets/workspace/inference/{model_info['type']}")
+            for file in files_changed
+        ):
+            models.add(model)
+            seen_model_types.add(model_info["type"])
     return list(models)
 
 
@@ -128,7 +131,9 @@ def check_modified_models():
     run_command(f"git remote add {GITREMOTE_TARGET} {KAITO_REPO_URL}")
     run_command(f"git fetch {GITREMOTE_TARGET}")
 
-    files = run_command(f"git diff --name-only {GITREMOTE_TARGET}/main")  # Returns each file on newline
+    files = run_command(
+        f"git diff --name-only {GITREMOTE_TARGET}/main"
+    )  # Returns each file on newline
     files = files.split("\n")
     print("Files Changed: ", files)
 
@@ -140,28 +145,40 @@ def check_modified_models():
 
 
 def main():
-    force_run_all = os.environ.get("FORCE_RUN_ALL", "false")  # If not specified default to False
-    force_run_all_phi = os.environ.get("FORCE_RUN_ALL_PHI", "false")  # If not specified default to False
+    force_run_all = os.environ.get(
+        "FORCE_RUN_ALL", "false"
+    )  # If not specified default to False
+    force_run_all_phi = os.environ.get(
+        "FORCE_RUN_ALL_PHI", "false"
+    )  # If not specified default to False
 
     affected_models = []
     if force_run_all != "false":
-        affected_models = [model['name'] for model in YAML_PR['models'] if model['name']]
+        affected_models = [
+            model["name"] for model in YAML_PR["models"] if model["name"]
+        ]
     elif force_run_all_phi != "false":
-        affected_models = [model['name'] for model in YAML_PR['models'] if 'phi-3' in model['name']]
+        affected_models = [
+            model["name"] for model in YAML_PR["models"] if "phi-3" in model["name"]
+        ]
     else:
         # Logic to determine affected models
         # Example: affected_models = ['model1', 'model2', 'model3']
         affected_models = check_modified_models()
 
-    download_at_runtime = {model['name'] for model in YAML_PR['models'] if model.get('downloadAtRuntime')}
-    affected_models = [model for model in affected_models if model not in download_at_runtime]
+    download_at_runtime = {
+        model["name"] for model in YAML_PR["models"] if model.get("downloadAtRuntime")
+    }
+    affected_models = [
+        model for model in affected_models if model not in download_at_runtime
+    ]
 
     # Convert the list of models into JSON matrix format
     matrix = create_matrix(affected_models)
     print("Create Matrix: ", matrix)
 
     # Set the matrix as an output for the GitHub Actions workflow
-    set_multiline_output('matrix', matrix)
+    set_multiline_output("matrix", matrix)
 
 
 if __name__ == "__main__":
diff --git a/.github/e2e-preset-configs.json b/.github/e2e-preset-configs.json
index fe6cff4..12a2c40 100644
--- a/.github/e2e-preset-configs.json
+++ b/.github/e2e-preset-configs.json
@@ -4,7 +4,7 @@
       {
         "name": "falcon-7b",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -23,7 +23,7 @@
       {
         "name": "falcon-7b-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -80,7 +80,7 @@
       {
         "name": "mistral-7b",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -99,7 +99,7 @@
       {
         "name": "mistral-7b-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -118,7 +118,7 @@
       {
         "name": "phi-2",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 50,
         "OSS": true,
         "loads_adapter": false,
@@ -129,7 +129,7 @@
             "gpu_count": 1
           },
           "vllm": {
-            "command": "python3 /workspace/vllm/inference_api.py --kaito-max-probe-steps 6 --dtype float16",
+            "command": "VLLM_USE_V1=0 python3 /workspace/vllm/inference_api.py --kaito-max-probe-steps 6 --dtype float16 --kaito-config-file /mnt/config/inference_config.yaml",
             "gpu_count": 1
           }
         }
@@ -137,7 +137,7 @@
       {
         "name": "phi-3-mini-4k-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 50,
         "OSS": true,
         "loads_adapter": false,
@@ -156,7 +156,7 @@
       {
         "name": "phi-3-mini-128k-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 50,
         "OSS": true,
         "loads_adapter": false,
@@ -175,7 +175,7 @@
       {
         "name": "phi-3-medium-4k-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC12s_v3",
+        "node-vm-size": "Standard_NV72ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -194,7 +194,7 @@
       {
         "name": "phi-3-medium-128k-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC12s_v3",
+        "node-vm-size": "Standard_NV72ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -232,7 +232,7 @@
       {
         "name": "phi-4-mini-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -252,7 +252,7 @@
         "name": "qwen2.5-coder-7b-instruct",
         "workload": "qwen2-5-coder-7b-instruct",
         "node-count": 1,
-        "node-vm-size": "Standard_NC12s_v3",
+        "node-vm-size": "Standard_NV72ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -310,7 +310,7 @@
       {
         "name": "deepseek-r1-distill-llama-8b",
         "node-count": 1,
-        "node-vm-size": "Standard_NC12s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
@@ -318,18 +318,18 @@
         "runtimes": {
           "hf": {
             "command": "accelerate launch --num_processes 1 --num_machines 1 --machine_rank 0 --gpu_ids all /workspace/tfs/inference_api.py --pipeline text-generation --torch_dtype auto --trust_remote_code",
-            "gpu_count": 2
+            "gpu_count": 1
           },
           "vllm": {
-            "command": "python3 /workspace/vllm/inference_api.py --kaito-config-file /mnt/config/inference_config.yaml --tensor-parallel-size 2 --max-model-len 32768",
-            "gpu_count": 2
+            "command": "python3 /workspace/vllm/inference_api.py --kaito-config-file /mnt/config/inference_config.yaml",
+            "gpu_count": 1
           }
         }
       },
       {
         "name": "vllm-distributed",
         "node-count": 2,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "resource-type": "statefulset",
         "node-osdisk-size": 100,
         "OSS": false,
@@ -339,7 +339,7 @@
       {
         "name": "tuning",
         "node-count": 1,
-        "node-vm-size": "Standard_NC6s_v3",
+        "node-vm-size": "Standard_NV36ads_A10_v5",
         "node-osdisk-size": 100,
         "OSS": true,
         "loads_adapter": false,
diff --git a/.github/workflows/aikit.yaml b/.github/workflows/aikit.yaml
new file mode 100644
index 0000000..85e1020
--- /dev/null
+++ b/.github/workflows/aikit.yaml
@@ -0,0 +1,207 @@
+name: AIKit Integration Test
+
+on:
+  workflow_dispatch:
+    inputs:
+      aikit_image:
+        description: 'AIKit container image to test'
+        required: false
+        default: 'ghcr.io/kaito-project/aikit/llama3.2:1b'
+        type: string
+  pull_request:
+    branches:
+      - main
+
+env:
+  KIND_VERSION: "0.29.0"
+  KUBECTL_VERSION: "1.33.0"
+  AIKIT_IMAGE: ${{ github.event.inputs.aikit_image || 'ghcr.io/kaito-project/aikit/llama3.2:1b' }}
+  KAITO_NAMESPACE: "kaito-system"
+
+jobs:
+  aikit-integration-test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+
+    steps:
+      - name: Harden Runner
+        uses: step-security/harden-runner@ec9f2d5744a09debf3a187a3f4f675c53b671911 # v2.13.0
+        with:
+          egress-policy: audit
+
+      - name: Checkout
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+
+      - name: Install kind
+        run: |
+          curl -Lo ./kind "https://github.com/kubernetes-sigs/kind/releases/download/v${{ env.KIND_VERSION }}/kind-linux-amd64"
+          chmod +x ./kind
+          sudo mv ./kind /usr/local/bin/kind
+
+      - name: Install Helm
+        uses: azure/setup-helm@b9e51907a09c216f16ebe8536097933489208112 # v4.3.0
+
+      - name: Create kind cluster
+        run: |
+          cat << EOF > kind-config.yaml
+          kind: Cluster
+          apiVersion: kind.x-k8s.io/v1alpha4
+          name: kaito
+          nodes:
+          - role: control-plane
+            kubeadmConfigPatches:
+            - |
+              kind: InitConfiguration
+              nodeRegistration:
+                kubeletExtraArgs:
+                  node-labels: "apps=aikit-test"
+          EOF
+          kind create cluster --config kind-config.yaml
+          kubectl cluster-info
+
+      - name: Build workspace controller
+        run: |
+          # Use the docker-build-workspace target with local configuration
+          make docker-build-workspace
+
+          # Load into kind cluster
+          kind load docker-image local/workspace:test --name kaito
+        env:
+          REGISTRY: "local"
+          IMG_NAME: "workspace"
+          IMG_TAG: "test"
+          OUTPUT_TYPE: "type=docker"
+          ARCH: "amd64"
+
+      - name: Deploy KAITO workspace controller
+        run: |
+          # Update the Helm values with our local image
+          yq -i '(.image.repository) = "local/workspace"' ./charts/kaito/workspace/values.yaml
+          yq -i '(.image.tag) = "test"' ./charts/kaito/workspace/values.yaml
+
+          kubectl create namespace ${{ env.KAITO_NAMESPACE }}
+
+          # Install via Helm
+          # CSI Local Node cannot be configured to not deploy and it crashes on kind on github runners for unknown reasons
+          # so we need to filter it out its components
+          helm template kaito-workspace ./charts/kaito/workspace --namespace ${{ env.KAITO_NAMESPACE }} --include-crds --debug --wait | awk '
+            BEGIN { RS="---\n"; ORS="---\n"; csi=0 }
+            /kind: ServiceAccount/ && /name: csi-local-node/ { csi=1 }
+            /kind: StorageClass/ && /name: kaito-local-nvme-disk/ { csi=1 }
+            /kind: ClusterRole/ && /name: csi-local-node-role/ { csi=1 }
+            /kind: ClusterRoleBinding/ && /name: csi-local-node-binding/ { csi=1 }
+            /kind: DaemonSet/ && /name: csi-local-node/ { csi=1 }
+            !csi { print }
+            { csi=0 }
+          ' | kubectl apply -f -
+
+          # Wait for controller to be ready
+          kubectl wait --for=condition=available --timeout=300s deployment/kaito-workspace -n ${{ env.KAITO_NAMESPACE }}
+
+      - name: Create AIKit workspace
+        run: |
+          cat << EOF > aikit-workspace.yaml
+          apiVersion: kaito.sh/v1beta1
+          kind: Workspace
+          metadata:
+            name: workspace-aikit-test
+          resource:
+            labelSelector:
+              matchLabels:
+                apps: aikit-test
+            preferredNodes:
+              - kaito-control-plane
+          inference:
+            template:
+              spec:
+                containers:
+                  - name: aikit-llama
+                    image: ${{ env.AIKIT_IMAGE }}
+                    args:
+                      - "run"
+                      - "--address=:5000"
+          EOF
+          kubectl apply -f aikit-workspace.yaml
+
+      - name: Wait for inference to be ready
+        run: |
+          echo "Waiting for inference to be ready..."
+          kubectl wait --for='jsonpath={.status.conditions[?(@.type=="InferenceReady")].status}=True' workspace/workspace-aikit-test --timeout 300s
+
+      - name: Validate response content
+        run: |
+          # Set up port forwarding to the workspace service
+          echo "Setting up port forwarding..."
+          kubectl port-forward service/workspace-aikit-test 8080:80 &
+          PORT_FORWARD_PID=$!
+
+          # Wait for port-forward to be ready
+          sleep 5
+
+          # Test that we get expected model information
+          echo 'Validating model response contains expected data...'
+          MODELS_RESPONSE=$(curl -s http://localhost:8080/v1/models)
+          echo "Models response: $MODELS_RESPONSE"
+
+          # Extract the first model name from the response
+          MODEL_NAME=$(echo "$MODELS_RESPONSE" | jq -r '.data[0].id')
+          echo "Extracted model name: $MODEL_NAME"
+
+          # Check if response contains 'llama'
+          if echo "$MODELS_RESPONSE" | grep -i 'llama'; then
+            echo 'Model response validation passed!'
+          else
+            echo 'Model response validation failed - no llama model found'
+            kill $PORT_FORWARD_PID
+            exit 1
+          fi
+
+          echo 'Testing completion response...'
+          COMPLETION_RESPONSE=$(curl -s -X POST \
+            -H 'Content-Type: application/json' \
+            -d "{
+              \"model\": \"$MODEL_NAME\",
+              \"messages\": [{\"role\": \"user\", \"content\": \"hello\"}]
+            }" \
+            http://localhost:8080/v1/chat/completions)
+          echo "Completion response: $COMPLETION_RESPONSE"
+
+          # Check if we get some response with expected fields
+          if echo "$COMPLETION_RESPONSE" | jq -e '(type=="object" and has("choices") and has("usage"))'; then
+            echo 'Completion response validation passed!'
+          else
+            echo 'Completion response validation failed - missing expected fields'
+            kill $PORT_FORWARD_PID
+            exit 1
+          fi
+
+          # Clean up port forwarding
+          kill $PORT_FORWARD_PID
+          echo 'Validation completed successfully!'
+
+      - name: Debug info
+        if: always()
+        run: |
+          echo "=== Workspace Status ==="
+          kubectl get workspace workspace-aikit-test -o yaml
+
+          echo "=== Workspace events ==="
+          kubectl describe workspace workspace-aikit-test
+
+          echo "=== Pods ==="
+          kubectl get pods -A -o wide
+
+          echo "=== Services ==="
+          kubectl get services workspace-aikit-test -o yaml
+
+          echo "=== Pod events ==="
+          kubectl describe pods -l kaito.sh/workspace=workspace-aikit-test
+
+          echo "=== Inference Pod Logs ==="
+          kubectl logs -l kaito.sh/workspace=workspace-aikit-test --tail=-1
+
+          echo "=== Controller logs ==="
+          kubectl logs -l app.kubernetes.io/name=workspace -n ${{ env.KAITO_NAMESPACE }} --tail=-1
+
+          echo "=== Node information ==="
+          kubectl describe nodes
diff --git a/.github/workflows/codeql.yaml b/.github/workflows/codeql.yaml
index 348d1d4..447a064 100644
--- a/.github/workflows/codeql.yaml
+++ b/.github/workflows/codeql.yaml
@@ -29,12 +29,12 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout repository
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
        
-      - uses: actions/setup-go@v5.4.0
+      - uses: actions/setup-go@v5.5.0
         with:
           go-version-file: 'go.mod'
 
diff --git a/.github/workflows/create-release.yaml b/.github/workflows/create-release.yaml
index a756d44..5389df7 100644
--- a/.github/workflows/create-release.yaml
+++ b/.github/workflows/create-release.yaml
@@ -7,7 +7,7 @@ on:
         required: true
         type: string
       release_branch:
-        description: 'Release branch (e.g., official/v0.2.2)'
+        description: 'Release branch (e.g., release-0.6)'
         required: true
         type: string
 
@@ -31,16 +31,16 @@ jobs:
           disable-telemetry: true
 
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@0aaccfd150d50ccaeb58ebd88d36e91967a5f35b # v5.4.0
+        uses: actions/setup-go@8e57b58e57be52ac95949151e2777ffda8501267 # v5.4.0
         with:
           go-version: ${{ env.GO_VERSION  }}
 
       - name: Checkout the repository at the release branch
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
-          ref: ${{ github.event.inputs.release_branch }}
+          ref: ${{ github.event.inputs.release_version }}
 
       - name: Validate inputs
         run: |
@@ -49,8 +49,8 @@ jobs:
             echo "Error: Release version must follow format vX.Y.Z"
             exit 1
           fi
-          if [[ ! "${{ github.event.inputs.release_branch }}" =~ ^official/ ]]; then
-            echo "Error: Release branch must start with 'official/'"
+          if [[ ! "${{ github.event.inputs.release_branch }}" =~ ^release- ]]; then
+            echo "Error: Release branch must start with 'release-'"
             exit 1
           fi
 
diff --git a/.github/workflows/dependency-review.yaml b/.github/workflows/dependency-review.yaml
index be1102a..6d40337 100644
--- a/.github/workflows/dependency-review.yaml
+++ b/.github/workflows/dependency-review.yaml
@@ -24,6 +24,6 @@ jobs:
           disable-telemetry: true
 
       - name: 'Checkout Repository'
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
       - name: 'Dependency Review'
         uses: actions/dependency-review-action@da24556b548a50705dd671f47852072ea4c105d9 # v4.7.1
diff --git a/.github/workflows/deploy-docs.yaml b/.github/workflows/deploy-docs.yaml
index eccc499..211441f 100644
--- a/.github/workflows/deploy-docs.yaml
+++ b/.github/workflows/deploy-docs.yaml
@@ -28,7 +28,7 @@ jobs:
       run:
         working-directory: website
     steps:
-    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
 
     - name: Harden Runner
       uses: step-security/harden-runner@ec9f2d5744a09debf3a187a3f4f675c53b671911
@@ -48,7 +48,7 @@ jobs:
       run: echo "dir=$(yarn cache dir)" > $GITHUB_OUTPUT
 
     - name: Cache dependencies
-      uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
+      uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
       with:
         path: ${{ steps.yarn-cache.outputs.dir }}
         key: ${{ runner.os }}-website-${{ hashFiles('**/yarn.lock') }}
diff --git a/.github/workflows/e2e-preset-test.yaml b/.github/workflows/e2e-preset-test.yaml
index 55b92af..66f2c95 100644
--- a/.github/workflows/e2e-preset-test.yaml
+++ b/.github/workflows/e2e-preset-test.yaml
@@ -42,7 +42,7 @@ jobs:
       full_matrix: ${{ steps.images.outputs.full_matrix }}
     steps:
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
@@ -113,6 +113,12 @@ jobs:
                 map(. as $model | $configs[] | select(.name == $model.name) | $model + .)
             ')
 
+            # Filter out items with A100 VMs
+            # FIXME: add it back when A100 is available
+            COMBINED_MATRIX=$(echo "$COMBINED_MATRIX" | jq -c '
+                map(select(."node-vm-size" | contains("A100") | not))
+            ')
+
             echo "full_matrix=$COMBINED_MATRIX" >> $GITHUB_OUTPUT
       
       - name: Print Combined Matrix
@@ -128,7 +134,7 @@ jobs:
     environment: preset-env
     strategy:
       fail-fast: false
-      max-parallel: 10
+      max-parallel: 5
       matrix:
         # Ex matrix element:
         # {"name":"falcon-40b","type":"text-generation","version":"#",
@@ -137,7 +143,7 @@ jobs:
         model: ${{fromJson(needs.determine-models.outputs.full_matrix)}}
     steps:
       - name: Checkout
-        uses: actions/checkout@v4.2.2
+        uses: actions/checkout@v5.0.0
         with:
             submodules: true
             fetch-depth: 0
@@ -267,63 +273,88 @@ jobs:
       - name: Test inference endpoint
         run: |
             echo "Testing inference for ${{ matrix.model.name }}"
-            if [[ "${{ env.RUNTIME }}" == *"vllm"* ]]; then
-                kubectl exec ${{steps.workload.outputs.RESOURCE_TYPE}}/${{steps.workload.outputs.WORKLOAD_NAME}} -- \
-                curl -s -X POST \
-                -H "accept: application/json" \
-                -H "Content-Type: application/json" \
-                -d '{
-                    "model": "test",
-                    "messages": [
-                    {
-                        "role": "system",
-                        "content": "You are a helpful assistant."
-                    },
-                    {
-                        "role": "user",
-                        "content": "Hello!"
-                    }
-                    ]
-                    }' \
-                http://localhost:5000/v1/chat/completions
-            else
-                kubectl exec ${{steps.workload.outputs.RESOURCE_TYPE}}/${{steps.workload.outputs.WORKLOAD_NAME}} -- \
-                curl -s -X POST \
-                -H "accept: application/json" \
-                -H "Content-Type: application/json" \
-                -d '{
-                    "prompt":"Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
-                    "return_full_text": false,
-                    "clean_up_tokenization_spaces": false, 
-                    "prefix": null,
-                    "handle_long_generation": null,
-                    "generate_kwargs": {
-                            "max_length":200,
-                            "min_length":0,
-                            "do_sample":true,
-                            "early_stopping":false,
-                            "num_beams":1,
-                            "num_beam_groups":1,
-                            "diversity_penalty":0.0,
-                            "temperature":1.0,
-                            "top_k":10,
-                            "top_p":1,
-                            "typical_p":1,
-                            "repetition_penalty":1,
-                            "length_penalty":1,
-                            "no_repeat_ngram_size":0,
-                            "encoder_no_repeat_ngram_size":0,
-                            "bad_words_ids":null,
-                            "num_return_sequences":1,
-                            "output_scores":false,
-                            "return_dict_in_generate":false,
-                            "forced_bos_token_id":null,
-                            "forced_eos_token_id":null,
-                            "remove_invalid_values":null
+            
+            # Retry logic - attempt up to 3 times
+            MAX_ATTEMPTS=3
+            ATTEMPT=1
+            
+            while [ $ATTEMPT -le $MAX_ATTEMPTS ]; do
+                echo "Attempt $ATTEMPT of $MAX_ATTEMPTS"
+                
+                if [[ "${{ env.RUNTIME }}" == *"vllm"* ]]; then
+                    if kubectl exec ${{steps.workload.outputs.RESOURCE_TYPE}}/${{steps.workload.outputs.WORKLOAD_NAME}} -- \
+                    curl -s -X POST \
+                    -H "accept: application/json" \
+                    -H "Content-Type: application/json" \
+                    -d '{
+                        "model": "test",
+                        "messages": [
+                        {
+                            "role": "system",
+                            "content": "You are a helpful assistant."
+                        },
+                        {
+                            "role": "user",
+                            "content": "Hello!"
                         }
-                    }' \
-                http://localhost:5000/chat
-            fi
+                        ]
+                        }' \
+                    http://localhost:5000/v1/chat/completions; then
+                        echo "Inference test succeeded"
+                        break
+                    fi
+                else
+                    if kubectl exec ${{steps.workload.outputs.RESOURCE_TYPE}}/${{steps.workload.outputs.WORKLOAD_NAME}} -- \
+                    curl -s -X POST \
+                    -H "accept: application/json" \
+                    -H "Content-Type: application/json" \
+                    -d '{
+                        "prompt":"Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
+                        "return_full_text": false,
+                        "clean_up_tokenization_spaces": false, 
+                        "prefix": null,
+                        "handle_long_generation": null,
+                        "generate_kwargs": {
+                                "max_length":200,
+                                "min_length":0,
+                                "do_sample":true,
+                                "early_stopping":false,
+                                "num_beams":1,
+                                "num_beam_groups":1,
+                                "diversity_penalty":0.0,
+                                "temperature":1.0,
+                                "top_k":10,
+                                "top_p":1,
+                                "typical_p":1,
+                                "repetition_penalty":1,
+                                "length_penalty":1,
+                                "no_repeat_ngram_size":0,
+                                "encoder_no_repeat_ngram_size":0,
+                                "bad_words_ids":null,
+                                "num_return_sequences":1,
+                                "output_scores":false,
+                                "return_dict_in_generate":false,
+                                "forced_bos_token_id":null,
+                                "forced_eos_token_id":null,
+                                "remove_invalid_values":null
+                            }
+                        }' \
+                    http://localhost:5000/chat; then
+                        echo "Inference test succeeded"
+                        break
+                    fi
+                fi
+                
+                if [ $ATTEMPT -lt $MAX_ATTEMPTS ]; then
+                    echo "Test failed, retrying in 5 seconds..."
+                    sleep 5
+                else
+                    echo "Test failed after $MAX_ATTEMPTS attempts"
+                    exit 1
+                fi
+                
+                ATTEMPT=$((ATTEMPT + 1))
+            done
       
       - name: Cleanup
         if: always()
diff --git a/.github/workflows/e2e-preset-tuning-test.yaml b/.github/workflows/e2e-preset-tuning-test.yaml
index 796f3e1..9ecb4f3 100644
--- a/.github/workflows/e2e-preset-tuning-test.yaml
+++ b/.github/workflows/e2e-preset-tuning-test.yaml
@@ -21,7 +21,7 @@ jobs:
     environment: preset-env
     steps:
       - name: Checkout
-        uses: actions/checkout@v4.2.2
+        uses: actions/checkout@v5.0.0
         with:
             submodules: true
             fetch-depth: 0
diff --git a/.github/workflows/e2e-workflow.yaml b/.github/workflows/e2e-workflow.yaml
index 2d64c1b..9ec48b4 100644
--- a/.github/workflows/e2e-workflow.yaml
+++ b/.github/workflows/e2e-workflow.yaml
@@ -49,7 +49,7 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           ref: ${{ inputs.git_sha }}
 
@@ -89,7 +89,7 @@ jobs:
         run: sudo rm -rf ~/go/pkg/mod
       
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@v5.4.0
+        uses: actions/setup-go@v5.5.0
         with:
           go-version: ${{ env.GO_VERSION  }}
 
diff --git a/.github/workflows/license-header.yaml b/.github/workflows/license-header.yaml
index a14196e..21be500 100644
--- a/.github/workflows/license-header.yaml
+++ b/.github/workflows/license-header.yaml
@@ -36,7 +36,7 @@ jobs:
           disable-sudo: true
           disable-telemetry: true
 
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/lint-go.yaml b/.github/workflows/lint-go.yaml
index 50889c7..b96e79b 100644
--- a/.github/workflows/lint-go.yaml
+++ b/.github/workflows/lint-go.yaml
@@ -1,9 +1,5 @@
 name: Go Lint
 
-concurrency:
-  group:  ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
-  cancel-in-progress: true
-
 on:
   push:
     branches:
@@ -31,13 +27,13 @@ jobs:
           disable-sudo: true
           disable-telemetry: true
 
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
 
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@0aaccfd150d50ccaeb58ebd88d36e91967a5f35b # v5.4.0
+        uses: actions/setup-go@8e57b58e57be52ac95949151e2777ffda8501267 # v5.4.0
         with:
           go-version: ${{ env.GO_VERSION }}
 
diff --git a/.github/workflows/lint-python.yaml b/.github/workflows/lint-python.yaml
new file mode 100644
index 0000000..c279cfe
--- /dev/null
+++ b/.github/workflows/lint-python.yaml
@@ -0,0 +1,34 @@
+name: Python Lint
+
+on:
+  push:
+    paths:
+      - '**.py'
+    branches:
+      - main
+  pull_request:
+    paths:
+      - '**.py'
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Ruff linter
+        uses: astral-sh/ruff-action@v3
+        with:
+          # Update output format to enable automatic inline annotations.
+          args: check --output-format=github .
+
+      # Check formatting without making changes.
+      - name: Ruff format check
+        run: ruff format --check .
diff --git a/.github/workflows/markdown-link-check.yaml b/.github/workflows/markdown-link-check.yaml
index 8e66a55..5e71850 100644
--- a/.github/workflows/markdown-link-check.yaml
+++ b/.github/workflows/markdown-link-check.yaml
@@ -20,7 +20,7 @@ jobs:
           disable-sudo: true
           disable-telemetry: true
 
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
       - uses: gaurav-nelson/github-action-markdown-link-check@5c5dfc0ac2e225883c0e5f03a85311ec2830d368 # v1
         with:
           # this will only show errors in the output
diff --git a/.github/workflows/preset-image-build-1ES.yaml b/.github/workflows/preset-image-build-1ES.yaml
index 6771b1b..dcd9fe2 100644
--- a/.github/workflows/preset-image-build-1ES.yaml
+++ b/.github/workflows/preset-image-build-1ES.yaml
@@ -50,7 +50,7 @@ jobs:
       is_matrix_empty: ${{ steps.check_matrix_empty.outputs.is_empty }}
     steps:
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
@@ -95,7 +95,7 @@ jobs:
       max-parallel: 10
     steps:
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/preset-image-build.yaml b/.github/workflows/preset-image-build.yaml
index 614495d..e7ed603 100644
--- a/.github/workflows/preset-image-build.yaml
+++ b/.github/workflows/preset-image-build.yaml
@@ -44,7 +44,7 @@ jobs:
       is_matrix_empty: ${{ steps.check_matrix_empty.outputs.is_empty }}
     steps:
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
@@ -82,7 +82,7 @@ jobs:
       IMAGE_REF: "test/preset-tfs:latest"
     steps:
       - name: Checkout code
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
 
       - name: Build image from Dockerfile
         env:
@@ -99,7 +99,7 @@ jobs:
             -f docker/presets/models/tfs/Dockerfile ../
 
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@76071ef0d7ec797419534a183b498b4d6366cf37 # v0.31.0
+        uses: aquasecurity/trivy-action@dc5a429b52fcf669ce959baa2c2dd26090d2a6c4 # v0.32.0
         with:
           image-ref: "${{ env.IMAGE_REF }}"
           format: "table"
@@ -121,7 +121,7 @@ jobs:
       max-parallel: 10
     steps:
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/publish-rag-controller-gh-image.yaml b/.github/workflows/publish-rag-controller-gh-image.yaml
index eb71a84..414c484 100644
--- a/.github/workflows/publish-rag-controller-gh-image.yaml
+++ b/.github/workflows/publish-rag-controller-gh-image.yaml
@@ -40,7 +40,7 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           fetch-depth: 0
 
@@ -82,7 +82,7 @@ jobs:
         run: |
           echo "IMG_TAG=$(echo ${{ needs.check-tag.outputs.tag }} | tr -d v)" >> $GITHUB_ENV
 
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/publish-rag-controller-mcr-image.yaml b/.github/workflows/publish-rag-controller-mcr-image.yaml
index 05f444c..0459233 100644
--- a/.github/workflows/publish-rag-controller-mcr-image.yaml
+++ b/.github/workflows/publish-rag-controller-mcr-image.yaml
@@ -18,7 +18,7 @@ jobs:
     environment: publish-mcr
     steps:
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@v5.4.0
+        uses: actions/setup-go@v5.5.0
         with:
           go-version: ${{ env.GO_VERSION  }}
 
@@ -27,7 +27,7 @@ jobs:
           ver=${{ github.event.client_payload.tag }}
           echo "IMG_TAG=${ver#"v"}" >> $GITHUB_ENV
 
-      - uses: actions/checkout@v4.2.2
+      - uses: actions/checkout@v5.0.0
         with:
           fetch-depth: 0
           submodules: true
diff --git a/.github/workflows/publish-rag-service-gh-image.yaml b/.github/workflows/publish-rag-service-gh-image.yaml
index 96109f7..98c1542 100644
--- a/.github/workflows/publish-rag-service-gh-image.yaml
+++ b/.github/workflows/publish-rag-service-gh-image.yaml
@@ -40,7 +40,7 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           fetch-depth: 0
 
@@ -82,7 +82,7 @@ jobs:
         run: |
           echo "IMG_TAG=$(echo ${{ needs.check-tag.outputs.tag }} | tr -d v)" >> $GITHUB_ENV
 
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/publish-rag-service-mcr-image.yaml b/.github/workflows/publish-rag-service-mcr-image.yaml
index 3d18dc6..552076a 100644
--- a/.github/workflows/publish-rag-service-mcr-image.yaml
+++ b/.github/workflows/publish-rag-service-mcr-image.yaml
@@ -18,7 +18,7 @@ jobs:
     environment: publish-mcr
     steps:
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@v5.4.0
+        uses: actions/setup-go@v5.5.0
         with:
           go-version: ${{ env.GO_VERSION  }}
 
@@ -27,7 +27,7 @@ jobs:
           ver=${{ github.event.client_payload.tag }}
           echo "IMG_TAG=${ver#"v"}" >> $GITHUB_ENV
 
-      - uses: actions/checkout@v4.2.2
+      - uses: actions/checkout@v5.0.0
         with:
           fetch-depth: 0
           submodules: true
diff --git a/.github/workflows/publish-ragengine-helm-chart.yaml b/.github/workflows/publish-ragengine-helm-chart.yaml
index 4568f99..260341f 100644
--- a/.github/workflows/publish-ragengine-helm-chart.yaml
+++ b/.github/workflows/publish-ragengine-helm-chart.yaml
@@ -24,7 +24,7 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/publish-workspace-gh-image.yaml b/.github/workflows/publish-workspace-gh-image.yaml
index 4bed558..05ca813 100644
--- a/.github/workflows/publish-workspace-gh-image.yaml
+++ b/.github/workflows/publish-workspace-gh-image.yaml
@@ -40,7 +40,7 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           fetch-depth: 0
 
@@ -82,7 +82,7 @@ jobs:
         run: |
           echo "IMG_TAG=$(echo ${{ needs.check-tag.outputs.tag }} | tr -d v)" >> $GITHUB_ENV
 
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/publish-workspace-helm-chart.yaml b/.github/workflows/publish-workspace-helm-chart.yaml
index 8c1c3d5..cf81edb 100644
--- a/.github/workflows/publish-workspace-helm-chart.yaml
+++ b/.github/workflows/publish-workspace-helm-chart.yaml
@@ -24,7 +24,7 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/publish-workspace-mcr-image.yaml b/.github/workflows/publish-workspace-mcr-image.yaml
index 293446d..436b51e 100644
--- a/.github/workflows/publish-workspace-mcr-image.yaml
+++ b/.github/workflows/publish-workspace-mcr-image.yaml
@@ -18,7 +18,7 @@ jobs:
     environment: publish-mcr
     steps:
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@v5.4.0
+        uses: actions/setup-go@v5.5.0
         with:
           go-version: ${{ env.GO_VERSION  }}
 
@@ -27,7 +27,7 @@ jobs:
           ver=${{ github.event.client_payload.tag }}
           echo "IMG_TAG=${ver#"v"}" >> $GITHUB_ENV
 
-      - uses: actions/checkout@v4.2.2
+      - uses: actions/checkout@v5.0.0
         with:
           fetch-depth: 0
           submodules: true
diff --git a/.github/workflows/ragengine-e2e-workflow.yaml b/.github/workflows/ragengine-e2e-workflow.yaml
index 3696812..7171a21 100644
--- a/.github/workflows/ragengine-e2e-workflow.yaml
+++ b/.github/workflows/ragengine-e2e-workflow.yaml
@@ -48,7 +48,7 @@ jobs:
           disable-telemetry: true
 
       - name: Checkout
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           ref: ${{ inputs.git_sha }}
       
@@ -90,7 +90,7 @@ jobs:
         run: sudo rm -rf ~/go/pkg/mod
       
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@v5.4.0
+        uses: actions/setup-go@v5.5.0
         with:
           go-version: ${{ env.GO_VERSION }}
 
diff --git a/.github/workflows/trivy.yaml b/.github/workflows/trivy.yaml
index 4d6f542..c55e728 100644
--- a/.github/workflows/trivy.yaml
+++ b/.github/workflows/trivy.yaml
@@ -25,7 +25,7 @@ jobs:
     steps:
 
       - name: Checkout code
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
 
       - name: Build image from Dockerfile
         run: |
@@ -36,7 +36,7 @@ jobs:
           make ${{ matrix.target }} BUILD_FLAGS="--target dependencies"
 
       - name: Run Trivy vulnerability scanner
-        uses: aquasecurity/trivy-action@76071ef0d7ec797419534a183b498b4d6366cf37 # v0.31.0
+        uses: aquasecurity/trivy-action@dc5a429b52fcf669ce959baa2c2dd26090d2a6c4 # v0.32.0
         with:
           image-ref: ${{ matrix.image }}
           format: "table"
diff --git a/.github/workflows/unit-tests-ragengine.yaml b/.github/workflows/unit-tests-ragengine.yaml
index e90bcd4..8e9f11f 100644
--- a/.github/workflows/unit-tests-ragengine.yaml
+++ b/.github/workflows/unit-tests-ragengine.yaml
@@ -32,7 +32,7 @@ jobs:
           disable-telemetry: true
 
       - name: Check out the code
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml
index a71bf55..6b85ad9 100644
--- a/.github/workflows/unit-tests.yaml
+++ b/.github/workflows/unit-tests.yaml
@@ -1,9 +1,5 @@
 name: Unit Tests
 
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
-  cancel-in-progress: true
-
 on:
   push:
     branches: [main]
@@ -32,13 +28,13 @@ jobs:
           disable-telemetry: true
 
       - name: Check out the code
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           submodules: true
           fetch-depth: 0
 
       - name: Set up Go ${{ env.GO_VERSION }}
-        uses: actions/setup-go@0aaccfd150d50ccaeb58ebd88d36e91967a5f35b # v5.4.0
+        uses: actions/setup-go@8e57b58e57be52ac95949151e2777ffda8501267 # v5.4.0
         with:
           go-version: ${{ env.GO_VERSION  }}
 
diff --git a/.gitignore b/.gitignore
index deb7cd0..cf27bbb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,8 @@
 bin/*
 Dockerfile.cross
 __pycache__/
+.ruff_cache/
+.venv/
 
 # Test binary, build with `go test -c`
 *.test
diff --git a/Makefile b/Makefile
index b6b8ba6..cd22fed 100644
--- a/Makefile
+++ b/Makefile
@@ -2,7 +2,7 @@
 # Image URL to use all building/pushing image targets
 REGISTRY ?= YOUR_REGISTRY
 IMG_NAME ?= workspace
-VERSION ?= v0.5.1
+VERSION ?= v0.6.0
 GPU_PROVISIONER_VERSION ?= 0.3.5
 RAGENGINE_IMG_NAME ?= ragengine
 IMG_TAG ?= $(subst v,,$(VERSION))
@@ -589,13 +589,24 @@ release-manifest: ## Update manifest and Helm charts for release.
 	@sed -i -e "s/tag: .*/tag: ${IMG_TAG}/" ./charts/kaito/ragengine/values.yaml
 	@sed -i -e "s/presetRagImageTag: .*/presetRagImageTag: ${IMG_TAG}/" ./charts/kaito/ragengine/values.yaml
 	@sed -i -e 's/IMG_TAG=.*/IMG_TAG=${IMG_TAG}/' ./charts/kaito/ragengine/README.md
-	@sed -i -e 's/export KAITO_WORKSPACE_VERSION=.*/export KAITO_WORKSPACE_VERSION=${IMG_TAG}/' ./website/docs/installation.md
-	@sed -i -e 's/$(shell grep 'default' ./terraform/variables.tf | sort | uniq -c | sort -nr | head -n1 | sed -E 's/^[ ]*[0-9]+[ ]+//')/default     = "${IMG_TAG}"/' ./terraform/variables.tf
 
 	git checkout -b release-${VERSION}
-	git add ./Makefile ./charts/kaito/workspace/Chart.yaml ./charts/kaito/workspace/values.yaml ./charts/kaito/workspace/README.md ./charts/kaito/ragengine/Chart.yaml ./charts/kaito/ragengine/values.yaml ./charts/kaito/ragengine/README.md ./website/docs/installation.md /terraform/variables.tf
+	git add ./Makefile ./charts/kaito/workspace/Chart.yaml ./charts/kaito/workspace/values.yaml ./charts/kaito/workspace/README.md ./charts/kaito/ragengine/Chart.yaml ./charts/kaito/ragengine/values.yaml ./charts/kaito/ragengine/README.md
 	git commit -s -m "release: update manifest and helm charts for ${VERSION}"
 
+## --------------------------------------
+## Post Release Doc Update
+## To update documents after a release, run `make post-release-doc-update VERSION=vx.y.z`
+## --------------------------------------
+
+##@ Post Release Doc Update
+
+.PHONY: post-release-doc-update
+post-release-doc-update:
+	@sed -i -e 's/export KAITO_WORKSPACE_VERSION=.*/export KAITO_WORKSPACE_VERSION=${IMG_TAG}/' ./website/docs/installation.md
+	@sed -i -e 's/$(shell grep 'default' ./terraform/variables.tf | sort | uniq -c | sort -nr | head -n1 | sed -E 's/^[ ]*[0-9]+[ ]+//')/default     = "${IMG_TAG}"/' ./terraform/variables.tf
+
+
 ## --------------------------------------
 ## Cleanup
 ## --------------------------------------
@@ -616,4 +627,4 @@ clean-bin: ## Remove all generated binaries.
 ##@ Help:
 
 help: ## Display this help.
-	@awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n  make \033[36m<target>\033[0m\n"} /^[0-9a-zA-Z_-]+:.*?##/ { printf "  \033[36m%-25s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST)
\ No newline at end of file
+	@awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n  make \033[36m<target>\033[0m\n"} /^[0-9a-zA-Z_-]+:.*?##/ { printf "  \033[36m%-25s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST)
diff --git a/README.md b/README.md
index a17d93c..d4b1cd3 100644
--- a/README.md
+++ b/README.md
@@ -8,8 +8,8 @@
 
 | ![notification](website/static/img/bell.svg) What is NEW!                                                                                                                                                                                                |
 | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about from [here](https://kaito-project.github.io/kaito/docs/rag)! |
-| Latest Release: July 18th, 2025. KAITO v0.5.1.                                                                                                                                                                                                  |
+| Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about it [here](https://kaito-project.github.io/kaito/docs/rag)! |
+| Latest Release: Aug 7th, 2025. KAITO v0.6.0                                                                                                                                                                                                  |
 | First Release: Nov 15th, 2023. KAITO v0.1.0.                                                                                                                                                                                                   |
 
 KAITO is an operator that automates the AI/ML model inference or tuning workload in a Kubernetes cluster.
@@ -43,7 +43,7 @@ The above figure presents the KAITO architecture overview. Its major components
 </div>
 
 As illustrated in the above figure, the **RAGEngine controller** reconciles the `ragengine` custom resource and creates a `RAGService` deployment. The `RAGService` provides the following capabilities:
-  - **Orchstration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
+  - **Orchestration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
   - **Embedding**: support both local and remote embedding services, to embed queries and documents in the vector database.
   - **Vector database**: support a built-in [faiss](https://github.com/facebookresearch/faiss) in-memory vector database. Remote vector database support will be added soon.
   - **Backend inference**: support any OAI compatible inference service.
@@ -78,7 +78,7 @@ inference:
 $ kubectl apply -f examples/inference/kaito_workspace_phi_3.5-instruct.yaml
 ```
 
-The workspace status can be tracked by running the following command. When the WORKSPACEREADY column becomes `True`, the model has been deployed successfully.
+The workspace status can be tracked by running the following command. When the WORKSPACESUCCEEDED column becomes `True`, the model has been deployed successfully.
 
 ```sh
 $ kubectl get workspace workspace-phi-3-5-mini
diff --git a/api/v1beta1/workspace_validation.go b/api/v1beta1/workspace_validation.go
index c3af2f3..7013f7d 100644
--- a/api/v1beta1/workspace_validation.go
+++ b/api/v1beta1/workspace_validation.go
@@ -29,6 +29,7 @@ import (
 	"k8s.io/klog/v2"
 	"knative.dev/pkg/apis"
 
+	"github.com/kaito-project/kaito/pkg/featuregates"
 	"github.com/kaito-project/kaito/pkg/model"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
@@ -57,6 +58,16 @@ func (w *Workspace) Validate(ctx context.Context) (errs *apis.FieldError) {
 	if len(errmsgs) > 0 {
 		errs = errs.Also(apis.ErrInvalidValue(strings.Join(errmsgs, ", "), "name"))
 	}
+
+	// Check if node auto-provisioning is disabled and validate preferred nodes
+	if featuregates.FeatureGates[consts.FeatureFlagDisableNodeAutoProvisioning] {
+		if len(w.Resource.PreferredNodes) < *w.Resource.Count {
+			errs = errs.Also(apis.ErrInvalidValue(
+				fmt.Sprintf("When node auto-provisioning is disabled, the number of preferred nodes (%d) must be at least the required amount (%d) set in count",
+					len(w.Resource.PreferredNodes), *w.Resource.Count), "preferredNodes"))
+		}
+	}
+
 	base := apis.GetBaseline(ctx)
 	if base == nil {
 		klog.InfoS("Validate creation", "workspace", fmt.Sprintf("%s/%s", w.Namespace, w.Name))
diff --git a/charts/kaito/ragengine/Chart.yaml b/charts/kaito/ragengine/Chart.yaml
index 9436de9..013a4d3 100644
--- a/charts/kaito/ragengine/Chart.yaml
+++ b/charts/kaito/ragengine/Chart.yaml
@@ -6,13 +6,13 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 0.5.1
+version: 0.6.0
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to
 # follow Semantic Versioning. They should reflect the version the application is using.
 # It is recommended to use it with quotes.
-appVersion: 0.5.1
+appVersion: 0.6.0
 home: https://github.com/kaito-project/kaito
 sources:
   - https://github.com/kaito-project/kaito
diff --git a/charts/kaito/ragengine/README.md b/charts/kaito/ragengine/README.md
index 87fa368..e550f90 100644
--- a/charts/kaito/ragengine/README.md
+++ b/charts/kaito/ragengine/README.md
@@ -7,7 +7,7 @@ KAITO RAGEngine provides Retrieval-Augmented Generation (RAG) capabilities for A
 ```bash
 export REGISTRY=mcr.microsoft.com/aks/kaito
 export IMG_NAME=ragengine
-export IMG_TAG=0.5.1
+export IMG_TAG=0.6.0
 helm install ragengine ./charts/kaito/ragengine  \
 --set image.repository=${REGISTRY}/${IMG_NAME} --set image.tag=${IMG_TAG} \
 --namespace kaito-ragengine --create-namespace
diff --git a/charts/kaito/ragengine/templates/deployment.yaml b/charts/kaito/ragengine/templates/deployment.yaml
index 7eea448..122ef62 100644
--- a/charts/kaito/ragengine/templates/deployment.yaml
+++ b/charts/kaito/ragengine/templates/deployment.yaml
@@ -33,6 +33,8 @@ spec:
           image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
           imagePullPolicy: {{ .Values.image.pullPolicy }}
           env:
+            - name: CONFIG_LOGGING_NAME
+              value: "kaito-logging-config"
             - name: WEBHOOK_SERVICE
               value: {{ include "kaito.fullname" . }}
             - name: WEBHOOK_PORT
diff --git a/charts/kaito/ragengine/templates/knative-logging-configmap.yaml b/charts/kaito/ragengine/templates/knative-logging-configmap.yaml
index ad06d09..c57bda1 100644
--- a/charts/kaito/ragengine/templates/knative-logging-configmap.yaml
+++ b/charts/kaito/ragengine/templates/knative-logging-configmap.yaml
@@ -1,7 +1,7 @@
 apiVersion: v1
 kind: ConfigMap
 metadata:
-  name: config-logging
+  name: kaito-logging-config
   namespace: {{ .Release.Namespace }}
 data:
   loglevel.controller: {{ .Values.logging.level | quote }}
diff --git a/charts/kaito/ragengine/values.yaml b/charts/kaito/ragengine/values.yaml
index 9d14c03..0e662a3 100644
--- a/charts/kaito/ragengine/values.yaml
+++ b/charts/kaito/ragengine/values.yaml
@@ -5,7 +5,7 @@ replicaCount: 1
 image:
   repository: mcr.microsoft.com/aks/kaito/ragengine
   pullPolicy: IfNotPresent
-  tag: 0.5.1
+  tag: 0.6.0
 imagePullSecrets: []
 podAnnotations: {}
 podSecurityContext:
@@ -34,4 +34,4 @@ affinity: {}
 cloudProviderName: "azure"
 presetRagRegistryName: "mcr.microsoft.com/aks/kaito"
 presetRagImageName: "kaito-rag-service"
-presetRagImageTag: 0.5.1
\ No newline at end of file
+presetRagImageTag: 0.6.0
\ No newline at end of file
diff --git a/charts/kaito/workspace/Chart.lock b/charts/kaito/workspace/Chart.lock
new file mode 100644
index 0000000..aca7f36
--- /dev/null
+++ b/charts/kaito/workspace/Chart.lock
@@ -0,0 +1,6 @@
+dependencies:
+- name: flux2
+  repository: https://fluxcd-community.github.io/helm-charts
+  version: 2.16.3
+digest: sha256:a4d238f7903ccdd446780a21c718d02cc3d6d3979b80c5b1bae848022e8132ac
+generated: "2025-08-06T19:41:44.694298719-04:00"
diff --git a/charts/kaito/workspace/Chart.yaml b/charts/kaito/workspace/Chart.yaml
index 043361a..c0a2df1 100644
--- a/charts/kaito/workspace/Chart.yaml
+++ b/charts/kaito/workspace/Chart.yaml
@@ -6,13 +6,13 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 0.5.1
+version: 0.6.0
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to
 # follow Semantic Versioning. They should reflect the version the application is using.
 # It is recommended to use it with quotes.
-appVersion: 0.5.1
+appVersion: 0.6.0
 home: https://github.com/kaito-project/kaito
 sources:
   - https://github.com/kaito-project/kaito
@@ -32,3 +32,8 @@ annotations:
     - name: Examples
       url: https://github.com/kaito-project/kaito/tree/main/examples
   artifacthub.io/prerelease: "false"
+dependencies:
+- name: flux2
+  version: "2.16.3"
+  repository: https://fluxcd-community.github.io/helm-charts
+  condition: flux2.enabled
diff --git a/charts/kaito/workspace/README.md b/charts/kaito/workspace/README.md
index b473252..b5ff493 100644
--- a/charts/kaito/workspace/README.md
+++ b/charts/kaito/workspace/README.md
@@ -5,7 +5,7 @@
 ```bash
 export REGISTRY=mcr.microsoft.com/aks/kaito
 export IMG_NAME=workspace
-export IMG_TAG=0.5.1
+export IMG_TAG=0.6.0
 helm install workspace ./charts/kaito/workspace  \
 --set image.repository=${REGISTRY}/${IMG_NAME} --set image.tag=${IMG_TAG} \
 --namespace kaito-workspace --create-namespace
diff --git a/charts/kaito/workspace/charts/flux2-2.16.3.tgz b/charts/kaito/workspace/charts/flux2-2.16.3.tgz
new file mode 100644
index 0000000..193480e
Binary files /dev/null and b/charts/kaito/workspace/charts/flux2-2.16.3.tgz differ
diff --git a/charts/kaito/workspace/templates/deployment.yaml b/charts/kaito/workspace/templates/deployment.yaml
index d8c73dd..8012903 100644
--- a/charts/kaito/workspace/templates/deployment.yaml
+++ b/charts/kaito/workspace/templates/deployment.yaml
@@ -35,6 +35,8 @@ spec:
           args:
             - --feature-gates={{ include "utils.joinKeyValuePairs" .Values.featureGates }}
           env:
+            - name: CONFIG_LOGGING_NAME
+              value: "kaito-logging-config"
             - name: WEBHOOK_SERVICE
               value: {{ include "kaito.fullname" . }}-svc
             - name: WEBHOOK_PORT
diff --git a/charts/kaito/workspace/templates/knative-logging-configmap.yaml b/charts/kaito/workspace/templates/knative-logging-configmap.yaml
index ad06d09..c57bda1 100644
--- a/charts/kaito/workspace/templates/knative-logging-configmap.yaml
+++ b/charts/kaito/workspace/templates/knative-logging-configmap.yaml
@@ -1,7 +1,7 @@
 apiVersion: v1
 kind: ConfigMap
 metadata:
-  name: config-logging
+  name: kaito-logging-config
   namespace: {{ .Release.Namespace }}
 data:
   loglevel.controller: {{ .Values.logging.level | quote }}
diff --git a/charts/kaito/workspace/templates/local-csi-driver-ds.yaml b/charts/kaito/workspace/templates/local-csi-driver-ds.yaml
index ef24e2d..c86d580 100644
--- a/charts/kaito/workspace/templates/local-csi-driver-ds.yaml
+++ b/charts/kaito/workspace/templates/local-csi-driver-ds.yaml
@@ -249,9 +249,6 @@ spec:
         volumeMounts:
         - mountPath: /dev
           name: device
-        - mountPath: /etc/kubernetes/
-          name: k8s-cfg
-          readOnly: true
         - mountPath: /csi
           name: csi-socket-dir
         - mountPath: /var/lib/kubelet/
@@ -391,10 +388,6 @@ spec:
           path: /dev
           type: Directory
         name: device
-      - hostPath:
-          path: /etc/kubernetes/
-          type: DirectoryOrCreate
-        name: k8s-cfg
       - hostPath:
           path: /var/lib/kubelet/plugins/localdisk.csi.acstor.io/
           type: DirectoryOrCreate
diff --git a/charts/kaito/workspace/templates/local-csi-storageclass.yaml b/charts/kaito/workspace/templates/local-csi-storageclass.yaml
index 03b5d33..f54e2dc 100644
--- a/charts/kaito/workspace/templates/local-csi-storageclass.yaml
+++ b/charts/kaito/workspace/templates/local-csi-storageclass.yaml
@@ -1,16 +1,12 @@
 # notes: This CSI driver is compatible with various environments, including Azure and AWS,
 # ensuring broad interoperability without vendor lock-in.
-{{- if not (lookup "storage.k8s.io/v1" "StorageClass" "" "kaito-local-nvme-disk") }}
 apiVersion: storage.k8s.io/v1
 kind: StorageClass
 metadata:
   name: kaito-local-nvme-disk
   labels:
     {{- include "kaito.labels" . | nindent 4 }}
-parameters:
-  {}
 provisioner: localdisk.csi.acstor.io
 reclaimPolicy: Delete
 volumeBindingMode: WaitForFirstConsumer
 allowVolumeExpansion: true
-{{- end }}
diff --git a/charts/kaito/workspace/templates/supported-models-configmap.yaml b/charts/kaito/workspace/templates/supported-models-configmap.yaml
index 4daba4c..21c8f11 100644
--- a/charts/kaito/workspace/templates/supported-models-configmap.yaml
+++ b/charts/kaito/workspace/templates/supported-models-configmap.yaml
@@ -9,7 +9,7 @@ data:
       - name: base
         type: text-generation
         runtime: tfs
-        tag: 0.0.4
+        tag: 0.0.5
       - name: llama-3.1-8b-instruct
         type: text-generation
         version: https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct/commit/0e9e39f249a16976918f6564b8830bc894c89659
diff --git a/charts/kaito/workspace/values.yaml b/charts/kaito/workspace/values.yaml
index d467fb1..88b2158 100644
--- a/charts/kaito/workspace/values.yaml
+++ b/charts/kaito/workspace/values.yaml
@@ -5,7 +5,7 @@ replicaCount: 1
 image:
   repository: mcr.microsoft.com/aks/kaito/workspace
   pullPolicy: IfNotPresent
-  tag: 0.5.1
+  tag: 0.6.0
 imagePullSecrets: []
 podAnnotations: {}
 podSecurityContext:
@@ -17,6 +17,7 @@ securityContext:
       - "ALL"
 featureGates:
   vLLM: "true"
+  disableNodeAutoProvisioning: "false"
 webhook:
   port: 9443
 # Knative logging configuration
@@ -36,3 +37,18 @@ affinity: {}
 # Values can be "azure" or "aws" or "arc"
 cloudProviderName: "azure"
 clusterName: "kaito"
+# Set values for the Flux controller to only install the Helm controller
+flux2:
+  enabled: false
+  helmController:
+    create: false
+  sourceController:
+    create: false
+  imageAutomationController:
+    create: false
+  imageReflectionController:
+    create: false
+  kustomizeController:
+    create: false
+  notificationController:
+    create: false
\ No newline at end of file
diff --git a/cmd/workspace/main.go b/cmd/workspace/main.go
index a190ec8..3a09451 100644
--- a/cmd/workspace/main.go
+++ b/cmd/workspace/main.go
@@ -87,7 +87,7 @@ func main() {
 			"Enabling this will ensure there is only one active controller manager.")
 	flag.BoolVar(&enableWebhook, "webhook", true,
 		"Enable webhook for controller manager. Default is true.")
-	flag.StringVar(&featureGates, "feature-gates", "vLLM=true", "Enable Kaito feature gates. Default,	vLLM=true.")
+	flag.StringVar(&featureGates, "feature-gates", "vLLM=true,disableNodeAutoProvisioning=false", "Enable Kaito feature gates. Default: vLLM=true,disableNodeAutoProvisioning=false.")
 	opts := zap.Options{
 		Development: true,
 	}
diff --git a/demo/inferenceUI/chainlit_openai.py b/demo/inferenceUI/chainlit_openai.py
index f79759f..316ef8a 100644
--- a/demo/inferenceUI/chainlit_openai.py
+++ b/demo/inferenceUI/chainlit_openai.py
@@ -14,10 +14,10 @@
 import os
 from urllib.parse import urljoin
 
-from openai import AsyncOpenAI
 import chainlit as cl
+from openai import AsyncOpenAI
 
-URL = os.environ.get('WORKSPACE_SERVICE_URL')
+URL = os.environ.get("WORKSPACE_SERVICE_URL")
 
 client = AsyncOpenAI(base_url=urljoin(URL, "v1"), api_key="YOUR_OPENAI_API_KEY")
 cl.instrument_openai()
@@ -30,6 +30,7 @@ settings = {
     "presence_penalty": 0,
 }
 
+
 @cl.on_chat_start
 async def start_chat():
     models = await client.models.list()
@@ -41,24 +42,17 @@ async def start_chat():
     model = models.data[0].id
     print(f"Using model: {model}")
 
+
 @cl.on_message
 async def main(message: cl.Message):
-    messages=[
-        {
-            "content": "You are a helpful assistant.",
-            "role": "system"
-        },
-        {
-            "content": message.content,
-            "role": "user"
-        }
+    messages = [
+        {"content": "You are a helpful assistant.", "role": "system"},
+        {"content": message.content, "role": "user"},
     ]
     msg = cl.Message(content="")
 
     stream = await client.chat.completions.create(
-        messages=messages, model=model,
-        stream=True,
-        **settings
+        messages=messages, model=model, stream=True, **settings
     )
 
     async for part in stream:
diff --git a/demo/inferenceUI/chainlit_transformers.py b/demo/inferenceUI/chainlit_transformers.py
index 9ec0b8f..8b7c577 100644
--- a/demo/inferenceUI/chainlit_transformers.py
+++ b/demo/inferenceUI/chainlit_transformers.py
@@ -17,7 +17,9 @@ from urllib.parse import urljoin
 import chainlit as cl
 import requests
 
-URL = os.environ.get('WORKSPACE_SERVICE_URL')
+URL = os.environ.get("WORKSPACE_SERVICE_URL")
+
+
 @cl.step
 def inference(prompt):
     # Endpoint URL
@@ -35,8 +37,8 @@ def inference(prompt):
             "temperature": 1.0,
             "top_p": 1,
             "typical_p": 1,
-            "repetition_penalty": 1
-        }
+            "repetition_penalty": 1,
+        },
     }
 
     response = requests.post(urljoin(URL, "chat"), json=data)
@@ -47,6 +49,7 @@ def inference(prompt):
     else:
         return f"Error: Received response code {response.status_code}"
 
+
 @cl.on_message
 async def main(message: cl.Message):
     """
diff --git a/docker/presets/models/tfs-onnx/convert_to_onnx.py b/docker/presets/models/tfs-onnx/convert_to_onnx.py
index 62fb71b..41f2128 100644
--- a/docker/presets/models/tfs-onnx/convert_to_onnx.py
+++ b/docker/presets/models/tfs-onnx/convert_to_onnx.py
@@ -12,9 +12,13 @@
 # limitations under the License.
 
 import sys
-import subprocess
-from optimum.onnxruntime import AutoOptimizationConfig, ORTModelForCausalLM, ORTOptimizer
-import glob
+
+from optimum.onnxruntime import (
+    AutoOptimizationConfig,
+    ORTModelForCausalLM,
+    ORTOptimizer,
+)
+
 
 def download_and_convert(repo_name):
     """
@@ -46,7 +50,9 @@ def download_and_convert(repo_name):
 
     # Try converting to ONNX with caching, then without if fails
     try:
-        model = ORTModelForCausalLM.from_pretrained(repo_name, export=True, provider="CUDAExecutionProvider")
+        model = ORTModelForCausalLM.from_pretrained(
+            repo_name, export=True, provider="CUDAExecutionProvider"
+        )
         model.save_pretrained(f"{repo_name}")
         print(f"Model converted and saved under {repo_name}")
         return model
@@ -54,7 +60,9 @@ def download_and_convert(repo_name):
         print(f"Failed to convert model to ONNX with caching: {e}")
 
     try:
-        model = ORTModelForCausalLM.from_pretrained(repo_name, use_cache=False, export=True, provider="CUDAExecutionProvider")
+        model = ORTModelForCausalLM.from_pretrained(
+            repo_name, use_cache=False, export=True, provider="CUDAExecutionProvider"
+        )
         model.save_pretrained(f"{repo_name}")
         print(f"Model converted without cache and saved under {repo_name}")
         return model
@@ -62,6 +70,7 @@ def download_and_convert(repo_name):
         print(f"Failed to convert model from {repo_name} without caching: {e}")
     return None
 
+
 def onnx_optimize_model(model, repo_name):
     try:
         optimizer = ORTOptimizer.from_pretrained(model)
@@ -72,9 +81,9 @@ def onnx_optimize_model(model, repo_name):
     except Exception as e:
         print("Optimizing model failed", e)
 
+
 if __name__ == "__main__":
     repo_name = sys.argv[1]
     model = download_and_convert(repo_name)
     if model:
         onnx_optimize_model(model, repo_name)
-
diff --git a/docs/Release_Management.md b/docs/Release_Management.md
new file mode 100644
index 0000000..488295b
--- /dev/null
+++ b/docs/Release_Management.md
@@ -0,0 +1,64 @@
+# Release Management
+
+## Overview
+
+This document describes KAITO project release management, which includes release versioning, supported releases, and supported upgrades.
+
+## Legend
+
+- **X.Y.Z** refers to the version (git tag) of KAITO that is released. This is the version of the KAITO images and the Charts version.
+- **Breaking changes** refer to schema changes, flag changes, and behavior changes of KAITO that may require a clean installation during upgrade, and it may introduce changes that could break backward compatibility.
+- **Milestone** should be designed to include feature sets to accommodate 2 months release cycles including test gates. GitHub's milestones are used by maintainers to manage each release. PRs and Issues for each release should be created as part of a corresponding milestone.
+- **Patch releases** refer to applicable fixes, including security fixes, may be backported to support releases, depending on severity and feasibility.
+- **Test gates** should include soak tests and upgrade tests from the last minor version.
+
+## Release Versioning
+
+All releases will be of the form _vX.Y.Z_ where X is the major version, Y is the minor version and Z is the patch version. This project strictly follows semantic versioning.
+
+The rest of the doc will cover the release process for the following kinds of releases:
+
+**Major Releases**
+
+No plan to move to 1.0.0 unless there is a major design change like an incompatible API change in the project
+
+**Minor Releases**
+
+- X.Y.0-rc.W, W >= 0 (Branch: release-X.Y)
+    - Released as needed before we cut a stable X.Y release
+    - Soak for a total of ~1 week before cutting a stable release
+    - Release candidate release, cut from release-X.Y branch
+- X.Y.0 (Branch: release-X.Y)
+    - Released every ~2 months (with the exception of emergency minor releases to support new models)
+    - Stable release, cut from release-X.Y branch when X.Y milestone is complete
+
+**Patch Releases**
+
+- Patch Releases X.Y.Z, Z > 0 (Branch: release-X.Y, only cut when a patch is needed)
+    - No breaking changes
+    - Applicable fixes, including security fixes, may be cherry-picked from main into the latest supported minor release-X.Y branches.
+    - Patch release, cut from a release-X.Y branch
+
+## Supported Releases
+
+Applicable fixes, including security fixes, may be cherry-picked into the release branch, depending on severity and feasibility. Patch releases are cut from that branch as needed.
+
+We expect users to stay reasonably up-to-date with the versions of KAITO they use in production, but understand that it may take time to upgrade. We expect users to be running approximately the latest patch release of a given minor release and encourage users to upgrade as soon as possible.
+
+We expect to "support" n (current) and n-1 major.minor releases. "Support" means we expect users to be running that version in production. For example, when v0.5.0 comes out, v0.3.x will no longer be supported for patches, and we encourage users to upgrade to a supported version as soon as possible.
+
+## Supported Kubernetes Versions
+
+KAITO is assumed to be compatible with the [current Kubernetes Supported Versions](https://kubernetes.io/releases/patch-releases/#detailed-release-history-for-active-branches) per [Kubernetes Supported Versions policy](https://kubernetes.io/releases/version-skew-policy/).
+
+For example, if KAITO _supported_ versions are v0.5 and v0.4, and Kubernetes _supported_ versions are v1.30, v1.31, v1.32, then all supported KAITO versions (v0.5, v0.4) are assumed to be compatible with all supported Kubernetes versions (v1.30, v1.31, v1.32). If Kubernetes v1.33 is released later, then KAITO v0.5 and v0.4 will be assumed to be compatible with v1.33 if those KAITO versions are still supported at that time.
+
+If you choose to use KAITO with a version of Kubernetes that it does not support, you are using it at your own risk.
+
+## Upgrades
+
+Users should be able to run both X.Y and X.Y + 1 simultaneously in order to support gradual rollouts.
+
+## Acknowledgement
+
+This document builds on the ideas and implementations of release processes from projects like Kubernetes and Helm.
diff --git a/go.mod b/go.mod
index b302468..12d06a0 100644
--- a/go.mod
+++ b/go.mod
@@ -23,7 +23,7 @@ require (
 	k8s.io/component-base v0.33.2
 	k8s.io/component-helpers v0.33.2
 	k8s.io/klog/v2 v2.130.1
-	k8s.io/utils v0.0.0-20241210054802-24370beab758
+	k8s.io/utils v0.0.0-20250321185631-1f6e0b77f77e
 	knative.dev/pkg v0.0.0-20240910170930-fdbc0b5adde7
 	sigs.k8s.io/controller-runtime v0.21.0
 	sigs.k8s.io/karpenter v1.0.4
@@ -40,8 +40,11 @@ require (
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/emicklei/go-restful/v3 v3.12.1 // indirect
 	github.com/evanphx/json-patch/v5 v5.9.11 // indirect
+	github.com/fluxcd/helm-controller/api v1.3.0 // indirect
+	github.com/fluxcd/pkg/apis/kustomize v1.10.0 // indirect
+	github.com/fluxcd/pkg/apis/meta v1.12.0 // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
-	github.com/fxamacker/cbor/v2 v2.7.0 // indirect
+	github.com/fxamacker/cbor/v2 v2.8.0 // indirect
 	github.com/go-kit/log v0.2.1 // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
 	github.com/go-logr/zapr v1.3.0 // indirect
@@ -97,7 +100,7 @@ require (
 	golang.org/x/term v0.32.0 // indirect
 	golang.org/x/text v0.25.0 // indirect
 	golang.org/x/time v0.9.0 // indirect
-	golang.org/x/tools v0.31.0 // indirect
+	golang.org/x/tools v0.32.0 // indirect
 	gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
 	google.golang.org/api v0.183.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20250324211829-b45e905df463 // indirect
@@ -111,7 +114,7 @@ require (
 	k8s.io/cloud-provider v0.33.2 // indirect
 	k8s.io/csi-translation-lib v0.33.2 // indirect
 	k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff // indirect
-	sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3 // indirect
+	sigs.k8s.io/json v0.0.0-20241014173422-cfa47c3a1cc8 // indirect
 	sigs.k8s.io/randfill v1.0.0 // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.7.0 // indirect
 	sigs.k8s.io/yaml v1.4.0 // indirect
diff --git a/go.sum b/go.sum
index 9e07944..ddc7614 100644
--- a/go.sum
+++ b/go.sum
@@ -162,10 +162,18 @@ github.com/evanphx/json-patch/v5 v5.9.11 h1:/8HVnzMq13/3x9TPvjG08wUGqBTmZBsCWzjT
 github.com/evanphx/json-patch/v5 v5.9.11/go.mod h1:3j+LviiESTElxA4p3EMKAB9HXj3/XEtnUf6OZxqIQTM=
 github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
 github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
+github.com/fluxcd/helm-controller/api v1.3.0 h1:PupXPuQbksmU0g2Lc6NjIYal2HJGL+6xohsf82eGVjo=
+github.com/fluxcd/helm-controller/api v1.3.0/go.mod h1:4b8PfdH0e/9Pfol2ogdMYbQ1nLjcVu9gAv27cQzIPK4=
+github.com/fluxcd/pkg/apis/kustomize v1.10.0 h1:47EeSzkQvlQZdH92vHMe2lK2iR8aOSEJq95avw5idts=
+github.com/fluxcd/pkg/apis/kustomize v1.10.0/go.mod h1:UsqMV4sqNa1Yg0pmTsdkHRJr7bafBOENIJoAN+3ezaQ=
+github.com/fluxcd/pkg/apis/meta v1.12.0 h1:XW15TKZieC2b7MN8VS85stqZJOx+/b8jATQ/xTUhVYg=
+github.com/fluxcd/pkg/apis/meta v1.12.0/go.mod h1:+son1Va60x2eiDcTwd7lcctbI6C+K3gM7R+ULmEq1SI=
 github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
 github.com/fsnotify/fsnotify v1.7.0/go.mod h1:40Bi/Hjc2AVfZrqy+aj+yEI+/bRxZnMJyTJwOpGvigM=
 github.com/fxamacker/cbor/v2 v2.7.0 h1:iM5WgngdRBanHcxugY4JySA0nk1wZorNOpTgCMedv5E=
 github.com/fxamacker/cbor/v2 v2.7.0/go.mod h1:pxXPTn3joSm21Gbwsv0w9OSA2y1HFR9qXEeXQVeNoDQ=
+github.com/fxamacker/cbor/v2 v2.8.0 h1:fFtUGXUzXPHTIUdne5+zzMPTfffl3RD5qYnkY40vtxU=
+github.com/fxamacker/cbor/v2 v2.8.0/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ=
 github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0=
 github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
@@ -712,6 +720,7 @@ golang.org/x/tools v0.0.0-20200825202427-b303f430e36d/go.mod h1:njjCfa9FT2d7l9Bc
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.31.0 h1:0EedkvKDbh+qistFTd0Bcwe/YLh4vHwWEkiI0toFIBU=
 golang.org/x/tools v0.31.0/go.mod h1:naFTU+Cev749tSJRXJlna0T3WxKvb1kWEx15xA4SdmQ=
+golang.org/x/tools v0.32.0/go.mod h1:ZxrU41P/wAbZD8EDa6dDCa6XfpkhJ7HFMjHJXfBDu8s=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
@@ -864,6 +873,8 @@ k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff h1:/usPimJzUKKu+m+TE36gUy
 k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff/go.mod h1:5jIi+8yX4RIb8wk3XwBo5Pq2ccx4FP10ohkbSKCZoK8=
 k8s.io/utils v0.0.0-20241210054802-24370beab758 h1:sdbE21q2nlQtFh65saZY+rRM6x6aJJI8IUa1AmH/qa0=
 k8s.io/utils v0.0.0-20241210054802-24370beab758/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
+k8s.io/utils v0.0.0-20250321185631-1f6e0b77f77e h1:KqK5c/ghOm8xkHYhlodbp6i6+r+ChV2vuAuVRdFbLro=
+k8s.io/utils v0.0.0-20250321185631-1f6e0b77f77e/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
 knative.dev/pkg v0.0.0-20240910170930-fdbc0b5adde7 h1:ClGJ7Q7iQC9qxov1lNys7xC0MqDGaw/ysGEuyqdIwAM=
 knative.dev/pkg v0.0.0-20240910170930-fdbc0b5adde7/go.mod h1:fL50zroI/eGDwkkWTJsOpto46G8T/xmTjBzkwnh/S48=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
@@ -879,6 +890,8 @@ sigs.k8s.io/controller-runtime v0.21.0 h1:CYfjpEuicjUecRk+KAeyYh+ouUBn4llGyDYytI
 sigs.k8s.io/controller-runtime v0.21.0/go.mod h1:OSg14+F65eWqIu4DceX7k/+QRAbTTvxeQSNSOQpukWM=
 sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3 h1:/Rv+M11QRah1itp8VhT6HoVx1Ray9eB4DBr+K+/sCJ8=
 sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3/go.mod h1:18nIHnGi6636UCz6m8i4DhaJ65T6EruyzmoQqI2BVDo=
+sigs.k8s.io/json v0.0.0-20241014173422-cfa47c3a1cc8 h1:gBQPwqORJ8d8/YNZWEjoZs7npUVDpVXUUOFfW6CgAqE=
+sigs.k8s.io/json v0.0.0-20241014173422-cfa47c3a1cc8/go.mod h1:mdzfpAEoE6DHQEN0uh9ZbOCuHbLK5wOm7dK4ctXE9Tg=
 sigs.k8s.io/karpenter v1.0.4 h1:tE4gbYjXZ2uny35QmrKFSuMhxkIfGCqCulXhwyN2D3c=
 sigs.k8s.io/karpenter v1.0.4/go.mod h1:3NLmsnHHw8p4VutpjTOPUZyhE3qH6yGTs8O94Lsu8uw=
 sigs.k8s.io/randfill v0.0.0-20250304075658-069ef1bbf016/go.mod h1:XeLlZ/jmk4i1HRopwe7/aU3H5n1zNUcX6TM94b3QxOY=
diff --git a/hack/local-csi-driver-charts-sync/local-csi-driver-charts-sync.sh b/hack/local-csi-driver-charts-sync/local-csi-driver-charts-sync.sh
index 9636892..9e8b37d 100755
--- a/hack/local-csi-driver-charts-sync/local-csi-driver-charts-sync.sh
+++ b/hack/local-csi-driver-charts-sync/local-csi-driver-charts-sync.sh
@@ -32,7 +32,7 @@ echo "Created temporary directory: $TEMP_DIR"
 echo "Cloning $REPO repository..."
 git clone "$REPO" "$REPO_DIR"
 echo "Checking out revision: $REVISION"
-(cd "$REPO_DIR" && git switch -c "$REVISION") || { echo "Failed to checkout revision $REVISION"; exit 1; }
+(cd "$REPO_DIR" && git checkout "$REVISION") || { echo "Failed to checkout revision $REVISION"; exit 1; }
 
 helm template --release-name local-csi-driver \
   --set image.driver.tag="0.0.1-latest" \
diff --git a/hack/local-csi-driver-charts-sync/local-csi-driver-charts.diff b/hack/local-csi-driver-charts-sync/local-csi-driver-charts.diff
index 0abe49d..906ce35 100644
--- a/hack/local-csi-driver-charts-sync/local-csi-driver-charts.diff
+++ b/hack/local-csi-driver-charts-sync/local-csi-driver-charts.diff
@@ -1,5 +1,5 @@
 diff --git a/./local-csi-driver-charts.yaml b/./charts/kaito/workspace/templates/local-csi-driver-ds.yaml
-index 93d0b57..54cb7e4 100644
+index 93d0b571..cd8173e2 100644
 --- a/./local-csi-driver-charts.yaml
 +++ b/./charts/kaito/workspace/templates/local-csi-driver-ds.yaml
 @@ -4,13 +4,9 @@ apiVersion: v1
@@ -89,3 +89,24 @@ index 93d0b57..54cb7e4 100644
          imagePullPolicy: IfNotPresent
          livenessProbe:
            httpGet:
+@@ -262,9 +242,6 @@ spec:
+         volumeMounts:
+         - mountPath: /dev
+           name: device
+-        - mountPath: /etc/kubernetes/
+-          name: k8s-cfg
+-          readOnly: true
+         - mountPath: /csi
+           name: csi-socket-dir
+         - mountPath: /var/lib/kubelet/
+@@ -404,10 +381,6 @@ spec:
+           path: /dev
+           type: Directory
+         name: device
+-      - hostPath:
+-          path: /etc/kubernetes/
+-          type: DirectoryOrCreate
+-        name: k8s-cfg
+       - hostPath:
+           path: /var/lib/kubelet/plugins/localdisk.csi.acstor.io/
+           type: DirectoryOrCreate
diff --git a/pkg/featuregates/featuregates.go b/pkg/featuregates/featuregates.go
index 0d69c2d..9885910 100644
--- a/pkg/featuregates/featuregates.go
+++ b/pkg/featuregates/featuregates.go
@@ -25,8 +25,9 @@ import (
 var (
 	// FeatureGates is a map that holds the feature gate names and their default values for Kaito.
 	FeatureGates = map[string]bool{
-		consts.FeatureFlagVLLM:            true,
-		consts.FeatureFlagEnsureNodeClass: false,
+		consts.FeatureFlagVLLM:                        true,
+		consts.FeatureFlagEnsureNodeClass:             false,
+		consts.FeatureFlagDisableNodeAutoProvisioning: false,
 		//	Add more feature gates here
 	}
 )
diff --git a/pkg/ragengine/controllers/preset-rag.go b/pkg/ragengine/controllers/preset_rag.go
similarity index 100%
rename from pkg/ragengine/controllers/preset-rag.go
rename to pkg/ragengine/controllers/preset_rag.go
diff --git a/pkg/utils/common-preset.go b/pkg/utils/common_preset.go
similarity index 100%
rename from pkg/utils/common-preset.go
rename to pkg/utils/common_preset.go
diff --git a/pkg/utils/consts/consts.go b/pkg/utils/consts/consts.go
index 26ea55f..80544d7 100644
--- a/pkg/utils/consts/consts.go
+++ b/pkg/utils/consts/consts.go
@@ -31,8 +31,9 @@ const (
 	NvidiaGPU                     = "nvidia.com/gpu"
 
 	// Feature flags
-	FeatureFlagVLLM            = "vLLM"
-	FeatureFlagEnsureNodeClass = "ensureNodeClass"
+	FeatureFlagVLLM                        = "vLLM"
+	FeatureFlagEnsureNodeClass             = "ensureNodeClass"
+	FeatureFlagDisableNodeAutoProvisioning = "disableNodeAutoProvisioning"
 
 	// Nodeclaim related consts
 	KaitoNodePoolName             = "kaito"
diff --git a/pkg/workspace/controllers/workspace_controller.go b/pkg/workspace/controllers/workspace_controller.go
index 8409eb0..1d573a9 100644
--- a/pkg/workspace/controllers/workspace_controller.go
+++ b/pkg/workspace/controllers/workspace_controller.go
@@ -47,6 +47,7 @@ import (
 	karpenterv1 "sigs.k8s.io/karpenter/pkg/apis/v1"
 
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
+	"github.com/kaito-project/kaito/pkg/featuregates"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
 	"github.com/kaito-project/kaito/pkg/utils/nodeclaim"
@@ -360,6 +361,11 @@ func (c *WorkspaceReconciler) applyWorkspaceResource(ctx context.Context, wObj *
 	newNodesCount := lo.FromPtr(wObj.Resource.Count) - len(selectedNodes)
 
 	if newNodesCount > 0 {
+		// Check if node auto-provisioning is disabled
+		if featuregates.FeatureGates[consts.FeatureFlagDisableNodeAutoProvisioning] {
+			return fmt.Errorf("node auto-provisioning is disabled but insufficient nodes available: need %d nodes, have %d selected nodes", lo.FromPtr(wObj.Resource.Count), len(selectedNodes))
+		}
+
 		klog.InfoS("need to create more nodes", "NodeCount", newNodesCount)
 		if err := c.updateStatusConditionIfNotMatch(ctx, wObj,
 			kaitov1beta1.ConditionTypeNodeClaimStatus, metav1.ConditionUnknown,
@@ -419,46 +425,77 @@ func (c *WorkspaceReconciler) applyWorkspaceResource(ctx context.Context, wObj *
 	return nil
 }
 
-func (c *WorkspaceReconciler) getAllQualifiedNodes(ctx context.Context, wObj *kaitov1beta1.Workspace) ([]*corev1.Node, error) {
-	var qualifiedNodes []*corev1.Node
+func nodeIsReadyAndNotDeleting(node corev1.Node) bool {
 
-	nodeList, err := resources.ListNodes(ctx, c.Client, wObj.Resource.LabelSelector.MatchLabels)
-	if err != nil {
-		return nil, err
+	// skip nodes that are being deleted
+	if node.DeletionTimestamp != nil {
+		return false
 	}
 
-	if len(nodeList.Items) == 0 {
-		klog.InfoS("no current nodes match the workspace resource spec", "workspace", klog.KObj(wObj))
-		return nil, nil
-	}
+	// skip nodes that are not ready
+	_, statusRunning := lo.Find(node.Status.Conditions, func(condition corev1.NodeCondition) bool {
+		return condition.Type == corev1.NodeReady && condition.Status == corev1.ConditionTrue
+	})
 
-	preferredNodeSet := sets.New(wObj.Resource.PreferredNodes...)
+	return statusRunning
+}
+
+func (c *WorkspaceReconciler) getAllQualifiedNodes(ctx context.Context, wObj *kaitov1beta1.Workspace) ([]*corev1.Node, error) {
+	var qualifiedNodes []*corev1.Node
 
-	for index := range nodeList.Items {
-		nodeObj := nodeList.Items[index]
-		// skip nodes that are being deleted
-		if nodeObj.DeletionTimestamp != nil {
-			continue
+	if featuregates.FeatureGates[consts.FeatureFlagDisableNodeAutoProvisioning] {
+		for _, name := range wObj.Resource.PreferredNodes {
+			// Get the node with a client call
+			node, err := resources.GetNode(ctx, name, c.Client)
+			if err != nil {
+				klog.ErrorS(err, "failed to get preferred node", "node", name)
+				continue
+			}
+
+			if nodeIsReadyAndNotDeleting(*node) {
+				// Check that all labels in wObj.Resource.LabelSelector.MatchLabels are present in the node's labels.
+				allLabelsMatch := true
+				for k, v := range wObj.Resource.LabelSelector.MatchLabels {
+					val, ok := node.Labels[k]
+					if !ok || v != val {
+						allLabelsMatch = false
+						break
+					}
+				}
+				if allLabelsMatch {
+					qualifiedNodes = append(qualifiedNodes, node)
+				}
+			}
 		}
 
-		// skip nodes that are not ready
-		_, statusRunning := lo.Find(nodeObj.Status.Conditions, func(condition corev1.NodeCondition) bool {
-			return condition.Type == corev1.NodeReady && condition.Status == corev1.ConditionTrue
-		})
-		if !statusRunning {
-			continue
+		if len(qualifiedNodes) < *wObj.Resource.Count {
+			return nil, fmt.Errorf("when node auto-provisioning is disabled, at least %d preferred nodes must match the label selector and be ready and not deleting, only have %d", *wObj.Resource.Count, len(qualifiedNodes))
+		}
+	} else {
+		nodeList, err := resources.ListNodes(ctx, c.Client, wObj.Resource.LabelSelector.MatchLabels)
+		if err != nil {
+			return nil, err
 		}
 
-		// match the preferred node
-		if preferredNodeSet.Has(nodeObj.Name) {
-			qualifiedNodes = append(qualifiedNodes, lo.ToPtr(nodeObj))
-			continue
+		if len(nodeList.Items) == 0 {
+			klog.InfoS("no current nodes match the workspace resource spec", "workspace", klog.KObj(wObj))
+			return nil, nil
 		}
 
-		// match the instanceType
-		if len(wObj.Resource.PreferredNodes) == 0 { // don't match in perferred nodes mode
-			if nodeObj.Labels[corev1.LabelInstanceTypeStable] == wObj.Resource.InstanceType {
-				qualifiedNodes = append(qualifiedNodes, lo.ToPtr(nodeObj))
+		preferredNodeSet := sets.New(wObj.Resource.PreferredNodes...)
+		for index := range nodeList.Items {
+			node := nodeList.Items[index]
+			if nodeIsReadyAndNotDeleting(node) {
+				// match the preferred node
+				if preferredNodeSet.Has(node.Name) {
+					qualifiedNodes = append(qualifiedNodes, lo.ToPtr(node))
+					continue
+				}
+
+				// match the instanceType
+				if node.Labels[corev1.LabelInstanceTypeStable] == wObj.Resource.InstanceType {
+					qualifiedNodes = append(qualifiedNodes, lo.ToPtr(node))
+				}
 			}
 		}
 	}
diff --git a/presets/ragengine/config.py b/presets/ragengine/config.py
index ae38585..5aa2aa0 100644
--- a/presets/ragengine/config.py
+++ b/presets/ragengine/config.py
@@ -25,14 +25,22 @@ import os
 """
 
 # Embedding configuration
-EMBEDDING_SOURCE_TYPE = os.getenv("EMBEDDING_SOURCE_TYPE", "local")  # Determines local or remote embedding source
+EMBEDDING_SOURCE_TYPE = os.getenv(
+    "EMBEDDING_SOURCE_TYPE", "local"
+)  # Determines local or remote embedding source
 
 # Local embedding model
-LOCAL_EMBEDDING_MODEL_ID = os.getenv("LOCAL_EMBEDDING_MODEL_ID", "BAAI/bge-small-en-v1.5")
+LOCAL_EMBEDDING_MODEL_ID = os.getenv(
+    "LOCAL_EMBEDDING_MODEL_ID", "BAAI/bge-small-en-v1.5"
+)
 
 # Remote embedding model (if not local)
-REMOTE_EMBEDDING_URL = os.getenv("REMOTE_EMBEDDING_URL", "http://localhost:5000/embedding")
-REMOTE_EMBEDDING_ACCESS_SECRET = os.getenv("REMOTE_EMBEDDING_ACCESS_SECRET", "default-access-secret")
+REMOTE_EMBEDDING_URL = os.getenv(
+    "REMOTE_EMBEDDING_URL", "http://localhost:5000/embedding"
+)
+REMOTE_EMBEDDING_ACCESS_SECRET = os.getenv(
+    "REMOTE_EMBEDDING_ACCESS_SECRET", "default-access-secret"
+)
 
 """
 =========================================================================
@@ -41,17 +49,25 @@ REMOTE_EMBEDDING_ACCESS_SECRET = os.getenv("REMOTE_EMBEDDING_ACCESS_SECRET", "de
 # Reranking Configuration
 # For now we support simple LLMReranker, future additions would include
 # FlagEmbeddingReranker, SentenceTransformerReranker, CohereReranker
-LLM_RERANKER_BATCH_SIZE = int(os.getenv("LLM_RERANKER_BATCH_SIZE", 5))  # Default LLM batch size
-LLM_RERANKER_TOP_N = int(os.getenv("LLM_RERANKER_TOP_N", 3))  # Default top 3 reranked nodes
+LLM_RERANKER_BATCH_SIZE = int(
+    os.getenv("LLM_RERANKER_BATCH_SIZE", 5)
+)  # Default LLM batch size
+LLM_RERANKER_TOP_N = int(
+    os.getenv("LLM_RERANKER_TOP_N", 3)
+)  # Default top 3 reranked nodes
 
 """
 =========================================================================
 """
 
 # LLM (Large Language Model) configuration
-LLM_INFERENCE_URL = os.getenv("LLM_INFERENCE_URL", "http://localhost:5000/v1/completions")
+LLM_INFERENCE_URL = os.getenv(
+    "LLM_INFERENCE_URL", "http://localhost:5000/v1/completions"
+)
 LLM_ACCESS_SECRET = os.getenv("LLM_ACCESS_SECRET", "default-access-secret")
-LLM_CONTEXT_WINDOW = int(os.getenv("LLM_CONTEXT_WINDOW", 64000))  # Default context window size
+LLM_CONTEXT_WINDOW = int(
+    os.getenv("LLM_CONTEXT_WINDOW", 64000)
+)  # Default context window size
 # LLM_RESPONSE_FIELD = os.getenv("LLM_RESPONSE_FIELD", "result")  # Uncomment if needed in the future
 
 """
diff --git a/presets/ragengine/embedding/base.py b/presets/ragengine/embedding/base.py
index b3b159c..ccbbeb1 100644
--- a/presets/ragengine/embedding/base.py
+++ b/presets/ragengine/embedding/base.py
@@ -12,19 +12,20 @@
 # limitations under the License.
 
 
-from typing import List
+import asyncio
 from abc import ABC, abstractmethod
+
 from llama_index.core.embeddings import BaseEmbedding
-import asyncio
+
 
 class BaseEmbeddingModel(BaseEmbedding, ABC):
-    async def _aget_text_embedding(self, text: str) -> List[float]:
+    async def _aget_text_embedding(self, text: str) -> list[float]:
         return await asyncio.to_thread(self._get_text_embedding, text)
 
-    async def _aget_query_embedding(self, query: str) -> List[float]:
+    async def _aget_query_embedding(self, query: str) -> list[float]:
         return await asyncio.to_thread(self._get_query_embedding, query)
 
     @abstractmethod
     def get_embedding_dimension(self) -> int:
         """Returns the embedding dimension for the model."""
-        pass
\ No newline at end of file
+        pass
diff --git a/presets/ragengine/embedding/huggingface_local_embedding.py b/presets/ragengine/embedding/huggingface_local_embedding.py
index 09ec38e..deca7f3 100644
--- a/presets/ragengine/embedding/huggingface_local_embedding.py
+++ b/presets/ragengine/embedding/huggingface_local_embedding.py
@@ -11,41 +11,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List, Union
 import time
 from io import BytesIO
+from typing import Any
+
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from .base import BaseEmbeddingModel
-from ragengine.metrics.helpers import record_embedding_metrics
+
 from ragengine.metrics.prometheus_metrics import (
-    rag_embedding_requests_total, 
-    rag_embedding_latency,
-    STATUS_SUCCESS,
+    MODE_LOCAL,
     STATUS_FAILURE,
-    MODE_LOCAL
+    STATUS_SUCCESS,
+    rag_embedding_latency,
+    rag_embedding_requests_total,
 )
 
+from .base import BaseEmbeddingModel
+
+
 class LocalHuggingFaceEmbedding(HuggingFaceEmbedding, BaseEmbeddingModel):
     """HuggingFace embedding model with metrics collection."""
-    
+
     def _embed_with_retry(
         self,
-        inputs: List[Union[str, BytesIO]],
+        inputs: list[str | BytesIO],
         prompt_name: Any = None,
-    ) -> List[List[float]]:
-        start_time = time.perf_counter() 
+    ) -> list[list[float]]:
+        start_time = time.perf_counter()
         status = STATUS_FAILURE  # Default to failure
-        
+
         try:
             result = super()._embed_with_retry(inputs, prompt_name)
-            status = STATUS_SUCCESS  
+            status = STATUS_SUCCESS
             return result
         except Exception:
-            raise 
+            raise
         finally:
             latency = time.perf_counter() - start_time
             rag_embedding_requests_total.labels(status=status, mode=MODE_LOCAL).inc()
-            rag_embedding_latency.labels(status=status, mode=MODE_LOCAL).observe(latency)
+            rag_embedding_latency.labels(status=status, mode=MODE_LOCAL).observe(
+                latency
+            )
 
     def get_embedding_dimension(self) -> int:
         """Infers the embedding dimension by making a local call to get the embedding of a dummy text."""
@@ -53,4 +58,4 @@ class LocalHuggingFaceEmbedding(HuggingFaceEmbedding, BaseEmbeddingModel):
         embedding = self.get_text_embedding(dummy_input)
         if embedding is None:
             raise ValueError("Unable to get embedding dimension due to None embedding.")
-        return len(embedding)
\ No newline at end of file
+        return len(embedding)
diff --git a/presets/ragengine/embedding/remote_embedding.py b/presets/ragengine/embedding/remote_embedding.py
index 7abd226..61754c2 100644
--- a/presets/ragengine/embedding/remote_embedding.py
+++ b/presets/ragengine/embedding/remote_embedding.py
@@ -11,12 +11,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from typing import Any
+
 import requests
-import json
-from .base import BaseEmbeddingModel
+
 from ragengine.metrics.helpers import record_embedding_metrics
 
+from .base import BaseEmbeddingModel
+
+
 class RemoteEmbeddingModel(BaseEmbeddingModel):
     def __init__(self, model_url: str, api_key: str, /, **data: Any):
         """
@@ -35,14 +39,14 @@ class RemoteEmbeddingModel(BaseEmbeddingModel):
         """Returns the text embedding for a given input string."""
         headers = {
             "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "inputs": text
+            "Content-Type": "application/json",
         }
+        payload = {"inputs": text}
 
         try:
-            response = requests.post(self.model_url, headers=headers, data=json.dumps(payload))
+            response = requests.post(
+                self.model_url, headers=headers, data=json.dumps(payload)
+            )
             response.raise_for_status()  # Raise an HTTPError for bad responses
             embedding = response.json()  # Assumes the API returns JSON
             if isinstance(embedding, list):
@@ -59,5 +63,5 @@ class RemoteEmbeddingModel(BaseEmbeddingModel):
         """Infers the embedding dimension by making a remote call to get the embedding of a dummy text."""
         dummy_input = "This is a dummy sentence."
         embedding = self._get_text_embedding(dummy_input)
-        
+
         return len(embedding)
diff --git a/presets/ragengine/inference/inference.py b/presets/ragengine/inference/inference.py
index f9f3629..2225da9 100644
--- a/presets/ragengine/inference/inference.py
+++ b/presets/ragengine/inference/inference.py
@@ -12,25 +12,38 @@
 # limitations under the License.
 
 
-from pydantic import PrivateAttr
-import logging
 import asyncio
-import httpx
-from typing import Any, Sequence
-from llama_index.core.llms import CustomLLM, CompletionResponse, LLMMetadata, CompletionResponseGen, ChatMessage, ChatResponse
-from llama_index.core.llms.callbacks import llm_completion_callback, llm_chat_callback
-import requests
-from requests.exceptions import HTTPError
-from urllib.parse import urlparse, urljoin
-from ragengine.config import LLM_INFERENCE_URL, LLM_ACCESS_SECRET, LLM_CONTEXT_WINDOW #, LLM_RESPONSE_FIELD
-from ragengine.models import ChatCompletionResponse
-from fastapi import HTTPException
 import concurrent.futures
 import json
+import logging
+from collections.abc import Sequence
+from typing import Any
+from urllib.parse import urljoin, urlparse
 
+import httpx
+import requests
+from fastapi import HTTPException
+from llama_index.core.llms import (
+    ChatMessage,
+    ChatResponse,
+    CompletionResponse,
+    CompletionResponseGen,
+    CustomLLM,
+    LLMMetadata,
+)
+from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback
 from openai.types.chat import (
     CompletionCreateParams,
 )
+from pydantic import PrivateAttr
+from requests.exceptions import HTTPError
+
+from ragengine.config import (
+    LLM_ACCESS_SECRET,
+    LLM_CONTEXT_WINDOW,
+    LLM_INFERENCE_URL,
+)
+from ragengine.models import ChatCompletionResponse
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -40,20 +53,21 @@ OPENAI_URL_PREFIX = "https://api.openai.com"
 HUGGINGFACE_URL_PREFIX = "https://api-inference.huggingface.co"
 DEFAULT_HEADERS = {
     "Authorization": f"Bearer {LLM_ACCESS_SECRET}",
-    "Content-Type": "application/json"
+    "Content-Type": "application/json",
 }
-DEFAULT_HTTP_TIMEOUT = 300.0 # Seconds
+DEFAULT_HTTP_TIMEOUT = 300.0  # Seconds
 DEFAULT_HTTP_SUCCESS_CODE = 200
 
+
 class Inference(CustomLLM):
     params: dict = {}
     _default_model: str = None
     _default_max_model_len: int = None
     _model_retrieval_attempted: bool = False
-    _async_http_client : httpx.AsyncClient = PrivateAttr(default=None)
+    _async_http_client: httpx.AsyncClient = PrivateAttr(default=None)
 
     async def _get_httpx_client(self):
-        """ Lazily initializes the HTTP client on first request. """
+        """Lazily initializes the HTTP client on first request."""
         if self._async_http_client is None:
             self._async_http_client = httpx.AsyncClient(timeout=DEFAULT_HTTP_TIMEOUT)
         return self._async_http_client
@@ -75,50 +89,69 @@ class Inference(CustomLLM):
             return future.result()
 
     @llm_completion_callback()
-    def complete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse:
+    def complete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponse:
         # Required implementation - Only called by LlamaIndex reranker because LLMRerank library doesn't use async call
         try:
-            result = self.run_async_coroutine(self.acomplete(prompt, formatted=formatted, **kwargs))
+            result = self.run_async_coroutine(
+                self.acomplete(prompt, formatted=formatted, **kwargs)
+            )
             if result.text == "Empty Response":
-                logger.error("LLMRerank Request returned an unparsable or invalid response")
-                raise HTTPException(status_code=422, detail="Rerank operation failed: Invalid response from LLM. This feature is experimental.")
+                logger.error(
+                    "LLMRerank Request returned an unparsable or invalid response"
+                )
+                raise HTTPException(
+                    status_code=422,
+                    detail="Rerank operation failed: Invalid response from LLM. This feature is experimental.",
+                )
             return result
         except HTTPException as http_exc:
             # If it's already an HTTPException (e.g., 422), re-raise it as is
             raise http_exc
         except Exception as e:
             logger.error(f"Unexpected exception in complete(): {e}")
-            raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"An unexpected error occurred: {str(e)}"
+            )
 
     @llm_completion_callback()
-    async def acomplete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse:
+    async def acomplete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponse:
         try:
             return await self._async_completions(prompt, **kwargs, **self.params)
         except HTTPException as http_exc:
             raise http_exc
         except Exception as e:
             logger.error(f"Unexpected exception in acomplete(): {e}")
-            raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"An unexpected error occurred: {str(e)}"
+            )
         finally:
             # Clear params after the completion is done
             self.params = {}
-    
+
     @llm_chat_callback()
     def chat(
         self,
         messages: Sequence[ChatMessage],
         **kwargs: Any,
     ) -> ChatResponse:
-        """ Perform a chat completion request. """
+        """Perform a chat completion request."""
         try:
-            logger.info(f"Sending chat request to {LLM_INFERENCE_URL} with messages: {messages}, and args: {kwargs}")
+            logger.info(
+                f"Sending chat request to {LLM_INFERENCE_URL} with messages: {messages}, and args: {kwargs}"
+            )
             return self.run_async_coroutine(self.achat(messages, **kwargs))
         except HTTPException as http_exc:
             logger.error(f"HTTP exception during chat(): {http_exc.detail}")
             raise http_exc
         except Exception as e:
             logger.error(f"Unexpected exception in chat(): {e}")
-            raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"An unexpected error occurred: {str(e)}"
+            )
 
     @llm_chat_callback()
     async def achat(
@@ -126,16 +159,21 @@ class Inference(CustomLLM):
         messages: Sequence[ChatMessage],
         **kwargs: Any,
     ) -> ChatResponse:
-        """ Perform an asynchronous chat completion request. """
+        """Perform an asynchronous chat completion request."""
         try:
             base_model, base_max_len = self._get_default_model_info()
             req = {
                 "model": self.get_param("model", base_model),
                 "max_tokens": self.get_param("max_tokens", base_max_len),
-                "messages": [{
-                    "role": message.role,
-                    "content": message.content if isinstance(message.content, str) else json.dumps(message.content)
-                } for message in messages if message.content is not None and message.content != ""
+                "messages": [
+                    {
+                        "role": message.role,
+                        "content": message.content
+                        if isinstance(message.content, str)
+                        else json.dumps(message.content),
+                    }
+                    for message in messages
+                    if message.content is not None and message.content != ""
                 ],
             }
 
@@ -143,60 +181,90 @@ class Inference(CustomLLM):
             for key, value in self.params.items():
                 if key not in req:
                     req[key] = value
-            
-            resp =  await self._async_post_request_raw(data=req, headers=DEFAULT_HEADERS)
+
+            resp = await self._async_post_request_raw(data=req, headers=DEFAULT_HEADERS)
             return ChatResponse(
                 logprobs=resp.get("logprobs", None),
                 delta=resp.get("delta", None),
                 raw=resp,
-                message=ChatMessage(content=resp.get("choices", [{}])[0].get("message", {}).get("content", ""))
+                message=ChatMessage(
+                    content=resp.get("choices", [{}])[0]
+                    .get("message", {})
+                    .get("content", "")
+                ),
             )
         except HTTPException as http_exc:
             logger.error(f"HTTP exception during achat(): {http_exc.detail}")
             raise http_exc
         except Exception as e:
             logger.error(f"Unexpected exception in achat(): {e}")
-            raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"An unexpected error occurred: {str(e)}"
+            )
         finally:
             # Clear params after the completion is done
             self.params = {}
-    
-    async def chat_completions_passthrough(self, chatCompletionsRequest: CompletionCreateParams, **kwargs: Any) -> ChatCompletionResponse:
+
+    async def chat_completions_passthrough(
+        self, chatCompletionsRequest: CompletionCreateParams, **kwargs: Any
+    ) -> ChatCompletionResponse:
         try:
             if "/chat/completions" not in LLM_INFERENCE_URL:
                 # If the URL does not support chat completions, raise an error
-                raise HTTPException(status_code=400, detail=f"Chat completions not supported through endpoint {LLM_INFERENCE_URL}.")
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Chat completions not supported through endpoint {LLM_INFERENCE_URL}.",
+                )
 
             client = await self._get_httpx_client()
-            response = await client.post(LLM_INFERENCE_URL, json=chatCompletionsRequest, headers=DEFAULT_HEADERS)
+            response = await client.post(
+                LLM_INFERENCE_URL, json=chatCompletionsRequest, headers=DEFAULT_HEADERS
+            )
             response.raise_for_status()  # Raise an exception for HTTP errors
             response_data = response.json()
             # Convert to ChatCompletionResponse with source_nodes=None for passthrough
             return ChatCompletionResponse(**response_data, source_nodes=None)
         except HTTPException as http_exc:
-            logger.error(f"HTTP exception during chat completions passthrough: {http_exc.detail}")
+            logger.error(
+                f"HTTP exception during chat completions passthrough: {http_exc.detail}"
+            )
             raise http_exc
         except httpx.HTTPStatusError as e:
-            logger.error(f"HTTP error {e.response.status_code} during POST request to {LLM_INFERENCE_URL}: {e.response.text}")
-            raise HTTPException(status_code=e.response.status_code, detail=f"{str(e.response.content)}")
+            logger.error(
+                f"HTTP error {e.response.status_code} during POST request to {LLM_INFERENCE_URL}: {e.response.text}"
+            )
+            raise HTTPException(
+                status_code=e.response.status_code, detail=f"{str(e.response.content)}"
+            )
         except httpx.RequestError as e:
             logger.error(f"Error during POST request to {LLM_INFERENCE_URL}: {e}")
-            raise HTTPException(status_code=500, detail=f"Error during POST request: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"Error during POST request: {str(e)}"
+            )
         except Exception as e:
             logger.error(f"Unexpected error during POST request: {e}")
-            raise HTTPException(status_code=500, detail=f"Error during POST request: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"Error during POST request: {str(e)}"
+            )
 
-    async def _async_completions(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+    async def _async_completions(
+        self, prompt: str, **kwargs: Any
+    ) -> CompletionResponse:
         model_name, model_max_len = self._get_default_model_info()
         if kwargs.get("model"):
             model_name = kwargs.pop("model")
         data = {"prompt": prompt, **kwargs}
         if model_name:
-            data["model"] = model_name # Include the model only if it is not None
-        if model_max_len and data.get("max_tokens"):
-            if data["max_tokens"] > model_max_len:
-                logger.error(f"Requested max_tokens ({data['max_tokens']}) exceeds model's max length ({model_max_len}).")
-                # vLLM will raise error ({"object":"error","message":"This model's maximum context length is 131072 tokens. However, you requested 500500500500505361 tokens (361 in the messages, 500500500500505000 in the completion). Please reduce the length of the messages or completion.","type":"BadRequestError","param":null,"code":400})
+            data["model"] = model_name  # Include the model only if it is not None
+        if (
+            model_max_len
+            and data.get("max_tokens")
+            and data["max_tokens"] > model_max_len
+        ):
+            logger.error(
+                f"Requested max_tokens ({data['max_tokens']}) exceeds model's max length ({model_max_len})."
+            )
+            # vLLM will raise error ({"object":"error","message":"This model's maximum context length is 131072 tokens. However, you requested 500500500500505361 tokens (361 in the messages, 500500500500505000 in the completion). Please reduce the length of the messages or completion.","type":"BadRequestError","param":null,"code":400})
 
         # DEBUG: Call the debugging function
         # self._debug_curl_command(data)
@@ -213,11 +281,17 @@ class Inference(CustomLLM):
                     f"Potential issue with 'model' parameter in API response. "
                     f"Response: {str(e)}. Attempting to update the model name as a mitigation..."
                 )
-                self._default_model, self._default_max_model_len = self._fetch_default_model_info()  # Fetch default model dynamically
+                self._default_model, self._default_max_model_len = (
+                    self._fetch_default_model_info()
+                )  # Fetch default model dynamically
                 if self._default_model:
-                    logger.info(f"Default model '{self._default_model}' fetched successfully. Retrying request...")
+                    logger.info(
+                        f"Default model '{self._default_model}' fetched successfully. Retrying request..."
+                    )
                     data["model"] = self._default_model
-                    resp = await self._async_post_request_raw(data, headers=DEFAULT_HEADERS)
+                    resp = await self._async_post_request_raw(
+                        data, headers=DEFAULT_HEADERS
+                    )
                     return self._completions_json_to_response(resp)
                 else:
                     logger.error("Failed to fetch a default model. Aborting retry.")
@@ -225,7 +299,7 @@ class Inference(CustomLLM):
         except Exception as e:
             logger.error(f"An unexpected error occurred: {e}")
             raise
-    
+
     def _completions_json_to_response(self, response_json: dict) -> CompletionResponse:
         """
         Converts the JSON response from the completions API to a CompletionResponse object.
@@ -257,7 +331,9 @@ class Inference(CustomLLM):
             return None, None
 
         except Exception as e:
-            logger.error(f"Error fetching models from {models_url}: {e}. \"model\" parameter will not be included with inference call.")
+            logger.error(
+                f'Error fetching models from {models_url}: {e}. "model" parameter will not be included with inference call.'
+            )
             return None
 
     def _get_default_model_info(self) -> (str, int):
@@ -266,7 +342,9 @@ class Inference(CustomLLM):
         """
         if not self._default_model and not self._model_retrieval_attempted:
             self._model_retrieval_attempted = True
-            self._default_model, self._default_max_model_len = self._fetch_default_model_info()
+            self._default_model, self._default_max_model_len = (
+                self._fetch_default_model_info()
+            )
         return self._default_model, self._default_max_model_len
 
     async def _async_post_request_raw(self, data: dict, headers: dict):
@@ -276,7 +354,9 @@ class Inference(CustomLLM):
             response.raise_for_status()  # Raise an exception for HTTP errors
             return response.json()
         except httpx.HTTPStatusError as e:
-            logger.error(f"HTTP error {e.response.status_code} during POST request to {LLM_INFERENCE_URL}: {e.response.text}")
+            logger.error(
+                f"HTTP error {e.response.status_code} during POST request to {LLM_INFERENCE_URL}: {e.response.text}"
+            )
             raise
         except httpx.RequestError as e:
             logger.error(f"Error during POST request to {LLM_INFERENCE_URL}: {e}")
@@ -290,14 +370,20 @@ class Inference(CustomLLM):
         Constructs and prints the equivalent curl command for debugging purposes.
         """
         import json
+
         # Construct curl command
         curl_command = (
-                f"curl -X POST {LLM_INFERENCE_URL} "
-                + " ".join([f'-H "{key}: {value}"' for key, value in {
-            "Authorization": f"Bearer {LLM_ACCESS_SECRET}",
-            "Content-Type": "application/json"
-        }.items()])
-                + f" -d '{json.dumps(data)}'"
+            f"curl -X POST {LLM_INFERENCE_URL} "
+            + " ".join(
+                [
+                    f'-H "{key}: {value}"'
+                    for key, value in {
+                        "Authorization": f"Bearer {LLM_ACCESS_SECRET}",
+                        "Content-Type": "application/json",
+                    }.items()
+                ]
+            )
+            + f" -d '{json.dumps(data)}'"
         )
         logger.info("Equivalent curl command:")
         logger.info(curl_command)
@@ -313,6 +399,6 @@ class Inference(CustomLLM):
         )
 
     async def aclose(self):
-        """ Closes the HTTP client when shutting down. """
+        """Closes the HTTP client when shutting down."""
         if self._async_http_client:
             await self._async_http_client.aclose()
diff --git a/presets/ragengine/main.py b/presets/ragengine/main.py
index 2140c85..ceaa427 100644
--- a/presets/ragengine/main.py
+++ b/presets/ragengine/main.py
@@ -12,89 +12,120 @@
 # limitations under the License.
 
 
-from typing import List, Optional
 import json
+import os
 import time
-from vector_store_manager.manager import VectorStoreManager
+from urllib.parse import unquote
+
 from embedding.huggingface_local_embedding import LocalHuggingFaceEmbedding
 from embedding.remote_embedding import RemoteEmbeddingModel
 from fastapi import FastAPI, HTTPException, Query, Request
-from models import (IndexRequest, ListDocumentsResponse, UpdateDocumentRequest,
-                    QueryRequest, QueryResponse, Document, HealthStatus, DeleteDocumentRequest,
-                    DeleteDocumentResponse, UpdateDocumentResponse, ChatCompletionResponse)
-
-from vector_store.faiss_store import FaissVectorStoreHandler
+from models import (
+    ChatCompletionResponse,
+    DeleteDocumentRequest,
+    DeleteDocumentResponse,
+    Document,
+    HealthStatus,
+    IndexRequest,
+    ListDocumentsResponse,
+    QueryRequest,
+    QueryResponse,
+    UpdateDocumentRequest,
+    UpdateDocumentResponse,
+)
 
-from ragengine.config import (REMOTE_EMBEDDING_URL, REMOTE_EMBEDDING_ACCESS_SECRET,
-                              EMBEDDING_SOURCE_TYPE, LOCAL_EMBEDDING_MODEL_ID, DEFAULT_VECTOR_DB_PERSIST_DIR)
-from urllib.parse import unquote
-import os
 # Import Prometheus client for metrics collection
-from prometheus_client import generate_latest, CONTENT_TYPE_LATEST
-import time
+from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
 from starlette.responses import Response
+from vector_store.faiss_store import FaissVectorStoreHandler
+from vector_store_manager.manager import VectorStoreManager
+
+from ragengine.config import (
+    DEFAULT_VECTOR_DB_PERSIST_DIR,
+    EMBEDDING_SOURCE_TYPE,
+    LOCAL_EMBEDDING_MODEL_ID,
+    REMOTE_EMBEDDING_ACCESS_SECRET,
+    REMOTE_EMBEDDING_URL,
+)
 from ragengine.metrics.prometheus_metrics import (
-    rag_query_latency,
-    rag_query_requests_total,
+    MODE_LOCAL,
+    MODE_REMOTE,
+    STATUS_FAILURE,
+    STATUS_SUCCESS,
+    e2e_request_latency_seconds,
+    e2e_request_total,
+    num_requests_running,
+    rag_avg_source_score,
     rag_chat_latency,
     rag_chat_requests_total,
+    rag_delete_latency,
+    rag_delete_requests_total,
     rag_index_latency,
     rag_index_requests_total,
-    rag_indexes_latency,
-    rag_indexes_requests_total,
-    rag_indexes_document_latency,
-    rag_indexes_document_requests_total,
     rag_indexes_delete_document_latency,
     rag_indexes_delete_document_requests_total,
+    rag_indexes_document_latency,
+    rag_indexes_document_requests_total,
+    rag_indexes_latency,
+    rag_indexes_requests_total,
     rag_indexes_update_document_latency,
     rag_indexes_update_document_requests_total,
-    rag_persist_latency,
-    rag_persist_requests_total,
     rag_load_latency,
     rag_load_requests_total,
-    rag_delete_latency,
-    rag_delete_requests_total,
-    e2e_request_total,
-    e2e_request_latency_seconds,
-    num_requests_running,
     rag_lowest_source_score,
-    rag_avg_source_score,
-    STATUS_SUCCESS,
-    STATUS_FAILURE,
-    MODE_LOCAL,
-    MODE_REMOTE
+    rag_persist_latency,
+    rag_persist_requests_total,
+    rag_query_latency,
+    rag_query_requests_total,
 )
 
+# Import Prometheus client for metrics collection
+
+
 app = FastAPI()
+
+
 @app.middleware("http")
 async def track_requests(request: Request, call_next):
-    tracked_paths = ["/query", "/index", "/indexes", "/persist", "/load", "/v1/chat/completions"]
+    tracked_paths = [
+        "/query",
+        "/index",
+        "/indexes",
+        "/persist",
+        "/load",
+        "/v1/chat/completions",
+    ]
 
     should_track = any(request.url.path.startswith(path) for path in tracked_paths)
-    
+
     if not should_track:
         return await call_next(request)
 
-    num_requests_running.inc() 
-    start_time = time.perf_counter() 
-    status = STATUS_FAILURE # Default status
-    
+    num_requests_running.inc()
+    start_time = time.perf_counter()
+    status = STATUS_FAILURE  # Default status
+
     try:
         response = await call_next(request)
         status = STATUS_SUCCESS
         return response
-    except Exception as e:
+    except Exception:
         raise
     finally:
         num_requests_running.dec()
-        e2e_request_latency_seconds.labels(status=status).observe(time.perf_counter() - start_time)
+        e2e_request_latency_seconds.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
         e2e_request_total.labels(status=status).inc()
 
+
 # Initialize embedding model
 if EMBEDDING_SOURCE_TYPE.lower() == MODE_LOCAL:
     embedding_manager = LocalHuggingFaceEmbedding(LOCAL_EMBEDDING_MODEL_ID)
 elif EMBEDDING_SOURCE_TYPE.lower() == MODE_REMOTE:
-    embedding_manager = RemoteEmbeddingModel(REMOTE_EMBEDDING_URL, REMOTE_EMBEDDING_ACCESS_SECRET)
+    embedding_manager = RemoteEmbeddingModel(
+        REMOTE_EMBEDDING_URL, REMOTE_EMBEDDING_ACCESS_SECRET
+    )
 else:
     raise ValueError("Invalid Embedding Type Specified (Must be Local or Remote)")
 
@@ -105,6 +136,7 @@ vector_store_handler = FaissVectorStoreHandler(embedding_manager)
 # Initialize RAG operations
 rag_ops = VectorStoreManager(vector_store_handler)
 
+
 @app.get("/metrics", tags=["Monitoring"])
 async def metrics():
     """
@@ -114,10 +146,11 @@ async def metrics():
     try:
         # Serialize and return all Prometheus metrics collected in this process
         return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)
-    
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @app.get(
     "/health",
     response_model=HealthStatus,
@@ -137,19 +170,24 @@ async def metrics():
 def health_check():
     try:
         if embedding_manager is None:
-            raise HTTPException(status_code=500, detail="Embedding manager not initialized")
-        
+            raise HTTPException(
+                status_code=500, detail="Embedding manager not initialized"
+            )
+
         if rag_ops is None:
-            raise HTTPException(status_code=500, detail="RAG operations not initialized")
+            raise HTTPException(
+                status_code=500, detail="RAG operations not initialized"
+            )
 
         return HealthStatus(status="Healthy")
-    
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @app.post(
     "/index",
-    response_model=List[Document],
+    response_model=list[Document],
     summary="Index Documents",
     description="""
     Add documents to an index or create a new index.
@@ -181,12 +219,12 @@ def health_check():
 async def index_documents(request: IndexRequest):
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
         doc_ids = await rag_ops.index(request.index_name, request.documents)
         documents = [
-            Document(doc_id=doc_id, text=doc.text, metadata=doc.metadata) 
-            for doc_id, doc in zip(doc_ids, request.documents)
+            Document(doc_id=doc_id, text=doc.text, metadata=doc.metadata)
+            for doc_id, doc in zip(doc_ids, request.documents, strict=False)
         ]
         status = STATUS_SUCCESS
         return documents
@@ -197,7 +235,10 @@ async def index_documents(request: IndexRequest):
     finally:
         # Record metrics once in finally block
         rag_index_requests_total.labels(status=status).inc()
-        rag_index_latency.labels(status=status).observe(time.perf_counter() - start_time)
+        rag_index_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
+
 
 @app.post(
     "/query",
@@ -242,30 +283,34 @@ async def index_documents(request: IndexRequest):
 async def query_index(request: QueryRequest):
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
-        llm_params = request.llm_params or {}  # Default to empty dict if no params provided
-        rerank_params = request.rerank_params or {}  # Default to empty dict if no params provided
-        
+        llm_params = (
+            request.llm_params or {}
+        )  # Default to empty dict if no params provided
+        rerank_params = (
+            request.rerank_params or {}
+        )  # Default to empty dict if no params provided
+
         result_dict = await rag_ops.query(
             request.index_name, request.query, request.top_k, llm_params, rerank_params
         )
-        
+
         result = QueryResponse(
             response=result_dict["response"],
-            source_nodes=result_dict["source_nodes"], 
-            metadata=result_dict.get("metadata", {})
+            source_nodes=result_dict["source_nodes"],
+            metadata=result_dict.get("metadata", {}),
         )
-        
+
         # Record source retrieval quality metrics
         if result.source_nodes and result.response:
             lowest_score_node = min(result.source_nodes, key=lambda x: x.score)
             rag_lowest_source_score.observe(lowest_score_node.score)
-            
+
             scores = [node.score for node in result.source_nodes]
             avg_score = sum(scores) / len(scores)
             rag_avg_source_score.observe(avg_score)
-   
+
         status = STATUS_SUCCESS
         return result
     except HTTPException as http_exc:
@@ -280,7 +325,10 @@ async def query_index(request: QueryRequest):
     finally:
         # Record metrics once in finally block
         rag_query_requests_total.labels(status=status).inc()
-        rag_query_latency.labels(status=status).observe(time.perf_counter() - start_time)
+        rag_query_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
+
 
 @app.post(
     "/v1/chat/completions",
@@ -353,9 +401,10 @@ async def chat_completions(request: dict):
         rag_chat_requests_total.labels(status=status).inc()
         rag_chat_latency.labels(status=status).observe(time.perf_counter() - start_time)
 
+
 @app.get(
     "/indexes",
-    response_model=List[str],
+    response_model=list[str],
     summary="List All Indexes",
     description="""
     Retrieve the names of all available indexes.
@@ -372,9 +421,9 @@ async def chat_completions(request: dict):
 def list_indexes():
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
-        result = rag_ops.list_indexes() 
+        result = rag_ops.list_indexes()
         status = STATUS_SUCCESS
         return result
     except Exception as e:
@@ -382,7 +431,9 @@ def list_indexes():
     finally:
         # Record metrics once in finally block
         rag_indexes_requests_total.labels(status=status).inc()
-        rag_indexes_latency.labels(status=status).observe(time.perf_counter() - start_time)
+        rag_indexes_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
 
 
 @app.get(
@@ -421,17 +472,19 @@ def list_indexes():
 )
 async def list_documents_in_index(
     index_name: str,
-    limit: int = Query(10, ge=1, le=100, description="Maximum number of documents to return"),
+    limit: int = Query(
+        10, ge=1, le=100, description="Maximum number of documents to return"
+    ),
     offset: int = Query(0, ge=0, description="Starting point for the document list"),
-    max_text_length: Optional[int] = Query(
-        1000, 
-        ge=1, 
-        description="Maximum text length to return **per document**. This does not impose a limit on the total length of all documents returned."
+    max_text_length: int | None = Query(
+        1000,
+        ge=1,
+        description="Maximum text length to return **per document**. This does not impose a limit on the total length of all documents returned.",
     ),
-    metadata_filter: Optional[str | None] = Query(
+    metadata_filter: str | None | None = Query(
         None,
-        description="Optional metadata filter to apply when listing documents. This should be a dictionary with key-value pairs to match against document metadata."
-    )
+        description="Optional metadata filter to apply when listing documents. This should be a dictionary with key-value pairs to match against document metadata.",
+    ),
 ):
     """
     Handles URL-encoded index names sent by the client.
@@ -445,14 +498,17 @@ async def list_documents_in_index(
     """
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
         if metadata_filter:
             # Attempt to parse the metadata filter as a JSON string
             try:
                 metadata_filter = json.loads(metadata_filter)
             except json.JSONDecodeError:
-                raise HTTPException(status_code=400, detail="Invalid metadata filter format. Must be a valid JSON string.")
+                raise HTTPException(
+                    status_code=400,
+                    detail="Invalid metadata filter format. Must be a valid JSON string.",
+                )
 
         # Decode the index_name in case it was URL-encoded by the client
         decoded_index_name = unquote(index_name)
@@ -461,13 +517,10 @@ async def list_documents_in_index(
             limit=limit,
             offset=offset,
             max_text_length=max_text_length,
-            metadata_filter=metadata_filter
+            metadata_filter=metadata_filter,
         )
 
-        result = ListDocumentsResponse(
-            documents=documents,
-            count=len(documents)
-        )
+        result = ListDocumentsResponse(documents=documents, count=len(documents))
         status = STATUS_SUCCESS
         return result
     except HTTPException as http_exc:
@@ -477,7 +530,10 @@ async def list_documents_in_index(
     finally:
         # Record metrics once in finally block
         rag_indexes_document_requests_total.labels(status=status).inc()
-        rag_indexes_document_latency.labels(status=status).observe(time.perf_counter() - start_time)
+        rag_indexes_document_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
+
 
 @app.post(
     "/indexes/{index_name}/documents",
@@ -508,11 +564,10 @@ async def update_documents_in_index(
 ):
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
         result = await rag_ops.update_documents(
-            index_name=index_name,
-            documents=request.documents
+            index_name=index_name, documents=request.documents
         )
         status = STATUS_SUCCESS
         return result
@@ -523,7 +578,10 @@ async def update_documents_in_index(
     finally:
         # Record metrics once in finally block
         rag_indexes_update_document_requests_total.labels(status=status).inc()
-        rag_indexes_update_document_latency.labels(status=status).observe(time.perf_counter() - start_time)
+        rag_indexes_update_document_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
+
 
 @app.post(
     "/indexes/{index_name}/documents/delete",
@@ -553,11 +611,10 @@ async def delete_documents_in_index(
 ):
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
         result = await rag_ops.delete_documents(
-            index_name=index_name,
-            doc_ids=request.doc_ids
+            index_name=index_name, doc_ids=request.doc_ids
         )
         status = STATUS_SUCCESS
         return result
@@ -568,7 +625,10 @@ async def delete_documents_in_index(
     finally:
         # Record metrics once in finally block
         rag_indexes_delete_document_requests_total.labels(status=status).inc()
-        rag_indexes_delete_document_latency.labels(status=status).observe(time.perf_counter() - start_time)
+        rag_indexes_delete_document_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
+
 
 @app.post(
     "/persist/{index_name}",
@@ -589,18 +649,25 @@ async def delete_documents_in_index(
       "message": "Successfully persisted index example_index to ./custom_path/example_index."
     }
     ```
-    """
+    """,
 )
 async def persist_index(
-        index_name: str,
-        path: str = Query(DEFAULT_VECTOR_DB_PERSIST_DIR, description="Path where the index will be persisted")
+    index_name: str,
+    path: str = Query(
+        DEFAULT_VECTOR_DB_PERSIST_DIR,
+        description="Path where the index will be persisted",
+    ),
 ):
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
         # Append index to save path to prevent saving conflicts/overwriting
-        path = os.path.join(path, index_name) if path == DEFAULT_VECTOR_DB_PERSIST_DIR else path
+        path = (
+            os.path.join(path, index_name)
+            if path == DEFAULT_VECTOR_DB_PERSIST_DIR
+            else path
+        )
         await rag_ops.persist(index_name, path)
         status = STATUS_SUCCESS
         return {"message": f"Successfully persisted index {index_name} to {path}."}
@@ -611,7 +678,9 @@ async def persist_index(
     finally:
         # Record metrics once in finally block
         rag_persist_requests_total.labels(status=status).inc()
-        rag_persist_latency.labels(status=status).observe(time.perf_counter() - start_time)
+        rag_persist_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
 
 
 @app.post(
@@ -633,20 +702,22 @@ async def persist_index(
       "message": "Successfully loaded index example_index from ./custom_path/example_index."
     }
     ```
-    """
+    """,
 )
 async def load_index(
-        index_name: str,
-        path: Optional[str] = Query(None, description="Path to load the index from"),
-        overwrite: bool = Query(False, description="Overwrite the existing index if it already exists")
+    index_name: str,
+    path: str | None = Query(None, description="Path to load the index from"),
+    overwrite: bool = Query(
+        False, description="Overwrite the existing index if it already exists"
+    ),
 ):
     # If no path is provided, use the default directory joined with index_name.
     if path is None:
         path = os.path.join(DEFAULT_VECTOR_DB_PERSIST_DIR, index_name)
-    
+
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
         await rag_ops.load(index_name, path, overwrite)
         status = STATUS_SUCCESS
@@ -660,6 +731,7 @@ async def load_index(
         rag_load_requests_total.labels(status=status).inc()
         rag_load_latency.labels(status=status).observe(time.perf_counter() - start_time)
 
+
 @app.delete(
     "/indexes/{index_name}",
     summary="Delete the Index",
@@ -677,17 +749,17 @@ async def load_index(
       "message": "Successfully deleted index example_index."
     }
     ```
-    """
+    """,
 )
 async def delete_index(index_name: str):
     """
     Deletes an index by its name.
-    
+
     This function is not exposed via an API endpoint but can be used internally.
     """
     start_time = time.perf_counter()
     status = STATUS_FAILURE  # Default status
-    
+
     try:
         await rag_ops.delete_index(index_name)
         status = STATUS_SUCCESS
@@ -699,16 +771,21 @@ async def delete_index(index_name: str):
     finally:
         # Record metrics once in finally block
         rag_delete_requests_total.labels(status=status).inc()
-        rag_delete_latency.labels(status=status).observe(time.perf_counter()- start_time)
+        rag_delete_latency.labels(status=status).observe(
+            time.perf_counter() - start_time
+        )
+
 
 @app.on_event("shutdown")
 async def shutdown_event():
-    """ Ensure the client is properly closed when the server shuts down. """
+    """Ensure the client is properly closed when the server shuts down."""
     await rag_ops.shutdown()
 
+
 if __name__ == "__main__":
     # DEBUG: Arize Phoenix
     # import llama_index.core
     # llama_index.core.set_global_handler("arize_phoenix")
     import uvicorn
+
     uvicorn.run(app, host="0.0.0.0", port=5000, loop="asyncio")
diff --git a/presets/ragengine/metrics/__init__.py b/presets/ragengine/metrics/__init__.py
index d93603b..84eb2a4 100644
--- a/presets/ragengine/metrics/__init__.py
+++ b/presets/ragengine/metrics/__init__.py
@@ -14,4 +14,4 @@
 
 """
 Metrics collection package.
-""" 
\ No newline at end of file
+"""
diff --git a/presets/ragengine/metrics/helpers.py b/presets/ragengine/metrics/helpers.py
index c8623a1..ee22fd9 100644
--- a/presets/ragengine/metrics/helpers.py
+++ b/presets/ragengine/metrics/helpers.py
@@ -14,32 +14,30 @@
 
 import time
 from functools import wraps
+
 from .prometheus_metrics import (
-    rag_embedding_requests_total,
-    rag_embedding_latency,
-    STATUS_SUCCESS,
+    MODE_REMOTE,
     STATUS_FAILURE,
-    MODE_REMOTE 
+    STATUS_SUCCESS,
+    rag_embedding_latency,
+    rag_embedding_requests_total,
 )
 
+
 def record_embedding_metrics(func):
     """
     Decorator to record embedding metrics for synchronous functions.
     Must be used within a context where the metrics are already imported.
     """
+
     @wraps(func)
     def wrapper(*args, **kwargs):
-        start_time = time.perf_counter()  
+        start_time = time.perf_counter()
         status = STATUS_FAILURE  # Default to failure
-        
+
         try:
             result = func(*args, **kwargs)
-            if result is None:
-                # None result is considered a failure
-                status = STATUS_FAILURE
-            else:
-                # Successful embedding
-                status = STATUS_SUCCESS
+            status = STATUS_FAILURE if result is None else STATUS_SUCCESS
             return result
         except Exception:
             # Status remains STATUS_FAILURE
@@ -48,6 +46,8 @@ def record_embedding_metrics(func):
             # Record metrics once in finally block
             latency = time.perf_counter() - start_time
             rag_embedding_requests_total.labels(status=status, mode=MODE_REMOTE).inc()
-            rag_embedding_latency.labels(status=status, mode=MODE_REMOTE).observe(latency)
+            rag_embedding_latency.labels(status=status, mode=MODE_REMOTE).observe(
+                latency
+            )
 
     return wrapper
diff --git a/presets/ragengine/metrics/prometheus_metrics.py b/presets/ragengine/metrics/prometheus_metrics.py
index 51bf689..8a862e3 100644
--- a/presets/ragengine/metrics/prometheus_metrics.py
+++ b/presets/ragengine/metrics/prometheus_metrics.py
@@ -12,96 +12,175 @@
 # limitations under the License.
 
 
-from prometheus_client import Counter, Histogram, Gauge
+from prometheus_client import Counter, Gauge, Histogram
 
 STATUS_LABEL = "status"
 MODE_LABEL = "mode"
 
 # These are the labels that will be used in the metrics
-STATUS_SUCCESS = "success"  
+STATUS_SUCCESS = "success"
 STATUS_FAILURE = "failure"
 MODE_LOCAL = "local"
 MODE_REMOTE = "remote"
 
 # Embedding metrics
 rag_embedding_latency = Histogram(
-    'rag_embedding_latency_seconds',
-    'Time to embed in seconds',
-    labelnames=[STATUS_LABEL, MODE_LABEL]
+    "rag_embedding_latency_seconds",
+    "Time to embed in seconds",
+    labelnames=[STATUS_LABEL, MODE_LABEL],
 )
 rag_embedding_requests_total = Counter(
-    'rag_embedding_requests_total', 
-    'Count of successful/failed embed requests', 
-    labelnames=[STATUS_LABEL, MODE_LABEL] 
+    "rag_embedding_requests_total",
+    "Count of successful/failed embed requests",
+    labelnames=[STATUS_LABEL, MODE_LABEL],
 )
 
 # Query API metrics
-rag_query_latency = Histogram('rag_query_latency_seconds', 'Time to call \'/query\' API in seconds', labelnames=[STATUS_LABEL])
-rag_query_requests_total = Counter('rag_query_requests_total', 'Count of successful/failed calling \'/query\' requests', labelnames=[STATUS_LABEL])
+rag_query_latency = Histogram(
+    "rag_query_latency_seconds",
+    "Time to call '/query' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_query_requests_total = Counter(
+    "rag_query_requests_total",
+    "Count of successful/failed calling '/query' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Chat API metrics
-rag_chat_latency = Histogram('rag_chat_latency_seconds', 'Time to call \'/v1/chat/completions\' API in seconds', labelnames=[STATUS_LABEL])
-rag_chat_requests_total = Counter('rag_chat_requests_total', 'Count of successful/failed calling \'/v1/chat/completions\' requests', labelnames=[STATUS_LABEL])
+rag_chat_latency = Histogram(
+    "rag_chat_latency_seconds",
+    "Time to call '/v1/chat/completions' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_chat_requests_total = Counter(
+    "rag_chat_requests_total",
+    "Count of successful/failed calling '/v1/chat/completions' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Index API metrics
-rag_index_latency = Histogram('rag_index_latency_seconds', 'Time to call \'/index\' API in seconds', labelnames=[STATUS_LABEL])
-rag_index_requests_total = Counter('rag_index_requests_total', 'Count of successful/failed calling \'/index\' requests', labelnames=[STATUS_LABEL])
+rag_index_latency = Histogram(
+    "rag_index_latency_seconds",
+    "Time to call '/index' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_index_requests_total = Counter(
+    "rag_index_requests_total",
+    "Count of successful/failed calling '/index' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Indexes API metrics
-rag_indexes_latency = Histogram('rag_indexes_latency_seconds', 'Time to call \'/indexes\' API in seconds', labelnames=[STATUS_LABEL])
-rag_indexes_requests_total = Counter('rag_indexes_requests_total', 'Count of successful/failed calling \'/indexes\' requests',  labelnames=[STATUS_LABEL])
+rag_indexes_latency = Histogram(
+    "rag_indexes_latency_seconds",
+    "Time to call '/indexes' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_indexes_requests_total = Counter(
+    "rag_indexes_requests_total",
+    "Count of successful/failed calling '/indexes' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Indexes document API metrics
-rag_indexes_document_latency = Histogram('rag_indexes_document_latency_seconds', 'Time to call get \'/indexes/{index_name}/documents\' API in seconds', labelnames=[STATUS_LABEL])
-rag_indexes_document_requests_total = Counter('rag_indexes_document_requests_total', 'Count of successful/failed calling get \'/indexes/{index_name}/documents\' requests', labelnames=[STATUS_LABEL])
+rag_indexes_document_latency = Histogram(
+    "rag_indexes_document_latency_seconds",
+    "Time to call get '/indexes/{index_name}/documents' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_indexes_document_requests_total = Counter(
+    "rag_indexes_document_requests_total",
+    "Count of successful/failed calling get '/indexes/{index_name}/documents' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Indexes update document API metrics
-rag_indexes_update_document_latency = Histogram('rag_indexes_update_document_latency_seconds', 'Time to call post \'/indexes/{index_name}/documents\' API in seconds', labelnames=[STATUS_LABEL])
-rag_indexes_update_document_requests_total = Counter('rag_indexes_update_document_requests_total', 'Count of successful/failed calling post \'/indexes/{index_name}/documents\' requests', labelnames=[STATUS_LABEL])
+rag_indexes_update_document_latency = Histogram(
+    "rag_indexes_update_document_latency_seconds",
+    "Time to call post '/indexes/{index_name}/documents' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_indexes_update_document_requests_total = Counter(
+    "rag_indexes_update_document_requests_total",
+    "Count of successful/failed calling post '/indexes/{index_name}/documents' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Indexes document API metrics
-rag_indexes_delete_document_latency = Histogram('rag_indexes_delete_document_latency_seconds', 'Time to call \'/indexes/{index_name}/documents/delete\' API in seconds', labelnames=[STATUS_LABEL])
-rag_indexes_delete_document_requests_total = Counter('rag_indexes_delete_document_requests_total', 'Count of successful/failed calling \'/indexes/{index_name}/documents/delete\' requests', labelnames=[STATUS_LABEL])
+rag_indexes_delete_document_latency = Histogram(
+    "rag_indexes_delete_document_latency_seconds",
+    "Time to call '/indexes/{index_name}/documents/delete' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_indexes_delete_document_requests_total = Counter(
+    "rag_indexes_delete_document_requests_total",
+    "Count of successful/failed calling '/indexes/{index_name}/documents/delete' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Persist API metrics
-rag_persist_latency = Histogram('rag_persist_latency_seconds', 'Time to call \'/persist/{index_name}\' API in seconds', labelnames=[STATUS_LABEL])
-rag_persist_requests_total = Counter('rag_persist_requests_total', 'Count of successful/failed calling \'/persist/{index_name}\' requests', labelnames=[STATUS_LABEL])
+rag_persist_latency = Histogram(
+    "rag_persist_latency_seconds",
+    "Time to call '/persist/{index_name}' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_persist_requests_total = Counter(
+    "rag_persist_requests_total",
+    "Count of successful/failed calling '/persist/{index_name}' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Load API metrics
-rag_load_latency = Histogram('rag_load_latency_seconds', 'Time to call \'/load/{index_name}\' API in seconds', labelnames=[STATUS_LABEL])
-rag_load_requests_total = Counter('rag_load_requests_total', 'Count of successful/failed calling \'/load/{index_name}\' requests', labelnames=[STATUS_LABEL])
+rag_load_latency = Histogram(
+    "rag_load_latency_seconds",
+    "Time to call '/load/{index_name}' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_load_requests_total = Counter(
+    "rag_load_requests_total",
+    "Count of successful/failed calling '/load/{index_name}' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # Delete API metrics
-rag_delete_latency = Histogram('rag_delete_index_latency_seconds', 'Time to call delete \'/indexes/{index_name}\' API in seconds', labelnames=[STATUS_LABEL])
-rag_delete_requests_total = Counter('rag_delete_index_requests_total', 'Count of successful/failed calling delete \'/indexes/{index_name}\' requests', labelnames=[STATUS_LABEL])
+rag_delete_latency = Histogram(
+    "rag_delete_index_latency_seconds",
+    "Time to call delete '/indexes/{index_name}' API in seconds",
+    labelnames=[STATUS_LABEL],
+)
+rag_delete_requests_total = Counter(
+    "rag_delete_index_requests_total",
+    "Count of successful/failed calling delete '/indexes/{index_name}' requests",
+    labelnames=[STATUS_LABEL],
+)
 
 # End-to-end request metrics
 e2e_request_total = Counter(
-    'e2e_request_total', 
-    'Total number of all processed requests', 
-    labelnames=[STATUS_LABEL]
+    "e2e_request_total",
+    "Total number of all processed requests",
+    labelnames=[STATUS_LABEL],
 )
 e2e_request_latency_seconds = Histogram(
-    'e2e_request_latency_seconds',
-    'End to end request latency in seconds',
-    labelnames=[STATUS_LABEL]
+    "e2e_request_latency_seconds",
+    "End to end request latency in seconds",
+    labelnames=[STATUS_LABEL],
 )
 
 # Active requests gauge
 num_requests_running = Gauge(
-    'num_requests_running',
-    'Number of requests currently being processed'
+    "num_requests_running", "Number of requests currently being processed"
 )
 
 # RAG source score metrics
 rag_lowest_source_score = Histogram(
-    'rag_lowest_source_score',
-    'Score of the lowest scoring source node (typically the most relevant)',
-    buckets=(0.1, 0.2, 0.3, 0.35, 0.4, 0.45, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0)
+    "rag_lowest_source_score",
+    "Score of the lowest scoring source node (typically the most relevant)",
+    buckets=(0.1, 0.2, 0.3, 0.35, 0.4, 0.45, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0),
 )
 
 rag_avg_source_score = Histogram(
-    'rag_avg_source_score',
-    'Average score of all retrieved source documents in RAG queries',
-    buckets=(0.1, 0.2, 0.3, 0.35, 0.4, 0.45, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0)
-)
\ No newline at end of file
+    "rag_avg_source_score",
+    "Average score of all retrieved source documents in RAG queries",
+    buckets=(0.1, 0.2, 0.3, 0.35, 0.4, 0.45, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0),
+)
diff --git a/presets/ragengine/models.py b/presets/ragengine/models.py
index 7490b00..c2242a0 100644
--- a/presets/ragengine/models.py
+++ b/presets/ragengine/models.py
@@ -12,81 +12,82 @@
 # limitations under the License.
 
 
-from typing import Any, Dict, List, Optional
-from typing_extensions import Literal, Required, Optional
+from typing import Any
+
 from openai.types.chat import (
     ChatCompletion,
-    CompletionCreateParams,
-    ChatCompletionMessageParam,
-    ChatCompletionUserMessageParam,
-    ChatCompletionSystemMessageParam,
-    ChatCompletionDeveloperMessageParam,
-    ChatCompletionContentPartTextParam,
 )
-
-from pydantic import BaseModel, Field, model_validator, ValidationError, create_model
+from pydantic import BaseModel, Field, model_validator
 
 
 class Document(BaseModel):
     doc_id: str = Field(default="")
     text: str
-    metadata: Optional[dict] = Field(default_factory=dict)
-    hash_value: Optional[str] = None
+    metadata: dict | None = Field(default_factory=dict)
+    hash_value: str | None = None
     is_truncated: bool = False
 
+
 class ListDocumentsResponse(BaseModel):
-    documents: List[Document] # List of DocumentResponses
+    documents: list[Document]  # List of DocumentResponses
     count: int  # Number of documents in the current response
 
+
 class IndexRequest(BaseModel):
     index_name: str
-    documents: List[Document]
+    documents: list[Document]
+
 
 class UpdateDocumentRequest(BaseModel):
-    documents: List[Document]
+    documents: list[Document]
+
 
 class UpdateDocumentResponse(BaseModel):
-    updated_documents: List[Document]
-    unchanged_documents: List[Document]
-    not_found_documents: List[Document]
+    updated_documents: list[Document]
+    unchanged_documents: list[Document]
+    not_found_documents: list[Document]
+
 
 class DeleteDocumentRequest(BaseModel):
-    doc_ids: List[str]
+    doc_ids: list[str]
+
 
 class DeleteDocumentResponse(BaseModel):
-    deleted_doc_ids: List[str]
-    not_found_doc_ids: List[str]
+    deleted_doc_ids: list[str]
+    not_found_doc_ids: list[str]
+
 
 class QueryRequest(BaseModel):
     index_name: str
     query: str
     top_k: int = 5
     # Accept a dictionary for our LLM parameters
-    llm_params: Optional[Dict[str, Any]] = Field(
+    llm_params: dict[str, Any] | None = Field(
         default_factory=dict,
         description="Optional parameters for the language model, e.g., temperature, top_p",
     )
     # Accept a dictionary for rerank parameters
-    rerank_params: Optional[Dict[str, Any]] = Field(
+    rerank_params: dict[str, Any] | None = Field(
         default_factory=dict,
         description="Experimental: Optional parameters for reranking. Only 'top_n' and 'choice_batch_size' are supported.",
     )
 
     @model_validator(mode="after")
-    def validate_params(cls, values: "QueryRequest") -> "QueryRequest":
-        # Access fields as attributes instead of treating as a dictionary
-        llm_params = values.llm_params # Validate LLM parameters on vLLM side
-        rerank_params = values.rerank_params
-        top_k = values.top_k
+    def validate_params(self) -> "QueryRequest":
+        # Access fields as attributes
+        rerank_params = self.rerank_params
+        top_k = self.top_k
 
         # Validate rerank parameters
         if "top_n" in rerank_params:
             if not isinstance(rerank_params["top_n"], int):
                 raise ValueError("Invalid type: 'top_n' must be an integer.")
             if rerank_params["top_n"] > top_k:
-                raise ValueError("Invalid configuration: 'top_n' for reranking cannot exceed 'top_k' from the RAG query.")
+                raise ValueError(
+                    "Invalid configuration: 'top_n' for reranking cannot exceed 'top_k' from the RAG query."
+                )
+        return self
 
-        return values
 
 # Define models for NodeWithScore, and QueryResponse
 class NodeWithScore(BaseModel):
@@ -94,21 +95,25 @@ class NodeWithScore(BaseModel):
     node_id: str
     text: str
     score: float
-    metadata: Optional[dict] = None
+    metadata: dict | None = None
+
 
 class QueryResponse(BaseModel):
     response: str
-    source_nodes: List[NodeWithScore]
-    metadata: Optional[dict] = None
+    source_nodes: list[NodeWithScore]
+    metadata: dict | None = None
+
 
 class HealthStatus(BaseModel):
     status: str
-    detail: Optional[str] = None
+    detail: str | None = None
+
 
 class ChatCompletionResponse(ChatCompletion):
-    source_nodes: Optional[List[NodeWithScore]] = None
+    source_nodes: list[NodeWithScore] | None = None
+
 
-def messages_to_prompt(messages: List[Dict]) -> str:
+def messages_to_prompt(messages: list[dict]) -> str:
     """Convert messages to a prompt string."""
     string_messages = []
     for message in messages:
@@ -116,7 +121,8 @@ def messages_to_prompt(messages: List[Dict]) -> str:
         string_messages.append(f"{message.get('role')}: {content}")
     return "\n".join(string_messages)
 
-def get_message_content(message: Dict) -> str:
+
+def get_message_content(message: dict) -> str:
     """Extract content from a ChatCompletionMessageParam."""
     if message.get("role") == "user":
         if message.get("content"):
@@ -135,4 +141,8 @@ def get_message_content(message: Dict) -> str:
                 return "\n".join(user_text_content)
         return ""
     else:
-        return message.get("content", {}).get("text", "") if isinstance(message.get("content"), dict) else message.get("content", "")
\ No newline at end of file
+        return (
+            message.get("content", {}).get("text", "")
+            if isinstance(message.get("content"), dict)
+            else message.get("content", "")
+        )
diff --git a/presets/ragengine/vector_store/base.py b/presets/ragengine/vector_store/base.py
index 7cb4b8c..7642928 100644
--- a/presets/ragengine/vector_store/base.py
+++ b/presets/ragengine/vector_store/base.py
@@ -12,39 +12,38 @@
 # limitations under the License.
 
 
-import logging
-from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Any
-from pydantic import ValidationError
+import asyncio
 import hashlib
+import logging
 import os
-import asyncio
-from itertools import islice
-import uuid
 import time
+import uuid
+from abc import ABC, abstractmethod
+from itertools import islice
+from typing import Any
 
+import aiorwlock
+from fastapi import HTTPException
 from llama_index.core import Document as LlamaDocument
-from llama_index.core import (StorageContext, VectorStoreIndex, load_index_from_storage)
+from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage
 from llama_index.core.chat_engine.types import ChatMode
 from llama_index.core.postprocessor import LLMRerank  # Query with LLM Reranking
-
+from llama_index.core.storage.docstore import SimpleDocumentStore
 from llama_index.vector_stores.faiss import FaissMapVectorStore
 from openai.types.chat import ChatCompletionContentPartTextParam, CompletionCreateParams
+from pydantic import ValidationError
 
-from ragengine.models import Document, ChatCompletionResponse, messages_to_prompt
+from ragengine.config import LLM_RERANKER_BATCH_SIZE, LLM_RERANKER_TOP_N
 from ragengine.embedding.base import BaseEmbeddingModel
 from ragengine.inference.inference import Inference
+from ragengine.models import ChatCompletionResponse, Document, messages_to_prompt
 from ragengine.vector_store.transformers.custom_transformer import CustomTransformer
-from ragengine.config import (LLM_RERANKER_BATCH_SIZE, LLM_RERANKER_TOP_N)
-from fastapi import HTTPException
-
-from llama_index.core.storage.docstore import SimpleDocumentStore
-import aiorwlock
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class BaseVectorStore(ABC):
     def __init__(self, embed_model: BaseEmbeddingModel, use_rwlock: bool = False):
         super().__init__()
@@ -59,46 +58,64 @@ class BaseVectorStore(ABC):
     @staticmethod
     def generate_doc_id(text: str) -> str:
         """Generates a unique document ID based on the hash of the document text."""
-        return hashlib.sha256(text.encode('utf-8')).hexdigest()
+        return hashlib.sha256(text.encode("utf-8")).hexdigest()
 
     async def shutdown(self):
         await self.llm.aclose()
 
-    async def index_documents(self, index_name: str, documents: List[Document]) -> List[str]:
+    async def index_documents(
+        self, index_name: str, documents: list[Document]
+    ) -> list[str]:
         """Common indexing logic for all vector stores."""
         if index_name in self.index_map:
             return await self._append_documents_to_index(index_name, documents)
         else:
             return await self._create_new_index(index_name, documents)
 
-    async def _append_documents_to_index(self, index_name: str, documents: List[Document]) -> List[str]:
+    async def _append_documents_to_index(
+        self, index_name: str, documents: list[Document]
+    ) -> list[str]:
         """Common logic for appending documents to existing index."""
-        logger.info(f"Index {index_name} already exists. Appending documents to existing index.")
+        logger.info(
+            f"Index {index_name} already exists. Appending documents to existing index."
+        )
         indexed_docs = [None] * len(documents)
 
         async def handle_document(doc_index: int, doc: Document):
             doc_id = self.generate_doc_id(doc.text)
             if self.use_rwlock:
                 async with self.rwlock.reader_lock:
-                    retrieved_doc = await self.index_map[index_name].docstore.aget_ref_doc_info(doc_id)
+                    retrieved_doc = await self.index_map[
+                        index_name
+                    ].docstore.aget_ref_doc_info(doc_id)
             else:
-                retrieved_doc = await self.index_map[index_name].docstore.aget_ref_doc_info(doc_id)
+                retrieved_doc = await self.index_map[
+                    index_name
+                ].docstore.aget_ref_doc_info(doc_id)
             if not retrieved_doc:
                 await self.add_document_to_index(index_name, doc, doc_id)
             else:
-                logger.info(f"Document {doc_id} already exists in index {index_name}. Skipping.")
+                logger.info(
+                    f"Document {doc_id} already exists in index {index_name}. Skipping."
+                )
             indexed_docs[doc_index] = doc_id
 
         # Gather all coroutines for processing documents
-        await asyncio.gather(*(handle_document(idx, doc) for idx, doc in enumerate(documents)))
+        await asyncio.gather(
+            *(handle_document(idx, doc) for idx, doc in enumerate(documents))
+        )
         return indexed_docs
-    
+
     @abstractmethod
-    async def _create_new_index(self, index_name: str, documents: List[Document]) -> List[str]:
+    async def _create_new_index(
+        self, index_name: str, documents: list[Document]
+    ) -> list[str]:
         """Create a new index - implementation specific to each vector store."""
         pass
-    
-    async def _create_index_common(self, index_name: str, documents: List[Document], vector_store) -> List[str]:
+
+    async def _create_index_common(
+        self, index_name: str, documents: list[Document], vector_store
+    ) -> list[str]:
         """Common logic for creating a new index with documents."""
         storage_context = StorageContext.from_defaults(vector_store=vector_store)
         llama_docs = []
@@ -136,12 +153,13 @@ class BaseVectorStore(ABC):
                 self.index_map[index_name] = index
         return indexed_doc_ids
 
-    async def query(self,
-              index_name: str,
-              query: str,
-              top_k: int,
-              llm_params: dict,
-              rerank_params: dict
+    async def query(
+        self,
+        index_name: str,
+        query: str,
+        top_k: int,
+        llm_params: dict,
+        rerank_params: dict,
     ):
         """
         Query the indexed documents
@@ -159,14 +177,18 @@ class BaseVectorStore(ABC):
             dict: A dictionary containing the response and source nodes.
         """
         if index_name not in self.index_map:
-            raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
+            raise HTTPException(
+                status_code=404, detail=f"No such index: '{index_name}' exists."
+            )
 
         node_postprocessors = []
         if rerank_params:
             # Set default reranking parameters and merge with provided params
             default_rerank_params = {
-                'choice_batch_size': LLM_RERANKER_BATCH_SIZE,  # Default batch size
-                'top_n': min(LLM_RERANKER_TOP_N, top_k)        # Limit top_n to top_k by default
+                "choice_batch_size": LLM_RERANKER_BATCH_SIZE,  # Default batch size
+                "top_n": min(
+                    LLM_RERANKER_TOP_N, top_k
+                ),  # Limit top_n to top_k by default
             }
             rerank_params = {**default_rerank_params, **rerank_params}
 
@@ -174,8 +196,8 @@ class BaseVectorStore(ABC):
             node_postprocessors.append(
                 LLMRerank(
                     llm=self.llm,
-                    choice_batch_size=rerank_params['choice_batch_size'],
-                    top_n=rerank_params['top_n']
+                    choice_batch_size=rerank_params["choice_batch_size"],
+                    top_n=rerank_params["top_n"],
                 )
             )
 
@@ -184,7 +206,7 @@ class BaseVectorStore(ABC):
             similarity_top_k=top_k,
             node_postprocessors=node_postprocessors,
             chat_mode=ChatMode.CONDENSE_PLUS_CONTEXT,
-            verbose=True
+            verbose=True,
         )
 
         if self.use_rwlock:
@@ -202,7 +224,7 @@ class BaseVectorStore(ABC):
                     "node_id": source_node.node_id,
                     "text": source_node.text,
                     "score": source_node.score,
-                    "metadata": source_node.metadata
+                    "metadata": source_node.metadata,
                 }
                 for source_node in query_result.source_nodes
             ],
@@ -219,8 +241,14 @@ class BaseVectorStore(ABC):
         Returns:
             ChatCompletionResponse: The response containing the generated chat completion.
         """
-        if request.get("index_name") and request.get("index_name") not in self.index_map:
-            raise HTTPException(status_code=404, detail=f"No such index: '{request.get('index_name')}' exists.")
+        if (
+            request.get("index_name")
+            and request.get("index_name") not in self.index_map
+        ):
+            raise HTTPException(
+                status_code=404,
+                detail=f"No such index: '{request.get('index_name')}' exists.",
+            )
 
         llm_params = {}
         if request.get("model") is not None:
@@ -244,29 +272,43 @@ class BaseVectorStore(ABC):
 
         if openai_request is None:
             logger.error(f"Invalid request format: {str(last_error)}")
-            raise HTTPException(status_code=400, detail=f"Invalid request format: {str(last_error)}")
-        
+            raise HTTPException(
+                status_code=400, detail=f"Invalid request format: {str(last_error)}"
+            )
+
         if not request.get("index_name"):
-            logger.info(f"Request does not specify an index, passing through to LLM directly.")
+            logger.info(
+                "Request does not specify an index, passing through to LLM directly."
+            )
             return await self.llm.chat_completions_passthrough(openai_request)
 
         if request.get("tools") or request.get("functions"):
-            logger.info(f"Request contains tools or functions, passing through to LLM directly.")
+            logger.info(
+                "Request contains tools or functions, passing through to LLM directly."
+            )
             return await self.llm.chat_completions_passthrough(openai_request)
 
         # Only support RAG usage on user/system/developer roles in messages and only text content
         for message in request.get("messages", []):
             # Every message must have a role
             if not message.get("role"):
-                raise HTTPException(status_code=400, detail=f"Invalid request format: messages must contain 'role'.")
-            
+                raise HTTPException(
+                    status_code=400,
+                    detail="Invalid request format: messages must contain 'role'.",
+                )
+
             # Every message must have content aside from assistant role messages
             if message.get("role") != "assistant" and message.get("content") is None:
-                raise HTTPException(status_code=400, detail=f"Invalid request format: messages must contain 'content' for role '{message.get('role')}'.")
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Invalid request format: messages must contain 'content' for role '{message.get('role')}'.",
+                )
 
             # Only user, system, and developer roles are supported for RAG
             if message.get("role") not in ["user", "system", "assistant", "developer"]:
-                logger.info(f"Request contains unsupported role '{message.get('role')}' in messages, passing through to LLM directly.")
+                logger.info(
+                    f"Request contains unsupported role '{message.get('role')}' in messages, passing through to LLM directly."
+                )
                 return await self.llm.chat_completions_passthrough(openai_request)
 
             # User message content can be a range of options, but we only support text content for RAG
@@ -276,27 +318,40 @@ class BaseVectorStore(ABC):
                     if isinstance(content, list):
                         for part in content:
                             if not isinstance(part, str) and part.get("type") != "text":
-                                logger.info(f"Request contains unsupported content type in user message, passing through to LLM directly.")
-                                return await self.llm.chat_completions_passthrough(openai_request)
-                    elif isinstance(content, str):
-                        pass
-                    elif isinstance(content, ChatCompletionContentPartTextParam):
+                                logger.info(
+                                    "Request contains unsupported content type in user message, passing through to LLM directly."
+                                )
+                                return await self.llm.chat_completions_passthrough(
+                                    openai_request
+                                )
+                    elif isinstance(content, str | ChatCompletionContentPartTextParam):
                         pass
                     else:
-                        logger.info(f"Request contains unsupported content type '{type(content)}' in messages, passing through to LLM directly.")
-                        return await self.llm.chat_completions_passthrough(openai_request)
+                        logger.info(
+                            f"Request contains unsupported content type '{type(content)}' in messages, passing through to LLM directly."
+                        )
+                        return await self.llm.chat_completions_passthrough(
+                            openai_request
+                        )
                 else:
-                    logger.error(f"Invalid request format: user messages must contain 'content'.")
-                    raise HTTPException(status_code=400, detail=f"Invalid request format: user messages must contain 'content'.")
+                    logger.error(
+                        "Invalid request format: user messages must contain 'content'."
+                    )
+                    raise HTTPException(
+                        status_code=400,
+                        detail="Invalid request format: user messages must contain 'content'.",
+                    )
 
         prompt = messages_to_prompt(request.get("messages", []))
 
-        logger.info(f"Creating chat engine for index '{request.get('index_name')}' with prompt: {prompt}")
+        logger.info(
+            f"Creating chat engine for index '{request.get('index_name')}' with prompt: {prompt}"
+        )
         chat_engine = self.index_map[request.get("index_name")].as_chat_engine(
             llm=self.llm,
             similarity_top_k=request.get("top_k", 5),
             chat_mode=ChatMode.CONDENSE_PLUS_CONTEXT,
-            verbose=True
+            verbose=True,
         )
 
         logger.info("Processing chat completion request with prompt.")
@@ -314,55 +369,74 @@ class BaseVectorStore(ABC):
                 object="chat.completion",
                 created=int(time.time()),
                 model=request.get("model"),
-                choices=[{
-                    "message": {
-                        "role": "assistant",
-                        "content": chat_result.response,
-                    },
-                    "finish_reason": "stop",
-                    "index": 0
-                }],
+                choices=[
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": chat_result.response,
+                        },
+                        "finish_reason": "stop",
+                        "index": 0,
+                    }
+                ],
                 source_nodes=[
                     {
                         "doc_id": source_node.node.ref_doc_id,
                         "node_id": source_node.node_id,
                         "text": source_node.text,
                         "score": source_node.score,
-                        "metadata": source_node.metadata
+                        "metadata": source_node.metadata,
                     }
                     for source_node in chat_result.source_nodes
                 ],
             )
         except Exception as e:
             logger.error(f"Error during chat completion: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Chat completion failed: {str(e)}")
+            raise HTTPException(
+                status_code=500, detail=f"Chat completion failed: {str(e)}"
+            )
 
-    async def add_document_to_index(self, index_name: str, document: Document, doc_id: str):
+    async def add_document_to_index(
+        self, index_name: str, document: Document, doc_id: str
+    ):
         """Common logic for adding a single document."""
         if index_name not in self.index_map:
-            raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
+            raise HTTPException(
+                status_code=404, detail=f"No such index: '{index_name}' exists."
+            )
 
-        llama_doc = LlamaDocument(id_=doc_id, text=document.text, metadata=document.metadata, excluded_llm_metadata_keys=[key for key in document.metadata.keys()])
+        llama_doc = LlamaDocument(
+            id_=doc_id,
+            text=document.text,
+            metadata=document.metadata,
+            excluded_llm_metadata_keys=[key for key in document.metadata],
+        )
 
         if self.use_rwlock:
             async with self.rwlock.writer_lock:
-                retrieved_doc = await self.index_map[index_name].docstore.aget_ref_doc_info(doc_id)
+                retrieved_doc = await self.index_map[
+                    index_name
+                ].docstore.aget_ref_doc_info(doc_id)
                 if retrieved_doc:
-                    logger.info(f"Document {doc_id} already exists in index {index_name} (double-check). Skipping insertion.")
+                    logger.info(
+                        f"Document {doc_id} already exists in index {index_name} (double-check). Skipping insertion."
+                    )
                     return
                 # Proceed with insertion only if the document is absent
                 await asyncio.to_thread(self.index_map[index_name].insert, llama_doc)
         else:
             await asyncio.to_thread(self.index_map[index_name].insert, llama_doc)
 
-    def list_indexes(self) -> List[str]:
-        return list(self.index_map.keys())
+    def list_indexes(self) -> list[str]:
+        return list(self.index_map)
 
-    async def delete_documents(self, index_name: str, doc_ids: List[str]):
+    async def delete_documents(self, index_name: str, doc_ids: list[str]):
         """Common logic for deleting a document."""
         if index_name not in self.index_map:
-            raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
-        
+            raise HTTPException(
+                status_code=404, detail=f"No such index: '{index_name}' exists."
+            )
+
         not_found_docs = []
         found_docs = []
         for doc_id in doc_ids:
@@ -374,11 +448,25 @@ class BaseVectorStore(ABC):
         try:
             if self.use_rwlock:
                 async with self.rwlock.writer_lock:
-                    await asyncio.gather(*(self.index_map[index_name].adelete_ref_doc(doc_id, delete_from_docstore=True) for doc_id in doc_ids),
-                                        return_exceptions=True)
+                    await asyncio.gather(
+                        *(
+                            self.index_map[index_name].adelete_ref_doc(
+                                doc_id, delete_from_docstore=True
+                            )
+                            for doc_id in doc_ids
+                        ),
+                        return_exceptions=True,
+                    )
             else:
-                await asyncio.gather(*(self.index_map[index_name].adelete_ref_doc(doc_id, delete_from_docstore=True) for doc_id in doc_ids),
-                                        return_exceptions=True)
+                await asyncio.gather(
+                    *(
+                        self.index_map[index_name].adelete_ref_doc(
+                            doc_id, delete_from_docstore=True
+                        )
+                        for doc_id in doc_ids
+                    ),
+                    return_exceptions=True,
+                )
 
             return {"deleted_doc_ids": found_docs, "not_found_doc_ids": not_found_docs}
         except NotImplementedError as e:
@@ -387,12 +475,14 @@ class BaseVectorStore(ABC):
         except Exception as e:
             logger.error(f"Error deleting documents from index {index_name}: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Loading failed: {str(e)}")
-    
-    async def update_documents(self, index_name: str, documents: List[Document]):
+
+    async def update_documents(self, index_name: str, documents: list[Document]):
         """Common logic for updating a document."""
         if index_name not in self.index_map:
-            raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
-        
+            raise HTTPException(
+                status_code=404, detail=f"No such index: '{index_name}' exists."
+            )
+
         not_found_docs = []
         found_docs = []
         llama_docs = []
@@ -400,7 +490,12 @@ class BaseVectorStore(ABC):
             if document.doc_id in self.index_map[index_name].ref_doc_info:
                 found_docs.append(document)
                 llama_docs.append(
-                    LlamaDocument(id_=document.doc_id, text=document.text, metadata=document.metadata, excluded_llm_metadata_keys=[key for key in document.metadata.keys()])
+                    LlamaDocument(
+                        id_=document.doc_id,
+                        text=document.text,
+                        metadata=document.metadata,
+                        excluded_llm_metadata_keys=[key for key in document.metadata],
+                    )
                 )
             else:
                 not_found_docs.append(document)
@@ -408,10 +503,12 @@ class BaseVectorStore(ABC):
         try:
             if self.use_rwlock:
                 async with self.rwlock.writer_lock:
-                    refreshed_docs = self.index_map[index_name].refresh_ref_docs(llama_docs)
+                    refreshed_docs = self.index_map[index_name].refresh_ref_docs(
+                        llama_docs
+                    )
             else:
                 refreshed_docs = self.index_map[index_name].refresh_ref_docs(llama_docs)
-            
+
             updated_docs = []
             unchanged_docs = []
             for doc, was_updated in zip(found_docs, refreshed_docs):
@@ -419,7 +516,11 @@ class BaseVectorStore(ABC):
                     updated_docs.append(doc)
                 else:
                     unchanged_docs.append(doc)
-            return {"updated_documents": updated_docs, "unchanged_documents": unchanged_docs, "not_found_documents": not_found_docs}
+            return {
+                "updated_documents": updated_docs,
+                "unchanged_documents": unchanged_docs,
+                "not_found_documents": not_found_docs,
+            }
         except NotImplementedError as e:
             logger.error(f"Update operation is not implemented for index {index_name}.")
             raise HTTPException(status_code=501, detail=f"Loading failed: {str(e)}")
@@ -427,16 +528,26 @@ class BaseVectorStore(ABC):
             logger.error(f"Error updating documents in index {index_name}: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Loading failed: {str(e)}")
 
-    async def _process_document(self, doc_id: str, doc_stub, doc_store, max_text_length: Optional[int]):
+    async def _process_document(
+        self, doc_id: str, doc_stub, doc_store, max_text_length: int | None
+    ):
         """
         Helper to process and format a single document.
         """
         try:
             if isinstance(doc_store, SimpleDocumentStore):
-                text, hash_value, ref_doc_id = doc_stub.text, doc_stub.hash, doc_stub.ref_doc_id
+                text, hash_value, ref_doc_id = (
+                    doc_stub.text,
+                    doc_stub.hash,
+                    doc_stub.ref_doc_id,
+                )
             else:
                 doc_info = await doc_store.aget_document(doc_id)
-                text, hash_value, ref_doc_id = doc_info.text, doc_info.hash, doc_stub.ref_doc_id
+                text, hash_value, ref_doc_id = (
+                    doc_info.text,
+                    doc_info.hash,
+                    doc_stub.ref_doc_id,
+                )
 
             # Truncate if needed
             is_truncated = bool(max_text_length and len(text) > max_text_length)
@@ -453,31 +564,39 @@ class BaseVectorStore(ABC):
             logger.error(f"Error processing document {doc_id}: {str(e)}")
             return None  # Explicitly return None for failed documents
 
-    async def list_documents_in_index(self, 
-            index_name: str, 
-            limit: int, 
-            offset: int, 
-            max_text_length: Optional[int] = None,
-            metadata_filter: Optional[Dict[str, Any]] = None
-        ) -> List[Dict[str, Any]]:
+    async def list_documents_in_index(
+        self,
+        index_name: str,
+        limit: int,
+        offset: int,
+        max_text_length: int | None = None,
+        metadata_filter: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
         """
         Return a dictionary of document metadata for the given index.
         """
         vector_store_index = self.index_map.get(index_name)
         if not vector_store_index:
-            raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
-        
+            raise HTTPException(
+                status_code=404, detail=f"No such index: '{index_name}' exists."
+            )
+
         doc_store = vector_store_index.docstore
         doc_store_items = doc_store.docs.items()
         if metadata_filter is not None:
-            docs_items = await self._filter_documents(doc_store_items, metadata_filter, offset, limit)
+            docs_items = await self._filter_documents(
+                doc_store_items, metadata_filter, offset, limit
+            )
         else:
             docs_items = islice(doc_store_items, offset, offset + limit)
 
         # Process documents concurrently, handling exceptions
         docs = await asyncio.gather(
-            *(self._process_document(doc_id, doc_stub, doc_store, max_text_length) for doc_id, doc_stub in docs_items),
-            return_exceptions=True
+            *(
+                self._process_document(doc_id, doc_stub, doc_store, max_text_length)
+                for doc_id, doc_stub in docs_items
+            ),
+            return_exceptions=True,
         )
 
         # Return list of valid documents
@@ -486,14 +605,16 @@ class BaseVectorStore(ABC):
     async def delete_index(self, index_name: str):
         """Common logic for deleting an index."""
         if index_name not in self.index_map:
-            raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
-        
+            raise HTTPException(
+                status_code=404, detail=f"No such index: '{index_name}' exists."
+            )
+
         if self.use_rwlock:
             async with self.rwlock.writer_lock:
                 del self.index_map[index_name]
         else:
             del self.index_map[index_name]
-        
+
         logger.info(f"Index {index_name} deleted successfully.")
 
     async def _filter_documents(self, doc_items, metadata_filter, offset, limit):
@@ -507,7 +628,9 @@ class BaseVectorStore(ABC):
                 filtered_docs.append((doc_id, doc_stub))
         return islice(filtered_docs, offset, offset + limit)
 
-    async def document_exists(self, index_name: str, doc: Document, doc_id: str) -> bool:
+    async def document_exists(
+        self, index_name: str, doc: Document, doc_id: str
+    ) -> bool:
         """Common logic for checking document existence."""
         if index_name not in self.index_map:
             logger.warning(f"No such index: '{index_name}' exists in vector store.")
@@ -518,21 +641,25 @@ class BaseVectorStore(ABC):
         """Common persistence logic for individual index."""
 
         if index_name not in self.index_map:
-            raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
+            raise HTTPException(
+                status_code=404, detail=f"No such index: '{index_name}' exists."
+            )
 
         if self.use_rwlock:
             async with self.rwlock.writer_lock:
                 await self._persist_internal(index_name, path)
         else:
             await self._persist_internal(index_name, path)
-            
+
     async def _persist_internal(self, index_name: str, path: str):
         """Common persistence logic for individual index."""
         try:
             # Ensure the directory exists
             os.makedirs(path, exist_ok=True)
             if index_name not in self.index_map:
-                raise HTTPException(status_code=404, detail=f"No such index: '{index_name}' exists.")
+                raise HTTPException(
+                    status_code=404, detail=f"No such index: '{index_name}' exists."
+                )
 
             # Persist the specific index
             storage_context = self.index_map[index_name].storage_context
@@ -559,30 +686,36 @@ class BaseVectorStore(ABC):
             if index_name in self.index_map and not overwrite:
                 raise HTTPException(
                     status_code=409,
-                    detail=f"Index '{index_name}' already exists. Use a different name or delete the existing index first."
+                    detail=f"Index '{index_name}' already exists. Use a different name or delete the existing index first.",
                 )
 
             logger.info(f"Loading index {index_name} from {path}.")
 
             try:
                 storage_context = StorageContext.from_defaults(persist_dir=path)
-            except UnicodeDecodeError as ude:
+            except UnicodeDecodeError:
                 # Failed to load the index in the default json format, trying faissdb
                 faiss_vs = FaissMapVectorStore.from_persist_dir(persist_dir=path)
-                storage_context = StorageContext.from_defaults(persist_dir=path, vector_store=faiss_vs)
+                storage_context = StorageContext.from_defaults(
+                    persist_dir=path, vector_store=faiss_vs
+                )
             except Exception as e:
                 logger.error(f"Failed to load index '{index_name}'. Error: {str(e)}")
                 raise HTTPException(status_code=500, detail=f"Loading failed: {str(e)}")
 
-            logger.info(f"Loading index '{index_name}' using the workspace's embedding model.")
+            logger.info(
+                f"Loading index '{index_name}' using the workspace's embedding model."
+            )
             # Load the index using the workspace's embedding model, assuming all indices
             # were created using the same embedding model currently in use.
-            loaded_index = await asyncio.to_thread(load_index_from_storage,
-                                                   storage_context,
-                                                   embed_model=self.embed_model,
-                                                   show_progress=True)
+            loaded_index = await asyncio.to_thread(
+                load_index_from_storage,
+                storage_context,
+                embed_model=self.embed_model,
+                show_progress=True,
+            )
             self.index_map[index_name] = loaded_index
             logger.info(f"Successfully loaded index {index_name}.")
         except Exception as e:
             logger.error(f"Failed to load index {index_name}. Error: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Loading failed: {str(e)}")
\ No newline at end of file
+            raise HTTPException(status_code=500, detail=f"Loading failed: {str(e)}")
diff --git a/presets/ragengine/vector_store/faiss_store.py b/presets/ragengine/vector_store/faiss_store.py
index 2c39270..f13ae92 100644
--- a/presets/ragengine/vector_store/faiss_store.py
+++ b/presets/ragengine/vector_store/faiss_store.py
@@ -12,12 +12,12 @@
 # limitations under the License.
 
 
-from typing import List
-
 import faiss
 from llama_index.vector_stores.faiss import FaissMapVectorStore
-from ragengine.models import Document
+
 from ragengine.embedding.base import BaseEmbeddingModel
+from ragengine.models import Document
+
 from .base import BaseVectorStore
 
 
@@ -26,7 +26,9 @@ class FaissVectorStoreHandler(BaseVectorStore):
         super().__init__(embed_model, use_rwlock=True)
         self.dimension = self.embed_model.get_embedding_dimension()
 
-    async def _create_new_index(self, index_name: str, documents: List[Document]) -> List[str]:
+    async def _create_new_index(
+        self, index_name: str, documents: list[Document]
+    ) -> list[str]:
         faiss_index = faiss.IndexFlatL2(self.dimension)
         # we cant use the IndexFlatL2 directly as its delete functionality changes document ids.
         # we can wrap it in the IDMap to keep the same functionality but also be able to index by ids and support delete with llama_index
diff --git a/presets/ragengine/vector_store/transformers/custom_transformer.py b/presets/ragengine/vector_store/transformers/custom_transformer.py
index ceffeff..7624736 100644
--- a/presets/ragengine/vector_store/transformers/custom_transformer.py
+++ b/presets/ragengine/vector_store/transformers/custom_transformer.py
@@ -13,18 +13,17 @@
 
 
 from llama_index.core.bridge.pydantic import PrivateAttr
-from llama_index.core.node_parser import SentenceSplitter, CodeSplitter
-from typing import Dict, List
+from llama_index.core.node_parser import CodeSplitter, SentenceSplitter
 from llama_index.core.schema import (
-    TransformComponent,
     BaseNode,
+    TransformComponent,
 )
-from llama_index.core.utils import get_tqdm_iterable
+
 
 class CustomTransformer(TransformComponent):
     """Custom transformer for splitting documents based on metadata input."""
 
-    _code_splitters: Dict[str, CodeSplitter] = PrivateAttr()
+    _code_splitters: dict[str, CodeSplitter] = PrivateAttr()
     _sentence_splitter: SentenceSplitter = PrivateAttr()
 
     def __init__(self):
@@ -32,14 +31,14 @@ class CustomTransformer(TransformComponent):
         self._code_splitters = {}
         self._sentence_splitter = SentenceSplitter()
 
-    def split_node(self, node: BaseNode) -> List[BaseNode]:
+    def split_node(self, node: BaseNode) -> list[BaseNode]:
         node_metadata = node.metadata
         split_type = node_metadata.get("split_type", "default")
         if split_type == "code":
             langauge = node_metadata.get("language", "")
             if not langauge:
                 raise ValueError("Language not specified in node metadata.")
-            
+
             if langauge not in self._code_splitters:
                 self._code_splitters[langauge] = CodeSplitter(
                     language=langauge,
@@ -48,12 +47,10 @@ class CustomTransformer(TransformComponent):
         else:
             # Default to sentence splitting
             return self._sentence_splitter([node])
-    
+
     def __call__(self, nodes, **kwargs):
-        all_nodes: List[BaseNode] = []
+        all_nodes: list[BaseNode] = []
         for node in nodes:
             nodes = self.split_node(node)
-            all_nodes.extend(
-                self.split_node(node)
-            )
+            all_nodes.extend(self.split_node(node))
         return all_nodes
diff --git a/presets/ragengine/vector_store_manager/manager.py b/presets/ragengine/vector_store_manager/manager.py
index 84bdb14..d613e1a 100644
--- a/presets/ragengine/vector_store_manager/manager.py
+++ b/presets/ragengine/vector_store_manager/manager.py
@@ -12,7 +12,7 @@
 # limitations under the License.
 
 
-from typing import Dict, List, Any
+from typing import Any
 
 from ragengine.models import Document
 from ragengine.vector_store.base import BaseVectorStore
@@ -22,19 +22,22 @@ class VectorStoreManager:
     def __init__(self, vector_store: BaseVectorStore):
         self.vector_store = vector_store
 
-    async def index(self, index_name: str, documents: List[Document]) -> List[str]:
+    async def index(self, index_name: str, documents: list[Document]) -> list[str]:
         """Index new documents."""
         return await self.vector_store.index_documents(index_name, documents)
 
-    async def query(self,
-              index_name: str,
-              query: str,
-              top_k: int,
-              llm_params: dict,
-              rerank_params: dict
+    async def query(
+        self,
+        index_name: str,
+        query: str,
+        top_k: int,
+        llm_params: dict,
+        rerank_params: dict,
     ):
         """Query the indexed documents."""
-        return await self.vector_store.query(index_name, query, top_k, llm_params, rerank_params)
+        return await self.vector_store.query(
+            index_name, query, top_k, llm_params, rerank_params
+        )
 
     async def chat_completion(self, request: dict):
         """Chat completion using the vector store."""
@@ -44,30 +47,24 @@ class VectorStoreManager:
         """List all indexes."""
         return self.vector_store.list_indexes()
 
-    async def list_documents_in_index(self,
-            index_name: str,
-            limit: int,
-            offset: int,
-            max_text_length: int,
-            metadata_filter: dict,
-    ) -> List[Dict[str, Any]]:
+    async def list_documents_in_index(
+        self,
+        index_name: str,
+        limit: int,
+        offset: int,
+        max_text_length: int,
+        metadata_filter: dict,
+    ) -> list[dict[str, Any]]:
         """List all documents in index."""
         return await self.vector_store.list_documents_in_index(
-            index_name,
-            limit,
-            offset,
-            max_text_length,
-            metadata_filter
+            index_name, limit, offset, max_text_length, metadata_filter
         )
 
-    async def update_documents(self,
-            index_name: str,
-            documents: List[Document]
-    ):
+    async def update_documents(self, index_name: str, documents: list[Document]):
         """Update documents in the index."""
         return await self.vector_store.update_documents(index_name, documents)
 
-    async def delete_documents(self, index_name: str, doc_ids: List[str]) -> List[str]:
+    async def delete_documents(self, index_name: str, doc_ids: list[str]) -> list[str]:
         """Delete documents from the index."""
         return await self.vector_store.delete_documents(index_name, doc_ids)
 
diff --git a/presets/workspace/dependencies/requirements.txt b/presets/workspace/dependencies/requirements.txt
index 0a73c61..69ba063 100644
--- a/presets/workspace/dependencies/requirements.txt
+++ b/presets/workspace/dependencies/requirements.txt
@@ -1,7 +1,7 @@
 # Dependencies for TFS
 
 # Core Dependencies
-vllm==0.9.0
+vllm==0.9.1
 transformers == 4.51.3
 torch==2.7.0
 accelerate==1.0.0
@@ -14,6 +14,7 @@ numpy<3.0,>=1.25.0
 sentencepiece==0.2.0
 jinja2>=3.1.0
 starlette==0.47.2
+lmcache==0.3.2
 
 # For accessing Ray dashboard. ray[default] version should be consistent with the one in vllm/vllm-openai:<vllm-version>.
 # Check with `docker run --entrypoint "" vllm/vllm-openai:<vllm-version> pip freeze | grep ray`
diff --git a/presets/workspace/inference/text-generation/inference_api.py b/presets/workspace/inference/text-generation/inference_api.py
index ff851b6..389f345 100644
--- a/presets/workspace/inference/text-generation/inference_api.py
+++ b/presets/workspace/inference/text-generation/inference_api.py
@@ -11,14 +11,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import codecs
 import logging
 import os
-import sys
 import signal
-import codecs
-from pathlib import Path
+import sys
 from dataclasses import asdict, dataclass, field
-from typing import Annotated, Any, Dict, List, Optional, Union
+from pathlib import Path
+from typing import Annotated, Any
 
 import GPUtil
 import psutil
@@ -26,56 +26,107 @@ import torch
 import transformers
 import uvicorn
 from fastapi import Body, FastAPI, HTTPException
-from fastapi.responses import Response
 from peft import PeftModel
-from pydantic import BaseModel, Extra, Field, validator
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          GenerationConfig, HfArgumentParser)
+from pydantic import BaseModel, Field
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    GenerationConfig,
+    HfArgumentParser,
+)
 
 # Initialize logger
 logger = logging.getLogger(__name__)
-debug_mode = os.environ.get('DEBUG_MODE', 'false').lower() == 'true'
+debug_mode = os.environ.get("DEBUG_MODE", "false").lower() == "true"
 logging.basicConfig(
     level=logging.DEBUG if debug_mode else logging.INFO,
-    format='%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s',
-    datefmt='%m-%d %H:%M:%S')
+    format="%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s",
+    datefmt="%m-%d %H:%M:%S",
+)
+
+ADAPTERS_DIR = "/mnt/adapter"
 
-ADAPTERS_DIR = '/mnt/adapter'
 
 @dataclass
 class ModelConfig:
     """
     Transformers Model Configuration Parameters
     """
-    pipeline: Optional[str] = field(default="text-generation", metadata={"help": "The model pipeline for the pre-trained model"})
-    pretrained_model_name_or_path: Optional[str] = field(default="/workspace/tfs/weights", metadata={"help": "Path to the pretrained model or model identifier from huggingface.co/models"})
-    combination_type: Optional[str]=field(default="svd", metadata={"help": "The combination type of multi adapters"})
-    state_dict: Optional[Dict[str, Any]] = field(default=None, metadata={"help": "State dictionary for the model"})
-    cache_dir: Optional[str] = field(default=None, metadata={"help": "Cache directory for the model"})
-    from_tf: bool = field(default=False, metadata={"help": "Load model from a TensorFlow checkpoint"})
-    force_download: bool = field(default=False, metadata={"help": "Force the download of the model"})
-    resume_download: bool = field(default=False, metadata={"help": "Resume an interrupted download"})
-    proxies: Optional[str] = field(default=None, metadata={"help": "Proxy configuration for downloading the model"})
-    output_loading_info: bool = field(default=False, metadata={"help": "Output additional loading information"})
-    allow_remote_files: bool = field(default=False, metadata={"help": "Allow using remote files, default is local only"})
-    revision: str = field(default="main", metadata={"help": "Specific model version to use"})
-    trust_remote_code: bool = field(default=False, metadata={"help": "Enable trusting remote code when loading the model"})
-    load_in_4bit: bool = field(default=False, metadata={"help": "Load model in 4-bit mode"})
-    load_in_8bit: bool = field(default=False, metadata={"help": "Load model in 8-bit mode"})
-    torch_dtype: Optional[str] = field(default=None, metadata={"help": "The torch dtype for the pre-trained model"})
-    device_map: str = field(default="auto", metadata={"help": "The device map for the pre-trained model"})
-    chat_template: Optional[str] = field(default=None, metadata={"help": "The file path to the chat template, or the template in single-line form for the specified model"})
+
+    pipeline: str | None = field(
+        default="text-generation",
+        metadata={"help": "The model pipeline for the pre-trained model"},
+    )
+    pretrained_model_name_or_path: str | None = field(
+        default="/workspace/tfs/weights",
+        metadata={
+            "help": "Path to the pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    combination_type: str | None = field(
+        default="svd", metadata={"help": "The combination type of multi adapters"}
+    )
+    state_dict: dict[str, Any] | None = field(
+        default=None, metadata={"help": "State dictionary for the model"}
+    )
+    cache_dir: str | None = field(
+        default=None, metadata={"help": "Cache directory for the model"}
+    )
+    from_tf: bool = field(
+        default=False, metadata={"help": "Load model from a TensorFlow checkpoint"}
+    )
+    force_download: bool = field(
+        default=False, metadata={"help": "Force the download of the model"}
+    )
+    resume_download: bool = field(
+        default=False, metadata={"help": "Resume an interrupted download"}
+    )
+    proxies: str | None = field(
+        default=None, metadata={"help": "Proxy configuration for downloading the model"}
+    )
+    output_loading_info: bool = field(
+        default=False, metadata={"help": "Output additional loading information"}
+    )
+    allow_remote_files: bool = field(
+        default=False,
+        metadata={"help": "Allow using remote files, default is local only"},
+    )
+    revision: str = field(
+        default="main", metadata={"help": "Specific model version to use"}
+    )
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={"help": "Enable trusting remote code when loading the model"},
+    )
+    load_in_4bit: bool = field(
+        default=False, metadata={"help": "Load model in 4-bit mode"}
+    )
+    load_in_8bit: bool = field(
+        default=False, metadata={"help": "Load model in 8-bit mode"}
+    )
+    torch_dtype: str | None = field(
+        default=None, metadata={"help": "The torch dtype for the pre-trained model"}
+    )
+    device_map: str = field(
+        default="auto", metadata={"help": "The device map for the pre-trained model"}
+    )
+    chat_template: str | None = field(
+        default=None,
+        metadata={
+            "help": "The file path to the chat template, or the template in single-line form for the specified model"
+        },
+    )
 
     # Method to process additional arguments
-    def process_additional_args(self, addt_args: List[str]):
+    def process_additional_args(self, addt_args: list[str]):
         """
         Process additional cmd line args and update the model configuration accordingly.
         """
         addt_args_dict = {}
         i = 0
         while i < len(addt_args):
-            key = addt_args[i].lstrip('-')  # Remove leading dashes
-            if i + 1 < len(addt_args) and not addt_args[i + 1].startswith('--'):
+            key = addt_args[i].lstrip("-")  # Remove leading dashes
+            if i + 1 < len(addt_args) and not addt_args[i + 1].startswith("--"):
                 value = addt_args[i + 1]
                 i += 2  # Move past the current key-value pair
             else:
@@ -87,22 +138,29 @@ class ModelConfig:
         # Update the ModelConfig instance with the additional args
         self.__dict__.update(addt_args_dict)
 
-    def __post_init__(self): # validate parameters 
+    def __post_init__(self):  # validate parameters
         """
         Post-initialization to validate some ModelConfig values
         """
         if self.torch_dtype == "auto":
             pass
-        elif self.torch_dtype and self.torch_dtype != "auto" and not hasattr(torch, self.torch_dtype):
+        elif (
+            self.torch_dtype
+            and self.torch_dtype != "auto"
+            and not hasattr(torch, self.torch_dtype)
+        ):
             raise ValueError(f"Invalid torch dtype: {self.torch_dtype}")
         else:
-            self.torch_dtype = getattr(torch, self.torch_dtype) if self.torch_dtype else None
+            self.torch_dtype = (
+                getattr(torch, self.torch_dtype) if self.torch_dtype else None
+            )
 
         supported_pipelines = {"conversational", "text-generation"}
         if self.pipeline not in supported_pipelines:
             raise ValueError(f"Unsupported pipeline: {self.pipeline}")
 
-def load_chat_template(chat_template: Optional[str]) -> Optional[str]:
+
+def load_chat_template(chat_template: str | None) -> str | None:
     logger.info(chat_template)
     if chat_template is None:
         return None
@@ -126,12 +184,12 @@ args, additional_args = parser.parse_args_into_dataclasses(
 args.process_additional_args(additional_args)
 
 model_args = asdict(args)
-model_args["local_files_only"] = not model_args.pop('allow_remote_files')
-model_pipeline = model_args.pop('pipeline')
-combination_type = model_args.pop('combination_type')
+model_args["local_files_only"] = not model_args.pop("allow_remote_files")
+model_pipeline = model_args.pop("pipeline")
+combination_type = model_args.pop("combination_type")
 
 app = FastAPI()
-resovled_chat_template = load_chat_template(model_args.pop('chat_template'))
+resovled_chat_template = load_chat_template(model_args.pop("chat_template"))
 tokenizer = AutoTokenizer.from_pretrained(**model_args)
 if resovled_chat_template is not None:
     tokenizer.chat_template = resovled_chat_template
@@ -141,7 +199,8 @@ if not os.path.exists(ADAPTERS_DIR):
     model = base_model
 else:
     valid_adapters_list = [
-        os.path.join(ADAPTERS_DIR, adapter) for adapter in os.listdir(ADAPTERS_DIR)
+        os.path.join(ADAPTERS_DIR, adapter)
+        for adapter in os.listdir(ADAPTERS_DIR)
         if os.path.isfile(os.path.join(ADAPTERS_DIR, adapter, "adapter_config.json"))
     ]
 
@@ -150,9 +209,11 @@ else:
         for adapter_path in valid_adapters_list:
             adapter_name = os.path.basename(adapter_path)
             adapter_names.append(adapter_name)
-            weights.append(float(os.getenv(adapter_name, '1.0')))
+            weights.append(float(os.getenv(adapter_name, "1.0")))
 
-        model = PeftModel.from_pretrained(base_model, valid_adapters_list[0], adapter_name=adapter_names[0])
+        model = PeftModel.from_pretrained(
+            base_model, valid_adapters_list[0], adapter_name=adapter_names[0]
+        )
         for i in range(1, len(valid_adapters_list)):
             model.load_adapter(valid_adapters_list[i], adapter_name=adapter_names[i])
 
@@ -168,11 +229,11 @@ else:
         # To avoid any potential future operations that use non-combined adapters
         for adapter in adapter_names:
             model.delete_adapter(adapter)
-        
+
         active_adapters = model.active_adapters
         if len(active_adapters) != 1 or active_adapters[0] != "combined_adapter":
-            raise ValueError(f"Adapters not merged correctly")
-        logger.info("Adapter added: %s", ', '.join(sorted(adapter_names)))
+            raise ValueError("Adapters not merged correctly")
+        logger.info("Adapter added: %s", ", ".join(sorted(adapter_names)))
     else:
         logger.warning("Did not find any valid adapters mounted, using base model")
         model = base_model
@@ -189,24 +250,23 @@ if args.torch_dtype:
     pipeline_kwargs["torch_dtype"] = args.torch_dtype
 
 pipeline = transformers.pipeline(
-    task="text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    **pipeline_kwargs
+    task="text-generation", model=model, tokenizer=tokenizer, **pipeline_kwargs
 )
 
 try:
     # Attempt to load the generation configuration
     default_generate_config = GenerationConfig.from_pretrained(
-        args.pretrained_model_name_or_path,
-        local_files_only=args.local_files_only
+        args.pretrained_model_name_or_path, local_files_only=args.local_files_only
     ).to_dict()
-except Exception as e:
+except Exception:
     default_generate_config = {}
 
+
 class HomeResponse(BaseModel):
     message: str = Field(..., example="Server is running")
-@app.get('/', response_model=HomeResponse, summary="Home Endpoint")
+
+
+@app.get("/", response_model=HomeResponse, summary="Home Endpoint")
 def home():
     """
     A simple endpoint that indicates the server is running.
@@ -214,8 +274,11 @@ def home():
     """
     return {"message": "Server is running"}
 
+
 class HealthStatus(BaseModel):
     status: str = Field(..., example="Healthy")
+
+
 @app.get(
     "/health",
     response_model=HealthStatus,
@@ -223,11 +286,7 @@ class HealthStatus(BaseModel):
     responses={
         200: {
             "description": "Successful Response",
-            "content": {
-                "application/json": {
-                    "example": {"status": "Healthy"}
-                }
-            }
+            "content": {"application/json": {"example": {"status": "Healthy"}}},
         },
         500: {
             "description": "Error Response",
@@ -236,17 +295,17 @@ class HealthStatus(BaseModel):
                     "examples": {
                         "model_uninitialized": {
                             "summary": "Model not initialized",
-                            "value": {"detail": "Model not initialized"}
+                            "value": {"detail": "Model not initialized"},
                         },
                         "pipeline_uninitialized": {
                             "summary": "Pipeline not initialized",
-                            "value": {"detail": "Pipeline not initialized"}
-                        }
+                            "value": {"detail": "Pipeline not initialized"},
+                        },
                     }
                 }
-            }
-        }
-    }
+            },
+        },
+    },
 )
 def health_check():
     if not model:
@@ -257,8 +316,9 @@ def health_check():
         raise HTTPException(status_code=500, detail="Pipeline not initialized")
     return {"status": "Healthy"}
 
+
 class GenerateKwargs(BaseModel):
-    max_length: int = 200 # Length of input prompt+max_new_tokens
+    max_length: int = 200  # Length of input prompt+max_new_tokens
     min_length: int = 0
     do_sample: bool = True
     early_stopping: bool = False
@@ -268,40 +328,63 @@ class GenerateKwargs(BaseModel):
     top_p: float = 1
     typical_p: float = 1
     repetition_penalty: float = 1
-    pad_token_id: Optional[int] = tokenizer.pad_token_id
-    eos_token_id: Optional[int] = tokenizer.eos_token_id
+    pad_token_id: int | None = tokenizer.pad_token_id
+    eos_token_id: int | None = tokenizer.eos_token_id
+
     class Config:
-        extra = 'allow' # Allows for additional fields not explicitly defined
+        extra = "allow"  # Allows for additional fields not explicitly defined
         json_schema_extra = {
             "example": {
                 "max_length": 200,
                 "temperature": 0.7,
                 "top_p": 0.9,
-                "additional_param": "Example value"
+                "additional_param": "Example value",
             }
         }
 
+
 class Message(BaseModel):
     role: str
     content: str
 
+
 class UnifiedRequestModel(BaseModel):
     # Fields for text generation
-    prompt: Optional[str] = Field(None, description="Prompt for text generation. Required for text-generation pipeline. Do not use with 'messages'.")
-    return_full_text: Optional[bool] = Field(True, description="Return full text if True, else only added text")
-    clean_up_tokenization_spaces: Optional[bool] = Field(False, description="Clean up extra spaces in text output")
-    prefix: Optional[str] = Field(None, description="Prefix added to prompt")
-    handle_long_generation: Optional[str] = Field(None, description="Strategy to handle long generation")
-    generate_kwargs: Optional[GenerateKwargs] = Field(default_factory=GenerateKwargs, description="Additional kwargs for generate method")
+    prompt: str | None = Field(
+        None,
+        description="Prompt for text generation. Required for text-generation pipeline. Do not use with 'messages'.",
+    )
+    return_full_text: bool | None = Field(
+        True, description="Return full text if True, else only added text"
+    )
+    clean_up_tokenization_spaces: bool | None = Field(
+        False, description="Clean up extra spaces in text output"
+    )
+    prefix: str | None = Field(None, description="Prefix added to prompt")
+    handle_long_generation: str | None = Field(
+        None, description="Strategy to handle long generation"
+    )
+    generate_kwargs: GenerateKwargs | None = Field(
+        default_factory=GenerateKwargs,
+        description="Additional kwargs for generate method",
+    )
 
     # Field for conversational model
-    messages: Optional[List[Message]] = Field(None, description="Messages for conversational model. Required for conversational pipeline. Do not use with 'prompt'.")
+    messages: list[Message] | None = Field(
+        None,
+        description="Messages for conversational model. Required for conversational pipeline. Do not use with 'prompt'.",
+    )
+
     def messages_to_dict_list(self):
-        return [message.model_dump() for message in self.messages] if self.messages else []
+        return (
+            [message.model_dump() for message in self.messages] if self.messages else []
+        )
+
 
 class ErrorResponse(BaseModel):
     detail: str
 
+
 @app.post(
     "/chat",
     summary="Chat Endpoint",
@@ -313,19 +396,17 @@ class ErrorResponse(BaseModel):
                     "examples": {
                         "text_generation": {
                             "summary": "Text Generation Response",
-                            "value": {
-                                "Result": "Generated text based on the prompt."
-                            }
+                            "value": {"Result": "Generated text based on the prompt."},
                         },
                         "conversation": {
                             "summary": "Conversation Response",
                             "value": {
                                 "Result": "Response to the last message in the conversation."
-                            }
-                        }
+                            },
+                        },
                     }
                 }
-            }
+            },
         },
         400: {
             "model": ErrorResponse,
@@ -335,70 +416,86 @@ class ErrorResponse(BaseModel):
                     "examples": {
                         "missing_prompt": {
                             "summary": "Missing Prompt",
-                            "value": {"detail": "Text generation parameter prompt required"}
+                            "value": {
+                                "detail": "Text generation parameter prompt required"
+                            },
                         },
                         "missing_messages": {
                             "summary": "Missing Messages",
-                            "value": {"detail": "Conversational parameter messages required"}
-                        }
+                            "value": {
+                                "detail": "Conversational parameter messages required"
+                            },
+                        },
                     }
                 }
-            }
+            },
         },
-        500: {
-            "model": ErrorResponse,
-            "description": "Internal Server Error"
-        }
-    }
+        500: {"model": ErrorResponse, "description": "Internal Server Error"},
+    },
 )
 def generate_text(
-        request_model: Annotated[
-            UnifiedRequestModel,
-            Body(
-                openapi_examples={
-                    "text_generation_example": {
-                        "summary": "Text Generation Example",
-                        "description": "An example of a text generation request.",
-                        "value": {
-                            "prompt": "Tell me a joke",
-                            "return_full_text": True,
-                            "clean_up_tokenization_spaces": False,
-                            "prefix": None,
-                            "handle_long_generation": None,
-                            "generate_kwargs": GenerateKwargs().model_dump(),
-                        },
+    request_model: Annotated[
+        UnifiedRequestModel,
+        Body(
+            openapi_examples={
+                "text_generation_example": {
+                    "summary": "Text Generation Example",
+                    "description": "An example of a text generation request.",
+                    "value": {
+                        "prompt": "Tell me a joke",
+                        "return_full_text": True,
+                        "clean_up_tokenization_spaces": False,
+                        "prefix": None,
+                        "handle_long_generation": None,
+                        "generate_kwargs": GenerateKwargs().model_dump(),
                     },
-                    "conversation_example": {
-                        "summary": "Conversation Example",
-                        "description": "An example of a conversational request.",
-                        "value": {
-                            "messages": [
-                                {"role": "user", "content": "What is your favourite condiment?"},
-                                {"role": "assistant", "content": "Well, im quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever im cooking up in the kitchen!"},
-                                {"role": "user", "content": "Do you have mayonnaise recipes?"}
-                            ],
-                            "return_full_text": True,
-                            "clean_up_tokenization_spaces": False,
-                            "prefix": None,
-                            "handle_long_generation": None,
-                            "generate_kwargs": GenerateKwargs().model_dump(),
-                        },
+                },
+                "conversation_example": {
+                    "summary": "Conversation Example",
+                    "description": "An example of a conversational request.",
+                    "value": {
+                        "messages": [
+                            {
+                                "role": "user",
+                                "content": "What is your favourite condiment?",
+                            },
+                            {
+                                "role": "assistant",
+                                "content": "Well, im quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever im cooking up in the kitchen!",
+                            },
+                            {
+                                "role": "user",
+                                "content": "Do you have mayonnaise recipes?",
+                            },
+                        ],
+                        "return_full_text": True,
+                        "clean_up_tokenization_spaces": False,
+                        "prefix": None,
+                        "handle_long_generation": None,
+                        "generate_kwargs": GenerateKwargs().model_dump(),
                     },
                 },
-            ),
-        ],
+            },
+        ),
+    ],
 ):
     """
     Processes chat requests, generating text based on the specified pipeline (text generation or conversational).
     Validates required parameters based on the pipeline and returns the generated text.
     """
-    user_generate_kwargs = request_model.generate_kwargs.model_dump() if request_model.generate_kwargs else {}
+    user_generate_kwargs = (
+        request_model.generate_kwargs.model_dump()
+        if request_model.generate_kwargs
+        else {}
+    )
     generate_kwargs = {**default_generate_config, **user_generate_kwargs}
 
     if args.pipeline == "text-generation":
         if not request_model.prompt:
             logger.error("Text generation parameter prompt required")
-            raise HTTPException(status_code=400, detail="Text generation parameter prompt required")
+            raise HTTPException(
+                status_code=400, detail="Text generation parameter prompt required"
+            )
         sequences = pipeline(
             request_model.prompt,
             # return_tensors=request_model.return_tensors,
@@ -407,25 +504,27 @@ def generate_text(
             clean_up_tokenization_spaces=request_model.clean_up_tokenization_spaces,
             prefix=request_model.prefix,
             handle_long_generation=request_model.handle_long_generation,
-            **generate_kwargs
+            **generate_kwargs,
         )
 
         result = ""
         for seq in sequences:
             logger.debug(f"Result: {seq['generated_text']}")
-            result += seq['generated_text']
+            result += seq["generated_text"]
 
         return {"Result": result}
 
     elif args.pipeline == "conversational":
         if not request_model.messages:
             logger.error("Conversational parameter messages required")
-            raise HTTPException(status_code=400, detail="Conversational parameter messages required")
+            raise HTTPException(
+                status_code=400, detail="Conversational parameter messages required"
+            )
 
         response = pipeline(
             request_model.messages_to_dict_list(),
             clean_up_tokenization_spaces=request_model.clean_up_tokenization_spaces,
-            **generate_kwargs
+            **generate_kwargs,
         )
         return {"Result": str(response[-1])}
 
@@ -433,16 +532,19 @@ def generate_text(
         logger.error("Invalid pipeline type")
         raise HTTPException(status_code=400, detail="Invalid pipeline type")
 
+
 class MemoryInfo(BaseModel):
     used: str
     total: str
 
+
 class CPUInfo(BaseModel):
     load_percentage: float
     physical_cores: int
     total_cores: int
     memory: MemoryInfo
 
+
 class GPUInfo(BaseModel):
     id: str
     name: str
@@ -450,9 +552,11 @@ class GPUInfo(BaseModel):
     temperature: str
     memory: MemoryInfo
 
+
 class MetricsResponse(BaseModel):
-    gpu_info: Optional[List[GPUInfo]] = None
-    cpu_info: Optional[CPUInfo] = None
+    gpu_info: list[GPUInfo] | None = None
+    cpu_info: CPUInfo | None = None
+
 
 @app.get(
     "/metrics",
@@ -467,26 +571,42 @@ class MetricsResponse(BaseModel):
                         "gpu_metrics": {
                             "summary": "Example when GPUs are available",
                             "value": {
-                                "gpu_info": [{"id": "GPU-1234", "name": "GeForce GTX 950", "load": "25.00%", "temperature": "55 C", "memory": {"used": "1.00 GB", "total": "2.00 GB"}}],
-                                "cpu_info": None  # Indicates CPUs info might not be present when GPUs are available
-                            }
+                                "gpu_info": [
+                                    {
+                                        "id": "GPU-1234",
+                                        "name": "GeForce GTX 950",
+                                        "load": "25.00%",
+                                        "temperature": "55 C",
+                                        "memory": {
+                                            "used": "1.00 GB",
+                                            "total": "2.00 GB",
+                                        },
+                                    }
+                                ],
+                                "cpu_info": None,  # Indicates CPUs info might not be present when GPUs are available
+                            },
                         },
                         "cpu_metrics": {
                             "summary": "Example when only CPU is available",
                             "value": {
                                 "gpu_info": None,  # Indicates GPU info might not be present when only CPU is available
-                                "cpu_info": {"load_percentage": 20.0, "physical_cores": 4, "total_cores": 8, "memory": {"used": "4.00 GB", "total": "16.00 GB"}}
-                            }
-                        }
+                                "cpu_info": {
+                                    "load_percentage": 20.0,
+                                    "physical_cores": 4,
+                                    "total_cores": 8,
+                                    "memory": {"used": "4.00 GB", "total": "16.00 GB"},
+                                },
+                            },
+                        },
                     }
                 }
-            }
+            },
         },
         500: {
             "description": "Internal Server Error",
             "model": ErrorResponse,
-        }
-    }
+        },
+    },
 )
 def get_metrics():
     """
@@ -496,16 +616,19 @@ def get_metrics():
     try:
         if torch.cuda.is_available():
             gpus = GPUtil.getGPUs()
-            gpu_info = [GPUInfo(
-                id=str(gpu.id),
-                name=gpu.name,
-                load=f"{gpu.load * 100:.2f}%",
-                temperature=f"{gpu.temperature} C",
-                memory=MemoryInfo(
-                    used=f"{gpu.memoryUsed / (1024 ** 3):.2f} GB",
-                    total=f"{gpu.memoryTotal / (1024 ** 3):.2f} GB"
+            gpu_info = [
+                GPUInfo(
+                    id=str(gpu.id),
+                    name=gpu.name,
+                    load=f"{gpu.load * 100:.2f}%",
+                    temperature=f"{gpu.temperature} C",
+                    memory=MemoryInfo(
+                        used=f"{gpu.memoryUsed / (1024**3):.2f} GB",
+                        total=f"{gpu.memoryTotal / (1024**3):.2f} GB",
+                    ),
                 )
-            ) for gpu in gpus]
+                for gpu in gpus
+            ]
             return MetricsResponse(gpu_info=gpu_info)
         else:
             # Gather CPU metrics
@@ -514,26 +637,28 @@ def get_metrics():
             total_cores = psutil.cpu_count(logical=True)
             virtual_memory = psutil.virtual_memory()
             memory = MemoryInfo(
-                used=f"{virtual_memory.used / (1024 ** 3):.2f} GB",
-                total=f"{virtual_memory.total / (1024 ** 3):.2f} GB"
+                used=f"{virtual_memory.used / (1024**3):.2f} GB",
+                total=f"{virtual_memory.total / (1024**3):.2f} GB",
             )
             cpu_info = CPUInfo(
                 load_percentage=cpu_usage,
                 physical_cores=physical_cores,
                 total_cores=total_cores,
-                memory=memory
+                memory=memory,
             )
             return MetricsResponse(cpu_info=cpu_info)
     except Exception as e:
         logger.error(f"Error fetching metrics: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+
 def shutdown_handler(sig, frame):
     sys.exit(0)
 
+
 if __name__ == "__main__":
     signal.signal(signal.SIGINT, shutdown_handler)
-    local_rank = int(os.environ.get("LOCAL_RANK", 0)) # Default to 0 if not set
-    port = 5000 + local_rank # Adjust port based on local rank
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))  # Default to 0 if not set
+    port = 5000 + local_rank  # Adjust port based on local rank
     logger.info(f"Starting server on port {port}")
-    uvicorn.run(app=app, host='0.0.0.0', port=port)
+    uvicorn.run(app=app, host="0.0.0.0", port=port)
diff --git a/presets/workspace/inference/vllm/inference_api.py b/presets/workspace/inference/vllm/inference_api.py
index 91dc6b2..f87e36b 100644
--- a/presets/workspace/inference/vllm/inference_api.py
+++ b/presets/workspace/inference/vllm/inference_api.py
@@ -11,30 +11,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
-import gc
-import os
 import argparse
-from typing import Callable, Optional, List, Any
-import yaml
 import copy
+import gc
+import logging
+import os
+from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Any
 
-import uvloop
+import psutil
 import torch
-from vllm.utils import FlexibleArgumentParser
+import uvloop
 import vllm.entrypoints.openai.api_server as api_server
+import vllm.envs as envs
+import yaml
+from vllm.engine.llm_engine import EngineArgs, LLMEngine, VllmConfig
 from vllm.entrypoints.openai.serving_models import LoRAModulePath
-from vllm.engine.llm_engine import (LLMEngine, EngineArgs, VllmConfig)
 from vllm.executor.executor_base import ExecutorBase
+from vllm.utils import FlexibleArgumentParser
 
 # Initialize logger
 logger = logging.getLogger(__name__)
-debug_mode = os.environ.get('DEBUG_MODE', 'false').lower() == 'true'
+debug_mode = os.environ.get("DEBUG_MODE", "false").lower() == "true"
 logging.basicConfig(
     level=logging.DEBUG if debug_mode else logging.INFO,
-    format='%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s',
-    datefmt='%m-%d %H:%M:%S')
+    format="%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s",
+    datefmt="%m-%d %H:%M:%S",
+)
+
 
 class KAITOArgumentParser(argparse.ArgumentParser):
     vllm_parser = FlexibleArgumentParser(description="vLLM serving server")
@@ -48,13 +53,32 @@ class KAITOArgumentParser(argparse.ArgumentParser):
 
         # KAITO only args
         # They should start with "kaito-" prefix to avoid conflict with vllm args
-        self.add_argument("--kaito-adapters-dir", type=str, default="/mnt/adapter", help="Directory where adapters are stored in KAITO preset.")
-        self.add_argument("--kaito-config-file", type=str, default="", help="Additional args for KAITO preset.")
-        self.add_argument("--kaito-max-probe-steps", type=int, help="Maximum number of steps to find the max available seq len fitting in the GPU memory.")
+        self.add_argument(
+            "--kaito-adapters-dir",
+            type=str,
+            default="/mnt/adapter",
+            help="Directory where adapters are stored in KAITO preset.",
+        )
+        self.add_argument(
+            "--kaito-config-file",
+            type=str,
+            default="",
+            help="Additional args for KAITO preset.",
+        )
+        self.add_argument(
+            "--kaito-max-probe-steps",
+            type=int,
+            help="Maximum number of steps to find the max available seq len fitting in the GPU memory.",
+        )
+        self.add_argument(
+            "--kaito-kv-cache-cpu-memory-utilization",
+            type=float,
+            default=0.5,
+            help="KV cache CPU memory utilization.",
+        )
 
     def _reset_vllm_defaults(self):
-        local_rank = int(os.environ.get("LOCAL_RANK",
-                                        0))  # Default to 0 if not set
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))  # Default to 0 if not set
         port = 5000 + local_rank  # Adjust port based on local rank
 
         server_default_args = {
@@ -70,20 +94,24 @@ class KAITOArgumentParser(argparse.ArgumentParser):
             "gpu_memory_utilization": 0.95,
             "swap_space": 4,
             "disable_log_stats": False,
-            "uvicorn_log_level": "error"
+            "uvicorn_log_level": "error",
         }
         self.vllm_parser.set_defaults(**engine_default_args)
 
     def parse_args(self, *args, **kwargs):
         args = super().parse_known_args(*args, **kwargs)
         kaito_args = args[0]
-        runtime_args = args[1] # Remaining args
+        runtime_args = args[1]  # Remaining args
 
         # Load KAITO config
         if kaito_args.kaito_config_file:
             file_config = KaitoConfig.from_yaml(kaito_args.kaito_config_file)
             if kaito_args.kaito_max_probe_steps is None:
                 kaito_args.kaito_max_probe_steps = file_config.max_probe_steps
+            if kaito_args.kaito_kv_cache_cpu_memory_utilization is None:
+                kaito_args.kaito_kv_cache_cpu_memory_utilization = (
+                    file_config.kv_cache_cpu_memory_utilization
+                )
 
             for key, value in file_config.vllm.items():
                 runtime_args.append(f"--{key}")
@@ -98,6 +126,7 @@ class KAITOArgumentParser(argparse.ArgumentParser):
         print("\norignal vLLM server arguments:\n")
         self.vllm_parser.print_help(file)
 
+
 @dataclass
 class KaitoConfig:
     # Extra arguments for the vllm serving server, will be forwarded to the vllm server.
@@ -107,20 +136,27 @@ class KaitoConfig:
     # Maximum number of steps to find the max available seq len fitting in the GPU memory.
     max_probe_steps: int
 
+    # Optional: CPU memory utilization for the vllm engine in kv cache offload mode. (default: 0.5, set to 0 to disable)
+    kv_cache_cpu_memory_utilization: float
+
     @staticmethod
-    def from_yaml(yaml_file: str) -> 'KaitoConfig':
-        with open(yaml_file, 'r') as file:
+    def from_yaml(yaml_file: str) -> "KaitoConfig":
+        with open(yaml_file) as file:
             config_data = yaml.safe_load(file)
         return KaitoConfig(
-            vllm=config_data.get('vllm', {}),
-            max_probe_steps=config_data.get('max_probe_steps', 6)
+            vllm=config_data.get("vllm", {}),
+            max_probe_steps=config_data.get("max_probe_steps", 6),
+            kv_cache_cpu_memory_utilization=config_data.get(
+                "kv_cache_cpu_memory_utilization", 0.5
+            ),
         )
 
     def to_yaml(self) -> str:
         return yaml.dump(self.__dict__)
 
-def load_lora_adapters(adapters_dir: str) -> Optional[LoRAModulePath]:
-    lora_list: List[LoRAModulePath] = []
+
+def load_lora_adapters(adapters_dir: str) -> LoRAModulePath | None:
+    lora_list: list[LoRAModulePath] = []
 
     if not os.path.exists(adapters_dir):
         return lora_list
@@ -133,6 +169,7 @@ def load_lora_adapters(adapters_dir: str) -> Optional[LoRAModulePath]:
 
     return lora_list
 
+
 def find_max_available_seq_len(vllm_config: VllmConfig, max_probe_steps: int) -> int:
     """
     Load model and run profiler to find max available seq len.
@@ -143,7 +180,11 @@ def find_max_available_seq_len(vllm_config: VllmConfig, max_probe_steps: int) ->
         return vllm_config.model_config.max_model_len
     executor = executor_class(vllm_config=vllm_config)
 
-    res = binary_search_with_limited_steps(vllm_config.model_config.max_model_len, max_probe_steps, lambda x: is_context_length_safe(executor, x))
+    res = binary_search_with_limited_steps(
+        vllm_config.model_config.max_model_len,
+        max_probe_steps,
+        lambda x: is_context_length_safe(executor, x),
+    )
 
     # release memory
     del executor
@@ -152,7 +193,10 @@ def find_max_available_seq_len(vllm_config: VllmConfig, max_probe_steps: int) ->
 
     return res
 
-def binary_search_with_limited_steps(upper: int, max_probe_steps: int, is_valid_fn: Callable[[int], bool]) -> int:
+
+def binary_search_with_limited_steps(
+    upper: int, max_probe_steps: int, is_valid_fn: Callable[[int], bool]
+) -> int:
     """
     Finds the maximum valid value with limited number of steps.
 
@@ -182,6 +226,7 @@ def binary_search_with_limited_steps(upper: int, max_probe_steps: int, is_valid_
 
     return low
 
+
 def is_context_length_safe(executor: ExecutorBase, context_length: int) -> bool:
     """
     Check if the avilable gpu blocks is enough for the given num_gpu_blocks.
@@ -190,16 +235,19 @@ def is_context_length_safe(executor: ExecutorBase, context_length: int) -> bool:
     # number of tokens equal to max_num_batched_tokens.
     executor.scheduler_config.max_num_batched_tokens = context_length
     try:
-        logger.info(f"Try to determine available gpu blocks for context length {context_length}")
+        logger.info(
+            f"Try to determine available gpu blocks for context length {context_length}"
+        )
         # see https://github.com/vllm-project/vllm/blob/v0.7.2/vllm/engine/llm_engine.py#L416
         available_gpu_blocks, _ = executor.determine_num_available_blocks()
-    except torch.OutOfMemoryError as e:
-        return False    
+    except torch.OutOfMemoryError:
+        return False
 
     num_gpu_blocks = context_length // executor.cache_config.block_size
     return available_gpu_blocks >= num_gpu_blocks
 
-def try_get_max_available_seq_len(args: argparse.Namespace) -> Optional[int]:
+
+def try_get_max_available_seq_len(args: argparse.Namespace) -> int | None:
     if args.max_model_len is not None:
         logger.info(f"max_model_len is set to {args.max_model_len}, skip probing.")
         return None
@@ -230,9 +278,11 @@ def try_get_max_available_seq_len(args: argparse.Namespace) -> Optional[int]:
     available_seq_len = find_max_available_seq_len(vllm_config, max_probe_steps)
     # see https://github.com/vllm-project/vllm/blob/v0.7.2/vllm/worker/worker.py#L539
     if available_seq_len <= 0:
-        raise ValueError("No available memory for the cache blocks. "
-                        "Try increasing `gpu_memory_utilization` when "
-                        "initializing the engine.")
+        raise ValueError(
+            "No available memory for the cache blocks. "
+            "Try increasing `gpu_memory_utilization` when "
+            "initializing the engine."
+        )
 
     if available_seq_len != max_model_len:
         logger.info(f"Set max_model_len from {max_model_len} to {available_seq_len}")
@@ -241,8 +291,53 @@ def try_get_max_available_seq_len(args: argparse.Namespace) -> Optional[int]:
         logger.info(f"Using model default max_model_len {max_model_len}")
         return None
 
+
+def set_kv_cache_offloading_if_appliable(args: argparse.Namespace) -> None:
+    """
+    Set KV cache offloading to CPU RAM if applicable.
+    This is only applicable when VLLM_USE_V1 is enabled and
+    kaito_kv_cache_cpu_memory_utilization is set.
+    """
+    if (
+        not envs.is_set("VLLM_USE_V1")
+        and args.kaito_kv_cache_cpu_memory_utilization > 0
+    ):
+        logger.info(
+            f"VLLM_USE_V1 is not set, but kaito_kv_cache_cpu_memory_utilization is set as {args.kaito_kv_cache_cpu_memory_utilization}, "
+            "run create_engine_config to check whether VLLM_USE_V1 should be set."
+        )
+        EngineArgs.from_cli_args(copy.deepcopy(args)).create_engine_config()
+
+    if (
+        envs.is_set("VLLM_USE_V1")
+        and envs.VLLM_USE_V1
+        and args.kaito_kv_cache_cpu_memory_utilization > 0
+    ):
+        os.environ["LMCACHE_CHUNK_SIZE"] = "256"
+        os.environ["LMCACHE_LOCAL_CPU"] = "True"
+        available_memory_gb = (
+            psutil.virtual_memory().total - psutil.virtual_memory().used
+        ) / (1024**3)
+        logger.info(
+            f"VLLM_USE_V1 is set, Offload KV cache to CPU RAM, size limit: {available_memory_gb} * {args.kaito_kv_cache_cpu_memory_utilization} GB"
+        )
+        os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = (
+            f"{available_memory_gb * args.kaito_kv_cache_cpu_memory_utilization}"
+        )
+
+        if args.kv_transfer_config is None:
+            args.kv_transfer_config = {
+                "kv_connector": "LMCacheConnectorV1",
+                "kv_role": "kv_both",
+            }
+    else:
+        logger.info(
+            "VLLM_USE_V1 or kv_cache_cpu_memory_utilization is not set, do not use KV cache offload to CPU RAM."
+        )
+
+
 if __name__ == "__main__":
-    parser = KAITOArgumentParser(description='KAITO wrapper of vLLM serving server')
+    parser = KAITOArgumentParser(description="KAITO wrapper of vLLM serving server")
     args = parser.parse_args()
 
     # set LoRA adapters
@@ -254,6 +349,8 @@ if __name__ == "__main__":
     if max_available_seq_len is not None:
         args.max_model_len = max_available_seq_len
 
+    set_kv_cache_offloading_if_appliable(args)
+
     # Run the serving server
     logger.info(f"Starting server on port {args.port}")
     # See https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html for more
diff --git a/presets/workspace/inference/vllm/multi-node-health-check.py b/presets/workspace/inference/vllm/multi-node-health-check.py
index 99a90c8..28fac76 100755
--- a/presets/workspace/inference/vllm/multi-node-health-check.py
+++ b/presets/workspace/inference/vllm/multi-node-health-check.py
@@ -20,15 +20,29 @@ import sys
 
 def liveness(args):
     from ray.util.state import list_actors, list_jobs
+
     gcs = f"{args.leader_address}:{args.ray_port}"
 
     # Checking the Ray job's entrypoint is the best we can do to identify the VLLM inference job
-    inference_job = next(filter(lambda job: job.entrypoint.startswith("python3 /workspace/vllm/inference_api.py"), list_jobs(address=gcs)), None)
+    inference_job = next(
+        filter(
+            lambda job: job.entrypoint.startswith(
+                "python3 /workspace/vllm/inference_api.py"
+            ),
+            list_jobs(address=gcs),
+        ),
+        None,
+    )
     assert inference_job is not None, "Inference job not found in Ray jobs"
 
     has_dead_actors = False
-    for actor in list_actors(address=gcs, filters=[("job_id", "=", inference_job.job_id), ("state", "!=", "ALIVE")]):
-        logging.error(f"Ray actor {actor.actor_id} is {actor.state}: {actor.death_cause}")
+    for actor in list_actors(
+        address=gcs,
+        filters=[("job_id", "=", inference_job.job_id), ("state", "!=", "ALIVE")],
+    ):
+        logging.error(
+            f"Ray actor {actor.actor_id} is {actor.state}: {actor.death_cause}"
+        )
         has_dead_actors = True
 
     if has_dead_actors:
@@ -37,22 +51,39 @@ def liveness(args):
 
 def readiness(args):
     import requests
+
     vllm_health_endpoint = f"http://{args.leader_address}:{args.vllm_port}/health"
     try:
         response = requests.get(vllm_health_endpoint)
         if response.status_code != 200:
             sys.exit(1)
     except requests.exceptions.RequestException as e:
-        print(f"Get \"{vllm_health_endpoint}\": {e}")
+        print(f'Get "{vllm_health_endpoint}": {e}')
         sys.exit(1)
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Health check for VLLM multi-node setup")
-    parser.add_argument("probe", type=str, choices=["liveness", "readiness"], help="Type of health check probe")
-    parser.add_argument("--leader-address", type=str, required=True, help="Leader address of the Ray cluster")
-    parser.add_argument("--ray-port", type=int, default=6379, help="Ray port of the cluster")
-    parser.add_argument("--vllm-port", type=int, default=5000, help="VLLM port for the API server")
+    parser = argparse.ArgumentParser(
+        description="Health check for VLLM multi-node setup"
+    )
+    parser.add_argument(
+        "probe",
+        type=str,
+        choices=["liveness", "readiness"],
+        help="Type of health check probe",
+    )
+    parser.add_argument(
+        "--leader-address",
+        type=str,
+        required=True,
+        help="Leader address of the Ray cluster",
+    )
+    parser.add_argument(
+        "--ray-port", type=int, default=6379, help="Ray port of the cluster"
+    )
+    parser.add_argument(
+        "--vllm-port", type=int, default=5000, help="VLLM port for the API server"
+    )
     args = parser.parse_args()
 
     is_leader = os.environ.get("POD_INDEX") == "0"
@@ -61,8 +92,12 @@ if __name__ == "__main__":
     # restarting the leader node as well. In other words, we are delegating the worker's
     # liveness check to the leader node.
     if args.probe == "liveness" and is_leader:
-        assert args.ray_port is not None, "Ray port must be specified for liveness probe"
+        assert args.ray_port is not None, (
+            "Ray port must be specified for liveness probe"
+        )
         liveness(args)
     elif args.probe == "readiness":
-        assert args.vllm_port is not None, "VLLM port must be specified for readiness probe"
+        assert args.vllm_port is not None, (
+            "VLLM port must be specified for readiness probe"
+        )
         readiness(args)
diff --git a/presets/workspace/models/phi2/model.go b/presets/workspace/models/phi2/model.go
index f380612..6db4167 100644
--- a/presets/workspace/models/phi2/model.go
+++ b/presets/workspace/models/phi2/model.go
@@ -64,7 +64,7 @@ func (*phi2) GetInferenceParameters() *model.PresetParam {
 				InferenceMainFile: inference.DefaultTransformersMainFile,
 			},
 			VLLM: model.VLLMParam{
-				BaseCommand:    inference.DefaultVLLMCommand,
+				BaseCommand:    "VLLM_USE_V1=0 python3 /workspace/vllm/inference_api.py",
 				ModelName:      PresetPhi2Model,
 				ModelRunParams: phiRunParamsVLLM,
 			},
diff --git a/presets/workspace/models/supported_models.yaml b/presets/workspace/models/supported_models.yaml
index fbb71c8..e8f2084 100644
--- a/presets/workspace/models/supported_models.yaml
+++ b/presets/workspace/models/supported_models.yaml
@@ -3,8 +3,9 @@ models:
   - name: base
     type: text-generation
     runtime: tfs
-    tag: 0.0.4
+    tag: 0.0.5
     # Tag history:
+    # 0.0.5 - update chat-template for new models
     # 0.0.4 - bump vLLM to 0.9.0.
     # 0.0.3 - add depdenencies to support vLLM multi-node distributed inference
     # 0.0.2 - bump vLLM to 0.8.5. Tool chat template added.
diff --git a/presets/workspace/tuning/text-generation/cli.py b/presets/workspace/tuning/text-generation/cli.py
index 5f81432..7c845fe 100644
--- a/presets/workspace/tuning/text-generation/cli.py
+++ b/presets/workspace/tuning/text-generation/cli.py
@@ -13,81 +13,183 @@
 
 import os
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import torch
 from peft import LoraConfig
-from transformers import (BitsAndBytesConfig, DataCollatorForLanguageModeling,
-                          PreTrainedTokenizer)
+from transformers import (
+    BitsAndBytesConfig,
+    DataCollatorForLanguageModeling,
+    PreTrainedTokenizer,
+)
 
-target_modules_env = os.environ.get('DEFAULT_TARGET_MODULES', None) 
-DEFAULT_TARGET_MODULES = [module.strip() for module in target_modules_env.split(",")] if target_modules_env else None
+target_modules_env = os.environ.get("DEFAULT_TARGET_MODULES", None)
+DEFAULT_TARGET_MODULES = (
+    [module.strip() for module in target_modules_env.split(",")]
+    if target_modules_env
+    else None
+)
 
 # Consider Future Support for other trainers
 # class TrainerTypes(Enum):
-    # TRAINER = "Trainer"
-    # SFT_TRAINER = "SFTTrainer"
-    # DPO_TRAINER = "DPOTrainer"
-    # REWARD_TRAINER = "RewardTrainer"
-    # PPO_TRAINER = "PPOTrainer"
-    # CPO_TRAINER = "CPOTrainer"
-    # ORPO_TRAINER = "ORPOTrainer"
+# TRAINER = "Trainer"
+# SFT_TRAINER = "SFTTrainer"
+# DPO_TRAINER = "DPOTrainer"
+# REWARD_TRAINER = "RewardTrainer"
+# PPO_TRAINER = "PPOTrainer"
+# CPO_TRAINER = "CPOTrainer"
+# ORPO_TRAINER = "ORPOTrainer"
 
 # @dataclass
 # class TrainerType:
 #     trainer_type: TrainerTypes = field(default=TrainerTypes.SFT_TRAINER, metadata={"help": "Type of trainer to use for fine-tuning."})
 
+
 @dataclass
 class ExtDataCollator(DataCollatorForLanguageModeling):
-    tokenizer: Optional[PreTrainedTokenizer] = field(default=PreTrainedTokenizer, metadata={"help": "Tokenizer for DataCollatorForLanguageModeling"})
+    tokenizer: PreTrainedTokenizer | None = field(
+        default=PreTrainedTokenizer,
+        metadata={"help": "Tokenizer for DataCollatorForLanguageModeling"},
+    )
+
 
 @dataclass
 class ExtLoraConfig(LoraConfig):
     """
     Lora Config
     """
-    init_lora_weights: bool = field(default=True, metadata={"help": "Enable initialization of LoRA weights"})
-    target_modules: Optional[List[str]] = field(default_factory=lambda: DEFAULT_TARGET_MODULES if DEFAULT_TARGET_MODULES else None, metadata={"help": "List of module names to replace with LoRA."})
-    layers_to_transform: Optional[List[int]] = field(default=None, metadata={"help": "Layer indices to apply LoRA"})
-    layers_pattern: Optional[List[str]] = field(default=None, metadata={"help": "Pattern to match layers for LoRA"})
-    loftq_config: Dict[str, any] = field(default_factory=dict, metadata={"help": "LoftQ configuration for quantization"})
+
+    init_lora_weights: bool = field(
+        default=True, metadata={"help": "Enable initialization of LoRA weights"}
+    )
+    target_modules: list[str] | None = field(
+        default_factory=lambda: DEFAULT_TARGET_MODULES
+        if DEFAULT_TARGET_MODULES
+        else None,
+        metadata={"help": "List of module names to replace with LoRA."},
+    )
+    layers_to_transform: list[int] | None = field(
+        default=None, metadata={"help": "Layer indices to apply LoRA"}
+    )
+    layers_pattern: list[str] | None = field(
+        default=None, metadata={"help": "Pattern to match layers for LoRA"}
+    )
+    loftq_config: dict[str, any] = field(
+        default_factory=dict, metadata={"help": "LoftQ configuration for quantization"}
+    )
+
 
 @dataclass
 class DatasetConfig:
     """
     Config for Dataset
     """
-    dataset_path: Optional[str] = field(default=None, metadata={"help": "Where dataset file is located in the /data folder. This path will be appended to /data. This path should be the absolute path in the image or host."})
-    dataset_extension: Optional[str] = field(default=None, metadata={"help": "Optional explicit file extension of the dataset. If not provided, the extension will be derived from the dataset_path."})
-    shuffle_dataset: bool = field(default=True, metadata={"help": "Whether to shuffle dataset"})
+
+    dataset_path: str | None = field(
+        default=None,
+        metadata={
+            "help": "Where dataset file is located in the /data folder. This path will be appended to /data. This path should be the absolute path in the image or host."
+        },
+    )
+    dataset_extension: str | None = field(
+        default=None,
+        metadata={
+            "help": "Optional explicit file extension of the dataset. If not provided, the extension will be derived from the dataset_path."
+        },
+    )
+    shuffle_dataset: bool = field(
+        default=True, metadata={"help": "Whether to shuffle dataset"}
+    )
     shuffle_seed: int = field(default=42, metadata={"help": "Seed for shuffling data"})
     # instruction_column: Optional[str] = field(default=None, metadata={"help": "Optional column for detailed instructions, used in more structured tasks like Alpaca-style setups."}) # Consider including in V2
-    context_column: Optional[str] = field(default=None, metadata={"help": "Column for additional context or prompts, used for generating responses based on scenarios."})
-    response_column: str = field(default="text", metadata={"help": "Main text column for standalone entries or the response part in prompt-response setups."})
-    messages_column: Optional[str] = field(default=None, metadata={"help": "Column containing structured conversational data in JSON format with roles and content, used for chatbot training."})
-    train_test_split: float = field(default=0.8, metadata={"help": "Split between test and training data (e.g. 0.8 means 80/20% train/test split)"})
+    context_column: str | None = field(
+        default=None,
+        metadata={
+            "help": "Column for additional context or prompts, used for generating responses based on scenarios."
+        },
+    )
+    response_column: str = field(
+        default="text",
+        metadata={
+            "help": "Main text column for standalone entries or the response part in prompt-response setups."
+        },
+    )
+    messages_column: str | None = field(
+        default=None,
+        metadata={
+            "help": "Column containing structured conversational data in JSON format with roles and content, used for chatbot training."
+        },
+    )
+    train_test_split: float = field(
+        default=0.8,
+        metadata={
+            "help": "Split between test and training data (e.g. 0.8 means 80/20% train/test split)"
+        },
+    )
+
 
 @dataclass
 class ModelConfig:
     """
     Transformers Model Configuration Parameters
     """
-    pretrained_model_name_or_path: Optional[str] = field(default="/workspace/tfs/weights", metadata={"help": "Path to the pretrained model or model identifier from huggingface.co/models"})
-    state_dict: Optional[Dict[str, Any]] = field(default=None, metadata={"help": "State dictionary for the model"})
-    cache_dir: Optional[str] = field(default=None, metadata={"help": "Cache directory for the model"})
-    from_tf: bool = field(default=False, metadata={"help": "Load model from a TensorFlow checkpoint"})
-    force_download: bool = field(default=False, metadata={"help": "Force the download of the model"})
-    resume_download: bool = field(default=False, metadata={"help": "Resume an interrupted download"})
-    proxies: Optional[str] = field(default=None, metadata={"help": "Proxy configuration for downloading the model"})
-    output_loading_info: bool = field(default=False, metadata={"help": "Output additional loading information"})
-    local_files_only: bool = field(default=False, metadata={"help": "Allow using remote files, default is local only"})
-    revision: str = field(default="main", metadata={"help": "Specific model version to use"})
-    trust_remote_code: bool = field(default=False, metadata={"help": "Enable trusting remote code when loading the model"})
-    load_in_4bit: bool = field(default=False, metadata={"help": "Load model in 4-bit mode"})
-    load_in_8bit: bool = field(default=False, metadata={"help": "Load model in 8-bit mode"})
-    torch_dtype: Optional[str] = field(default=None, metadata={"help": "The torch dtype for the pre-trained model"})
-    device_map: str = field(default="auto", metadata={"help": "The device map for the pre-trained model"})
-    chat_template: Optional[str] = field(default=None, metadata={"help": "The file path to the chat template, or the template in single-line form for the specified model"})
+
+    pretrained_model_name_or_path: str | None = field(
+        default="/workspace/tfs/weights",
+        metadata={
+            "help": "Path to the pretrained model or model identifier from huggingface.co/models"
+        },
+    )
+    state_dict: dict[str, Any] | None = field(
+        default=None, metadata={"help": "State dictionary for the model"}
+    )
+    cache_dir: str | None = field(
+        default=None, metadata={"help": "Cache directory for the model"}
+    )
+    from_tf: bool = field(
+        default=False, metadata={"help": "Load model from a TensorFlow checkpoint"}
+    )
+    force_download: bool = field(
+        default=False, metadata={"help": "Force the download of the model"}
+    )
+    resume_download: bool = field(
+        default=False, metadata={"help": "Resume an interrupted download"}
+    )
+    proxies: str | None = field(
+        default=None, metadata={"help": "Proxy configuration for downloading the model"}
+    )
+    output_loading_info: bool = field(
+        default=False, metadata={"help": "Output additional loading information"}
+    )
+    local_files_only: bool = field(
+        default=False,
+        metadata={"help": "Allow using remote files, default is local only"},
+    )
+    revision: str = field(
+        default="main", metadata={"help": "Specific model version to use"}
+    )
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={"help": "Enable trusting remote code when loading the model"},
+    )
+    load_in_4bit: bool = field(
+        default=False, metadata={"help": "Load model in 4-bit mode"}
+    )
+    load_in_8bit: bool = field(
+        default=False, metadata={"help": "Load model in 8-bit mode"}
+    )
+    torch_dtype: str | None = field(
+        default=None, metadata={"help": "The torch dtype for the pre-trained model"}
+    )
+    device_map: str = field(
+        default="auto", metadata={"help": "The device map for the pre-trained model"}
+    )
+    chat_template: str | None = field(
+        default=None,
+        metadata={
+            "help": "The file path to the chat template, or the template in single-line form for the specified model"
+        },
+    )
 
     def __post_init__(self):
         """
@@ -100,39 +202,67 @@ class ModelConfig:
                 raise ValueError(f"Invalid torch dtype: {self.torch_dtype}")
 
     def get_tokenizer_args(self):
-        return {k: v for k, v in self.__dict__.items() if k not in ["torch_dtype", "chat_template"]}
+        return {
+            k: v
+            for k, v in self.__dict__.items()
+            if k not in ["torch_dtype", "chat_template"]
+        }
 
     def get_model_args(self):
         return {k: v for k, v in self.__dict__.items() if k not in ["chat_template"]}
 
+
 @dataclass
 class QuantizationConfig(BitsAndBytesConfig):
     """
     Quanitization Configuration
     """
-    quant_method: str = field(default="bitsandbytes", metadata={"help": "Quantization Method {bitsandbytes,gptq,awq}"})
-    load_in_8bit: bool = field(default=False, metadata={"help": "Enable 8-bit quantization"})
-    load_in_4bit: bool = field(default=False, metadata={"help": "Enable 4-bit quantization"})
-    llm_int8_threshold: float = field(default=6.0, metadata={"help": "LLM.int8 threshold"})
-    llm_int8_skip_modules: List[str] = field(default=None, metadata={"help": "Modules to skip for 8-bit conversion"})
-    llm_int8_enable_fp32_cpu_offload: bool = field(default=False, metadata={"help": "Enable FP32 CPU offload for 8-bit"})
-    llm_int8_has_fp16_weight: bool = field(default=False, metadata={"help": "Use FP16 weights for LLM.int8"})
-    bnb_4bit_compute_dtype: str = field(default="float32", metadata={"help": "Compute dtype for 4-bit quantization"})
-    bnb_4bit_quant_type: str = field(default="fp4", metadata={"help": "Quantization type for 4-bit"})
-    bnb_4bit_use_double_quant: bool = field(default=False, metadata={"help": "Use double quantization for 4-bit"})
+
+    quant_method: str = field(
+        default="bitsandbytes",
+        metadata={"help": "Quantization Method {bitsandbytes,gptq,awq}"},
+    )
+    load_in_8bit: bool = field(
+        default=False, metadata={"help": "Enable 8-bit quantization"}
+    )
+    load_in_4bit: bool = field(
+        default=False, metadata={"help": "Enable 4-bit quantization"}
+    )
+    llm_int8_threshold: float = field(
+        default=6.0, metadata={"help": "LLM.int8 threshold"}
+    )
+    llm_int8_skip_modules: list[str] = field(
+        default=None, metadata={"help": "Modules to skip for 8-bit conversion"}
+    )
+    llm_int8_enable_fp32_cpu_offload: bool = field(
+        default=False, metadata={"help": "Enable FP32 CPU offload for 8-bit"}
+    )
+    llm_int8_has_fp16_weight: bool = field(
+        default=False, metadata={"help": "Use FP16 weights for LLM.int8"}
+    )
+    bnb_4bit_compute_dtype: str = field(
+        default="float32", metadata={"help": "Compute dtype for 4-bit quantization"}
+    )
+    bnb_4bit_quant_type: str = field(
+        default="fp4", metadata={"help": "Quantization type for 4-bit"}
+    )
+    bnb_4bit_use_double_quant: bool = field(
+        default=False, metadata={"help": "Use double quantization for 4-bit"}
+    )
+
 
 # class CheckpointCallback(TrainerCallback):
 #     def on_train_end(self, args, state, control, **kwargs):
 #         model_path = args.output_dir
 #         timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
 #         img_tag = f"ghcr.io/YOUR_USERNAME/LoRA-Adapter:{timestamp}"
-        
+
 #         # Write a file to indicate fine_tuning completion
 #         completion_indicator_path = os.path.join(model_path, "training_completed.txt")
 #         with open(completion_indicator_path, 'w') as f:
 #             f.write(f"Training completed at {timestamp}\n")
 #             f.write(f"Image Tag: {img_tag}\n")
 
-    # This method is called whenever a checkpoint is saved.
-    # def on_save(self, args, state, control, **kwargs):
-    #     docker_build_and_push()
+# This method is called whenever a checkpoint is saved.
+# def on_save(self, args, state, control, **kwargs):
+#     docker_build_and_push()
diff --git a/presets/workspace/tuning/text-generation/dataset.py b/presets/workspace/tuning/text-generation/dataset.py
index 47f9745..890d717 100644
--- a/presets/workspace/tuning/text-generation/dataset.py
+++ b/presets/workspace/tuning/text-generation/dataset.py
@@ -12,17 +12,19 @@
 # limitations under the License.
 
 import os
-from typing import Optional
 
 from datasets import load_dataset
 
-SUPPORTED_EXTENSIONS = {'csv', 'json', 'parquet', 'arrow', 'webdataset'}
+SUPPORTED_EXTENSIONS = {"csv", "json", "parquet", "arrow", "webdataset"}
+
 
 class DatasetManager:
     def __init__(self, config):
         self.config = config
         self.dataset = None
-        self.dataset_text_field = None # Set this field if dataset consists of singular text column
+        self.dataset_text_field = (
+            None  # Set this field if dataset consists of singular text column
+        )
 
     def check_dataset_loaded(self):
         if self.dataset is None:
@@ -30,7 +32,9 @@ class DatasetManager:
 
     def check_column_exists(self, column_name):
         if column_name not in self.dataset.column_names:
-            raise ValueError(f"Column '{column_name}' does not exist in the dataset. Available columns: {self.dataset.column_names}")
+            raise ValueError(
+                f"Column '{column_name}' does not exist in the dataset. Available columns: {self.dataset.column_names}"
+            )
 
     def select_and_rename_columns(self, columns_to_select, rename_map=None):
         self.dataset = self.dataset.select_columns(columns_to_select)
@@ -46,20 +50,32 @@ class DatasetManager:
         if self.config.dataset_path:
             dataset_path = os.path.join("/mnt", self.config.dataset_path.strip("/"))
         else:
-            dataset_path = self.find_valid_dataset(os.environ.get('DATASET_FOLDER_PATH', '/mnt/data'))
+            dataset_path = self.find_valid_dataset(
+                os.environ.get("DATASET_FOLDER_PATH", "/mnt/data")
+            )
             if not dataset_path:
                 raise ValueError("Unable to find a valid dataset file.")
 
-        file_ext = self.config.dataset_extension if self.config.dataset_extension else self.get_file_extension(dataset_path)
+        file_ext = (
+            self.config.dataset_extension
+            if self.config.dataset_extension
+            else self.get_file_extension(dataset_path)
+        )
         try:
-            self.dataset = load_dataset(file_ext, data_files=dataset_path, split="train")
-            print(f"Dataset loaded successfully from {dataset_path} with file type '{file_ext}'.")
+            self.dataset = load_dataset(
+                file_ext, data_files=dataset_path, split="train"
+            )
+            print(
+                f"Dataset loaded successfully from {dataset_path} with file type '{file_ext}'."
+            )
         except Exception as e:
             print(f"Error loading dataset: {e}")
-            raise ValueError(f"Unable to load dataset {dataset_path} with file type '{file_ext}'")
+            raise ValueError(
+                f"Unable to load dataset {dataset_path} with file type '{file_ext}'"
+            )
 
     def find_valid_dataset(self, data_dir):
-        """ Searches for a file with a valid dataset type in the given directory. """
+        """Searches for a file with a valid dataset type in the given directory."""
         for root, dirs, files in os.walk(data_dir):
             for file in files:
                 filename_lower = file.lower()  # Convert to lowercase once per filename
@@ -69,7 +85,7 @@ class DatasetManager:
         return None
 
     def get_file_extension(self, file_path):
-        """ Returns the file extension based on filetype guess or filename. """
+        """Returns the file extension based on filetype guess or filename."""
         filename_lower = os.path.basename(file_path).lower()
         for ext in SUPPORTED_EXTENSIONS:
             if ext in filename_lower:
@@ -87,10 +103,10 @@ class DatasetManager:
             raise ValueError("Train/Test split needs to be between 0 and 1")
         if self.config.train_test_split < 1:
             split_dataset = self.dataset.train_test_split(
-                test_size=1-self.config.train_test_split,
-                seed=self.config.shuffle_seed
+                test_size=1 - self.config.train_test_split,
+                seed=self.config.shuffle_seed,
             )
-            return split_dataset['train'], split_dataset['test']
+            return split_dataset["train"], split_dataset["test"]
         else:
             return self.dataset, None
 
diff --git a/presets/workspace/tuning/text-generation/fine_tuning.py b/presets/workspace/tuning/text-generation/fine_tuning.py
index 9ed63b6..ddf6eb8 100644
--- a/presets/workspace/tuning/text-generation/fine_tuning.py
+++ b/presets/workspace/tuning/text-generation/fine_tuning.py
@@ -15,35 +15,42 @@ import logging
 import os
 from dataclasses import asdict
 from datetime import datetime
-from parser import parse_configs, load_chat_template
-from cli import ModelConfig, ExtDataCollator, ExtLoraConfig, DatasetConfig
 
 import torch
 from accelerate import Accelerator
+from cli import DatasetConfig, ExtDataCollator, ExtLoraConfig, ModelConfig
 from dataset import DatasetManager
+from parser import load_chat_template, parse_configs
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          BitsAndBytesConfig, TrainingArguments,
-                          TrainerCallback, TrainerControl, TrainerState)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+    TrainingArguments,
+)
 from trl import SFTTrainer
 
 # Initialize logger
 logger = logging.getLogger(__name__)
-debug_mode = os.environ.get('DEBUG_MODE', 'false').lower() == 'true'
+debug_mode = os.environ.get("DEBUG_MODE", "false").lower() == "true"
 logging.basicConfig(
     level=logging.DEBUG if debug_mode else logging.INFO,
-    format='%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s',
-    datefmt='%m-%d %H:%M:%S')
+    format="%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s",
+    datefmt="%m-%d %H:%M:%S",
+)
 
-CONFIG_YAML = os.environ.get('YAML_FILE_PATH', '/mnt/config/training_config.yaml')
+CONFIG_YAML = os.environ.get("YAML_FILE_PATH", "/mnt/config/training_config.yaml")
 parsed_configs = parse_configs(CONFIG_YAML)
 
-model_config: ModelConfig = parsed_configs.get('ModelConfig')
-bnb_config: BitsAndBytesConfig = parsed_configs.get('QuantizationConfig') 
-ext_lora_config: ExtLoraConfig = parsed_configs.get('LoraConfig')
-ta_args: TrainingArguments = parsed_configs.get('TrainingArguments')
-ds_config: DatasetConfig = parsed_configs.get('DatasetConfig')
-dc_args: ExtDataCollator = parsed_configs.get('DataCollator')
+model_config: ModelConfig = parsed_configs.get("ModelConfig")
+bnb_config: BitsAndBytesConfig = parsed_configs.get("QuantizationConfig")
+ext_lora_config: ExtLoraConfig = parsed_configs.get("LoraConfig")
+ta_args: TrainingArguments = parsed_configs.get("TrainingArguments")
+ds_config: DatasetConfig = parsed_configs.get("DatasetConfig")
+dc_args: ExtDataCollator = parsed_configs.get("DataCollator")
 
 accelerator = Accelerator()
 
@@ -113,27 +120,32 @@ if ds_config.shuffle_dataset:
 
 train_dataset, eval_dataset = dm.split_dataset()
 
+
 class EmptyCacheCallback(TrainerCallback):
     def on_step_end(self, args, state: TrainerState, control: TrainerControl, **kwargs):
         torch.cuda.empty_cache()
         return control
+
+
 empty_cache_callback = EmptyCacheCallback()
 
 # Prepare for training
 torch.cuda.set_device(accelerator.process_index)
 torch.cuda.empty_cache()
 # Training the Model
-trainer: SFTTrainer = accelerator.prepare(SFTTrainer(
-    model=model,
-    tokenizer=tokenizer,
-    train_dataset=train_dataset,
-    eval_dataset=eval_dataset,
-    args=ta_args,
-    data_collator=dc_args,
-    dataset_text_field=dm.dataset_text_field,
-    callbacks=[empty_cache_callback]
-    # metrics = "tensorboard" or "wandb" # TODO
-))
+trainer: SFTTrainer = accelerator.prepare(
+    SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        args=ta_args,
+        data_collator=dc_args,
+        dataset_text_field=dm.dataset_text_field,
+        callbacks=[empty_cache_callback],
+        # metrics = "tensorboard" or "wandb" # TODO
+    )
+)
 trainer.train()
 os.makedirs(ta_args.output_dir, exist_ok=True)
 # only save the adapter weights
@@ -142,6 +154,8 @@ trainer.model.save_pretrained(ta_args.output_dir)
 # Write file to signify training completion
 timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
 logger.info("Fine-Tuning completed\n")
-completion_indicator_path = os.path.join(ta_args.output_dir, "fine_tuning_completed.txt")
-with open(completion_indicator_path, 'w') as f:
+completion_indicator_path = os.path.join(
+    ta_args.output_dir, "fine_tuning_completed.txt"
+)
+with open(completion_indicator_path, "w") as f:
     f.write(f"Fine-Tuning completed at {timestamp}\n")
diff --git a/presets/workspace/tuning/text-generation/metrics/metrics_server.py b/presets/workspace/tuning/text-generation/metrics/metrics_server.py
index c016e7b..3558a23 100644
--- a/presets/workspace/tuning/text-generation/metrics/metrics_server.py
+++ b/presets/workspace/tuning/text-generation/metrics/metrics_server.py
@@ -14,7 +14,6 @@
 
 import logging
 import os
-from typing import List, Optional
 
 import GPUtil
 import psutil
@@ -25,24 +24,28 @@ from pydantic import BaseModel
 
 # Initialize logger
 logger = logging.getLogger(__name__)
-debug_mode = os.environ.get('DEBUG_MODE', 'false').lower() == 'true'
+debug_mode = os.environ.get("DEBUG_MODE", "false").lower() == "true"
 logging.basicConfig(level=logging.DEBUG if debug_mode else logging.INFO)
 
 app = FastAPI()
 
+
 class ErrorResponse(BaseModel):
     detail: str
 
+
 class MemoryInfo(BaseModel):
     used: str
     total: str
 
+
 class CPUInfo(BaseModel):
     load_percentage: float
     physical_cores: int
     total_cores: int
     memory: MemoryInfo
 
+
 class GPUInfo(BaseModel):
     id: str
     name: str
@@ -50,9 +53,11 @@ class GPUInfo(BaseModel):
     temperature: str
     memory: MemoryInfo
 
+
 class MetricsResponse(BaseModel):
-    gpu_info: Optional[List[GPUInfo]] = None
-    cpu_info: Optional[CPUInfo] = None
+    gpu_info: list[GPUInfo] | None = None
+    cpu_info: CPUInfo | None = None
+
 
 @app.get(
     "/metrics",
@@ -67,26 +72,42 @@ class MetricsResponse(BaseModel):
                         "gpu_metrics": {
                             "summary": "Example when GPUs are available",
                             "value": {
-                                "gpu_info": [{"id": "GPU-1234", "name": "GeForce GTX 950", "load": "25.00%", "temperature": "55 C", "memory": {"used": "1.00 GB", "total": "2.00 GB"}}],
-                                "cpu_info": None  # Indicates CPUs info might not be present when GPUs are available
-                            }
+                                "gpu_info": [
+                                    {
+                                        "id": "GPU-1234",
+                                        "name": "GeForce GTX 950",
+                                        "load": "25.00%",
+                                        "temperature": "55 C",
+                                        "memory": {
+                                            "used": "1.00 GB",
+                                            "total": "2.00 GB",
+                                        },
+                                    }
+                                ],
+                                "cpu_info": None,  # Indicates CPUs info might not be present when GPUs are available
+                            },
                         },
                         "cpu_metrics": {
                             "summary": "Example when only CPU is available",
                             "value": {
                                 "gpu_info": None,  # Indicates GPU info might not be present when only CPU is available
-                                "cpu_info": {"load_percentage": 20.0, "physical_cores": 4, "total_cores": 8, "memory": {"used": "4.00 GB", "total": "16.00 GB"}}
-                            }
-                        }
+                                "cpu_info": {
+                                    "load_percentage": 20.0,
+                                    "physical_cores": 4,
+                                    "total_cores": 8,
+                                    "memory": {"used": "4.00 GB", "total": "16.00 GB"},
+                                },
+                            },
+                        },
                     }
                 }
-            }
+            },
         },
         500: {
             "description": "Internal Server Error",
             "model": ErrorResponse,
-        }
-    }
+        },
+    },
 )
 def get_metrics():
     """
@@ -96,16 +117,19 @@ def get_metrics():
     try:
         if torch.cuda.is_available():
             gpus = GPUtil.getGPUs()
-            gpu_info = [GPUInfo(
-                id=str(gpu.id),
-                name=gpu.name,
-                load=f"{gpu.load * 100:.2f}%",
-                temperature=f"{gpu.temperature} C",
-                memory=MemoryInfo(
-                    used=f"{gpu.memoryUsed / (1024 ** 3):.2f} GB",
-                    total=f"{gpu.memoryTotal / (1024 ** 3):.2f} GB"
+            gpu_info = [
+                GPUInfo(
+                    id=str(gpu.id),
+                    name=gpu.name,
+                    load=f"{gpu.load * 100:.2f}%",
+                    temperature=f"{gpu.temperature} C",
+                    memory=MemoryInfo(
+                        used=f"{gpu.memoryUsed / (1024**3):.2f} GB",
+                        total=f"{gpu.memoryTotal / (1024**3):.2f} GB",
+                    ),
                 )
-            ) for gpu in gpus]
+                for gpu in gpus
+            ]
             return MetricsResponse(gpu_info=gpu_info)
         else:
             # Gather CPU metrics
@@ -114,22 +138,23 @@ def get_metrics():
             total_cores = psutil.cpu_count(logical=True)
             virtual_memory = psutil.virtual_memory()
             memory = MemoryInfo(
-                used=f"{virtual_memory.used / (1024 ** 3):.2f} GB",
-                total=f"{virtual_memory.total / (1024 ** 3):.2f} GB"
+                used=f"{virtual_memory.used / (1024**3):.2f} GB",
+                total=f"{virtual_memory.total / (1024**3):.2f} GB",
             )
             cpu_info = CPUInfo(
                 load_percentage=cpu_usage,
                 physical_cores=physical_cores,
                 total_cores=total_cores,
-                memory=memory
+                memory=memory,
             )
             return MetricsResponse(cpu_info=cpu_info)
     except Exception as e:
         logger.error(f"Error fetching metrics: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+
 if __name__ == "__main__":
-    local_rank = int(os.environ.get("LOCAL_RANK", 0)) # Default to 0 if not set
-    port = 5000 + local_rank # Adjust port based on local rank
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))  # Default to 0 if not set
+    port = 5000 + local_rank  # Adjust port based on local rank
     logger.info(f"Starting server on port {port}")
-    uvicorn.run(app=app, host='0.0.0.0', port=port)
+    uvicorn.run(app=app, host="0.0.0.0", port=port)
diff --git a/presets/workspace/tuning/text-generation/metrics/test_metrics_server.py b/presets/workspace/tuning/text-generation/metrics/test_metrics_server.py
index 61ff0aa..8c8bddd 100644
--- a/presets/workspace/tuning/text-generation/metrics/test_metrics_server.py
+++ b/presets/workspace/tuning/text-generation/metrics/test_metrics_server.py
@@ -17,42 +17,50 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 from fastapi.testclient import TestClient
-from metrics_server import CPUInfo, GPUInfo, MemoryInfo, MetricsResponse, app
+from metrics_server import app
 
 client = TestClient(app)
 
+
 def mock_getGPUs():
     class MockGPU:
         id = "GPU-1234"
         name = "GeForce GTX 950"
         load = 0.25  # 25%
         temperature = 55  # 55 C
-        memoryUsed = 1 * (1024 ** 3)  # 1 GB
-        memoryTotal = 2 * (1024 ** 3)  # 2 GB
+        memoryUsed = 1 * (1024**3)  # 1 GB
+        memoryTotal = 2 * (1024**3)  # 2 GB
 
     return [MockGPU()]
 
+
 def mock_cpu_percent(interval=1, percpu=False):
     return 20.0
 
+
 def mock_cpu_count(logical=True):
     return 8 if logical else 4
 
+
 def mock_virtual_memory():
     svmem = MagicMock()
-    svmem.used = 4 * (1024 ** 3)  # 4 GB
-    svmem.total = 16 * (1024 ** 3)  # 16 GB
+    svmem.used = 4 * (1024**3)  # 4 GB
+    svmem.total = 16 * (1024**3)  # 16 GB
     return svmem
 
+
 @pytest.fixture
 def set_env_vars():
     os.environ["DEBUG_MODE"] = "true"
     yield
     os.environ.pop("DEBUG_MODE")
 
+
 def test_metrics_endpoint_with_gpu(set_env_vars):
-    with patch("GPUtil.getGPUs", side_effect=mock_getGPUs), \
-         patch("torch.cuda.is_available", return_value=True):
+    with (
+        patch("GPUtil.getGPUs", side_effect=mock_getGPUs),
+        patch("torch.cuda.is_available", return_value=True),
+    ):
         response = client.get("/metrics")
         print("response:", response)
         print("response data:", response.json())
@@ -67,11 +75,14 @@ def test_metrics_endpoint_with_gpu(set_env_vars):
         assert data["gpu_info"][0]["memory"]["total"] == "2.00 GB"
         assert data["cpu_info"] is None
 
+
 def test_metrics_endpoint_with_cpu(set_env_vars):
-    with patch("torch.cuda.is_available", return_value=False), \
-         patch("psutil.cpu_percent", side_effect=mock_cpu_percent), \
-         patch("psutil.cpu_count", side_effect=mock_cpu_count), \
-         patch("psutil.virtual_memory", side_effect=mock_virtual_memory):
+    with (
+        patch("torch.cuda.is_available", return_value=False),
+        patch("psutil.cpu_percent", side_effect=mock_cpu_percent),
+        patch("psutil.cpu_count", side_effect=mock_cpu_count),
+        patch("psutil.virtual_memory", side_effect=mock_virtual_memory),
+    ):
         response = client.get("/metrics")
         assert response.status_code == 200
         data = response.json()
@@ -83,6 +94,7 @@ def test_metrics_endpoint_with_cpu(set_env_vars):
         assert data["cpu_info"]["memory"]["total"] == "16.00 GB"
         assert data["gpu_info"] is None
 
+
 def test_metrics_endpoint_failure(set_env_vars):
     with patch("torch.cuda.is_available", side_effect=Exception("Test Exception")):
         response = client.get("/metrics")
diff --git a/presets/workspace/tuning/text-generation/parser.py b/presets/workspace/tuning/text-generation/parser.py
index 2d949da..99ca71f 100644
--- a/presets/workspace/tuning/text-generation/parser.py
+++ b/presets/workspace/tuning/text-generation/parser.py
@@ -11,39 +11,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import codecs
 import logging
 from dataclasses import asdict, fields
-import codecs
 from pathlib import Path
-from typing import Optional
 
 import yaml
-from cli import (DatasetConfig, ExtDataCollator, ExtLoraConfig, ModelConfig, QuantizationConfig)
+from cli import (
+    DatasetConfig,
+    ExtDataCollator,
+    ExtLoraConfig,
+    ModelConfig,
+    QuantizationConfig,
+)
 from transformers import HfArgumentParser, TrainingArguments
 
 logger = logging.getLogger(__name__)
 
 # Mapping from config section names to data classes
 CONFIG_CLASS_MAP = {
-    'ModelConfig': ModelConfig,
-    'QuantizationConfig': QuantizationConfig,
-    'LoraConfig': ExtLoraConfig,
-    'TrainingArguments': TrainingArguments,
-    'DatasetConfig': DatasetConfig,
-    'DataCollator': ExtDataCollator,
+    "ModelConfig": ModelConfig,
+    "QuantizationConfig": QuantizationConfig,
+    "LoraConfig": ExtLoraConfig,
+    "TrainingArguments": TrainingArguments,
+    "DatasetConfig": DatasetConfig,
+    "DataCollator": ExtDataCollator,
 }
 
-def flatten_config_to_cli_args(config, prefix=''):
+
+def flatten_config_to_cli_args(config, prefix=""):
     cli_args = []
     for key, value in config.items():
         if isinstance(value, dict):
-            cli_args.extend(flatten_config_to_cli_args(value, prefix=f'{prefix}{key}_'))
+            cli_args.extend(flatten_config_to_cli_args(value, prefix=f"{prefix}{key}_"))
         elif isinstance(value, list):
-            cli_arg = f'--{prefix}{key}'
+            cli_arg = f"--{prefix}{key}"
             cli_args.append(cli_arg)
             cli_args.extend(map(str, value))
         else:
-            cli_arg = f'--{prefix}{key}'
+            cli_arg = f"--{prefix}{key}"
             cli_args.append(cli_arg)
             cli_args.append(str(value))
     return cli_args
@@ -60,15 +66,17 @@ def parse_section(section_name, section_config):
     parser = HfArgumentParser((CONFIG_CLASS_MAP[section_name],))
     try:
         # Try parsing normally
-        cli_args = flatten_config_to_cli_args(section_config, prefix='')
+        cli_args = flatten_config_to_cli_args(section_config, prefix="")
         instance = parser.parse_args_into_dataclasses(cli_args)[0]
     except AttributeError as e:
         if section_name == "DataCollator" and section_config.get("mlm", False):
-            print("Warning: Tokenizer does not have a mask token. Retrying with mlm=False.")
+            print(
+                "Warning: Tokenizer does not have a mask token. Retrying with mlm=False."
+            )
             # Update the section_config to set mlm to False
             section_config["mlm"] = False
             # Regenerate the CLI arguments from the updated config
-            cli_args = flatten_config_to_cli_args(section_config, prefix='')
+            cli_args = flatten_config_to_cli_args(section_config, prefix="")
             instance = parser.parse_args_into_dataclasses(cli_args)[0]
         else:
             # Reraise any other Error
@@ -81,25 +89,30 @@ def parse_section(section_name, section_config):
 
 def parse_configs(config_yaml):
     # Load the YAML configuration
-    with open(config_yaml, 'r') as file:
+    with open(config_yaml) as file:
         full_config = yaml.safe_load(file)
-    training_config = full_config.get('training_config', {})
+    training_config = full_config.get("training_config", {})
     print("training_config:", training_config)
 
     # Parse and merge configurations
     parsed_configs = {}
     for section_name, class_type in CONFIG_CLASS_MAP.items():
         # Parse section from YAML
-        yaml_parsed_instance = parse_section(section_name, training_config.get(section_name, {}))
+        yaml_parsed_instance = parse_section(
+            section_name, training_config.get(section_name, {})
+        )
         yaml_parsed_dict = asdict(yaml_parsed_instance)
         merged_config = yaml_parsed_dict
 
-        filtered_config = filter_unsupported_init_args(CONFIG_CLASS_MAP[section_name], merged_config)
+        filtered_config = filter_unsupported_init_args(
+            CONFIG_CLASS_MAP[section_name], merged_config
+        )
         parsed_configs[section_name] = CONFIG_CLASS_MAP[section_name](**filtered_config)
 
     return parsed_configs
 
-def load_chat_template(chat_template: Optional[str]) -> Optional[str]:
+
+def load_chat_template(chat_template: str | None) -> str | None:
     logger.info(chat_template)
     if chat_template is None:
         return None
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..e61700f
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,33 @@
+[tool.ruff]
+
+# Same as Black.
+line-length = 88
+indent-width = 4
+
+# Assume Python 3.12
+target-version = "py312"
+
+[tool.ruff.lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default.
+select = ["E", "F", "UP", "SIM", "I"]
+ignore = ["E501", "F821"]
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+
+[tool.ruff.format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Exclude Jupyter notebooks
+exclude = ["*.ipynb"]
+
diff --git a/terraform/variables.tf b/terraform/variables.tf
index d29ea46..958e211 100644
--- a/terraform/variables.tf
+++ b/terraform/variables.tf
@@ -12,7 +12,7 @@ variable "kaito_gpu_provisioner_version" {
 
 variable "kaito_workspace_version" {
   type        = string
-  default     = "0.5.1"
+  default     = "0.6.0"
   description = "kaito workspace version"
 }
 
@@ -30,6 +30,6 @@ variable "deploy_kaito_ragengine" {
 
 variable "kaito_ragengine_version" {
   type        = string
-  default     = "0.5.1"
+  default     = "0.6.0"
   description = "KAITO RAGEngine version"
 }
diff --git a/website/docs/aws.md b/website/docs/aws.md
index 7c5b72d..820a516 100644
--- a/website/docs/aws.md
+++ b/website/docs/aws.md
@@ -24,7 +24,7 @@ KAITO can use [Karpenter](https://karpenter.sh/) to automatically provision GPU
 Alternative: If you already have GPU nodes or manage them separately, use the preferred nodes approach in the [Quick Start](quick-start) instead.
 :::
 
-## Setup Auto-Provisioning
+## Set Up Auto-Provisioning
 
 ### Create EKS Cluster and install Karpenter
 
diff --git a/website/docs/azure.md b/website/docs/azure.md
index 6eb9114..4197a67 100644
--- a/website/docs/azure.md
+++ b/website/docs/azure.md
@@ -8,7 +8,7 @@ This guide covers setting up auto-provisioning capabilities for KAITO on Azure K
 
 - An AKS cluster with KAITO workspace controller installed
   - See [Step 1](#step-1-create-and-configure-an-aks-cluster) to create an AKS cluster
-  - See [Installation](installation) to installing the KAITO workspace controller
+  - See [Installation](installation) to install the KAITO workspace controller
 - [Azure CLI](https://learn.microsoft.com/cli/azure/install-azure-cli) for managing Azure resources
 - [kubectl](https://kubernetes.io/docs/tasks/tools/) configured to access your AKS cluster
 
@@ -32,7 +32,7 @@ Choose auto-provisioning when:
 Alternative: If you already have GPU nodes or manage them separately, use the [preferred nodes approach](quick-start#option-1-using-preferred-nodes-existing-gpu-nodes) instead.
 :::
 
-## Setup Auto-Provisioning
+## Set Up Auto-Provisioning
 
 ### Step 1: Create and configure an AKS Cluster
 
diff --git a/website/docs/contributing.md b/website/docs/contributing.md
index 2766aae..3aec963 100644
--- a/website/docs/contributing.md
+++ b/website/docs/contributing.md
@@ -40,7 +40,7 @@ EOF
 ```
 
 - `default_registry`: The container registry where Tilt pushes development images (e.g. `<acr_name>.azurecr.io`). Ensure you're logged in with `docker login` and have push/pull permissions.
-if you want to use Azure Container Registry (ACR), set up a new ACR instance and attach it to your AKS cluster with these commands:
+If you want to use Azure Container Registry (ACR), set up a new ACR instance and attach it to your AKS cluster with these commands:
 
 ```bash
 export ACR_NAME=<acr_name>
diff --git a/website/docs/example_rag_client.py b/website/docs/example_rag_client.py
index 60d5e83..e2a5760 100644
--- a/website/docs/example_rag_client.py
+++ b/website/docs/example_rag_client.py
@@ -11,9 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import requests
 import json
 
+import requests
+
+
 class KAITORAGClient:
     """
     A bare bones client for interacting with the RAGEngine API.
@@ -25,7 +27,7 @@ class KAITORAGClient:
         self.base_url = base_url
         self.headers = {
             "Content-Type": "application/json",
-            "Accept": "application/json"
+            "Accept": "application/json",
         }
 
     def index_documents(self, index_name, documents):
@@ -34,10 +36,7 @@ class KAITORAGClient:
         documents: list of dicts, each with 'text' and optional 'metadata'
         """
         url = f"{self.base_url}/index"
-        payload = {
-            "index_name": index_name,
-            "documents": documents
-        }
+        payload = {"index_name": index_name, "documents": documents}
         resp = requests.post(url, json=payload, headers=self.headers)
         resp.raise_for_status()
         return resp.json()
@@ -56,8 +55,8 @@ class KAITORAGClient:
             "top_k": top_k,
             "llm_params": {
                 "temperature": llm_temperature,
-                "max_tokens": llm_max_tokens
-            }
+                "max_tokens": llm_max_tokens,
+            },
         }
         resp = requests.post(url, json=payload, headers=self.headers)
         resp.raise_for_status()
@@ -89,7 +88,7 @@ class KAITORAGClient:
         List documents in the RAGEngine.
         """
         url = f"{self.base_url}/indexes/{index_name}/documents"
-        params = {"limit": limit, "offset": offset }
+        params = {"limit": limit, "offset": offset}
         if metadata_filter:
             params["metadata_filter"] = json.dumps(metadata_filter)
         resp = requests.get(url, headers=self.headers, params=params)
@@ -132,4 +131,4 @@ class KAITORAGClient:
         url = f"{self.base_url}/indexes/{index_name}"
         resp = requests.delete(url, headers=self.headers)
         resp.raise_for_status()
-        return resp.json()
\ No newline at end of file
+        return resp.json()
diff --git a/website/docs/installation.md b/website/docs/installation.md
index 7d15e3a..d76f8c0 100644
--- a/website/docs/installation.md
+++ b/website/docs/installation.md
@@ -20,7 +20,7 @@ Before you begin, ensure you have the following tools installed:
 Install the KAITO workspace controller using Helm:
 
 ```bash
-export KAITO_WORKSPACE_VERSION=0.5.1
+export KAITO_WORKSPACE_VERSION=0.6.0
 export CLUSTER_NAME=kaito
 
 helm install kaito-workspace \
@@ -52,12 +52,12 @@ You need to create GPU nodes in order to run a Workspace with KAITO. There are t
 ### Option 1: Bring your own GPU nodes
 
 :::tip
-This is recommended if you just want to try KAITO with the least amount of setup required or if you only plan to use 1-2 GPU nodes. If your cloud provider is not listed in the [Auto-provisioning section](#option-2-auto-provision-gpu-nodes), you must create your own GPU nodes. If it is, you can still manually create the nodes if you don't want to setup auto-provisioning.
+This is recommended if you just want to try KAITO with the least amount of setup required or if you only plan to use 1-2 GPU nodes. If your cloud provider is not listed in the [Auto-provisioning section](#option-2-auto-provision-gpu-nodes), you must create your own GPU nodes. If it is, you can still manually create the nodes if you don't want to set up auto-provisioning.
 :::
 
 Create your GPU nodes and label them for quick access. In these docs, we will use the label `accelerator=nvidia` but any label can work.
 
-List your your GPU nodes and verify that they are all present and ready.
+List your GPU nodes and verify that they are all present and ready.
 
 ```bash
 kubectl get nodes -l accelerator=nvidia
diff --git a/website/docs/intro.md b/website/docs/intro.md
index 9127b38..7b87767 100644
--- a/website/docs/intro.md
+++ b/website/docs/intro.md
@@ -5,7 +5,7 @@ slug: /
 
 :::info What's NEW!
 
-Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about from [here](https://kaito-project.github.io/kaito/docs/rag)!
+Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about it [here](https://kaito-project.github.io/kaito/docs/rag)!
 **Latest Release:** July 18th, 2025. KAITO v0.5.1.
 
 **First Release:** Nov 15th, 2023. KAITO v0.1.0.
@@ -46,7 +46,7 @@ The [*gpu-provisioner*](https://github.com/Azure/gpu-provisioner) is an open sou
 
 
 As illustrated in the above figure, the **RAGEngine controller** reconciles the `ragengine` custom resource and creates a `RAGService` deployment. The `RAGService` provides the following capabilities:
-  - **Orchstration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
+  - **Orchestration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
   - **Embedding**: support both local and remote embedding services, to embed queries and documents in the vector database.
   - **Vector database**: support a built-in [faiss](https://github.com/facebookresearch/faiss) in-memory vector database. Remote vector database support will be added soon.
   - **Backend inference**: support any OAI compatible inference service.
diff --git a/website/docs/kaito-on-byo-gpu-nodes.md b/website/docs/kaito-on-byo-gpu-nodes.md
index ef3c664..054871a 100644
--- a/website/docs/kaito-on-byo-gpu-nodes.md
+++ b/website/docs/kaito-on-byo-gpu-nodes.md
@@ -210,7 +210,7 @@ EOF
 ```
 
 :::note
-In the above configuration you can see we have use a node labelSelector value as `apps: gpu`, this is the same label we have applied when we added the GPU node pool earlier.
+In the above configuration you can see we have used a node labelSelector value as `apps: gpu`, this is the same label we have applied when we added the GPU node pool earlier.
 :::
 
 Ensure that the workspace is ready by running the following command:
@@ -240,7 +240,7 @@ export CLUSTERIP=$(kubectl get \
 Let's send a request to the workspace to get an inference response. Modify the prompt as you see fit:
 
 ```bash
-export QUESTION="What's are LLMs?"
+export QUESTION="What are LLMs?"
 ```
 
 Run the following command to send the request:
diff --git a/website/docs/monitoring.md b/website/docs/monitoring.md
index 0f74ef6..2fc88ca 100644
--- a/website/docs/monitoring.md
+++ b/website/docs/monitoring.md
@@ -23,7 +23,7 @@ vLLM V1 engine is now enabled by default. Check [here](https://docs.vllm.ai/en/l
 | System Stats - Cache | `vllm:gpu_cache_usage_perc` | Gauge | GPU KV-cache usage. 1 means 100 percent usage |
 | System Stats - Cache | `vllm:cpu_cache_usage_perc` | Gauge | CPU KV-cache usage. 1 means 100 percent usage (DEPRECATED: KV cache offloading is not used in V1) |
 | System Stats - Cache | `vllm:cpu_prefix_cache_hit_rate` | Gauge | CPU prefix cache block hit rate (DEPRECATED: KV cache offloading is not used in V1) |
-| System Stats - Cache | `vllm:gpu_prefix_cache_hit_rate` | Gauge | GPU prefix cache block hit rat (DEPRECATED: use `vllm:gpu_prefix_cache_queries` and `vllm:gpu_prefix_cache_hits in V1`) |
+| System Stats - Cache | `vllm:gpu_prefix_cache_hit_rate` | Gauge | GPU prefix cache block hit rate (DEPRECATED: use `vllm:gpu_prefix_cache_queries` and `vllm:gpu_prefix_cache_hits in V1`) |
 | System Stats - Cache | `vllm:gpu_cache_usage_perc` | Counter | GPU prefix cache queries, in terms of number of queried blocks (V1 only) |
 | System Stats - Cache | `vllm:gpu_prefix_cache_hits` | Counter | GPU prefix cache hits, in terms of number of cached blocks (V1 only) |
 | Iteration Stats | `vllm:num_preemptions_total` | Counter | Cumulative number of preemption from the engine |
diff --git a/website/docs/presets.md b/website/docs/presets.md
index 5688c7f..97c39dd 100644
--- a/website/docs/presets.md
+++ b/website/docs/presets.md
@@ -21,7 +21,7 @@ Each preset model has its own hardware requirements in terms of GPU count and GP
 
 ## Distributed inference
 
-For models that support distributed inference, when the node count is larger than one, [torch distributed elastic](https://pytorch.org/docs/stable/distributed.elastic.html) is configured with master/worker pods running in multiple nodes and the service endpoint is the master pod.
+For models that support distributed inference, when the node count is larger than one, [Torch Distributed Elastic](https://pytorch.org/docs/stable/distributed.elastic.html) is configured with master/worker pods running in multiple nodes and the service endpoint is the master pod.
 
 The following preset models support multi-node distributed inference:
 
diff --git a/website/docs/quick-start.mdx b/website/docs/quick-start.mdx
index 11ccc5e..978abc8 100644
--- a/website/docs/quick-start.mdx
+++ b/website/docs/quick-start.mdx
@@ -142,7 +142,7 @@ Track the workspace status to see when the model has been deployed successfully:
 kubectl get workspace workspace-phi-4-mini
 ```
 
-When the `WORKSPACEREADY` column becomes `True`, the model has been deployed successfully:
+When the `WORKSPACESUCCEEDED` column becomes `True`, the model has been deployed successfully:
 
 ```bash
 NAME                   INSTANCE                   RESOURCEREADY   INFERENCEREADY   JOBSTARTED   WORKSPACESUCCEEDED   AGE
diff --git a/website/docs/rag-api.md b/website/docs/rag-api.md
new file mode 100644
index 0000000..5550caf
--- /dev/null
+++ b/website/docs/rag-api.md
@@ -0,0 +1,423 @@
+---
+title: API Definitions and Examples
+---
+
+A **RAGEngine index** is a logical collection that organizes and stores your documents for retrieval-augmented generation workflows. The relationship between indexes, documents, and document nodes is as follows:
+
+- **Index**: An index is a named container that holds a set of documents. Each index is independent and can be created, updated, queried, persisted, loaded, or deleted via the API.
+
+- **Documents**: Documents are the primary units of content that you add to an index. Each document contains a `text` field (the content to be indexed) and optional `metadata` (such as author, source, or custom tags). When you add documents to an index, each is assigned a unique `doc_id`.
+
+- **Document Nodes**: When a document is ingested, it is automatically split into smaller chunks called *nodes*. The splitting strategy depends on the document type and metadata:
+  - By default, documents are split into sentences.
+  - If you specify code-aware splitting (using the `split_type` and `language` metadata), the document is split into code blocks or logical code units.
+  - Each node represents a chunk of text that is indexed and can be retrieved as part of a query.
+
+**How it works in practice:**
+- When you index a document, it is divided into nodes for efficient retrieval and semantic search.
+- When you query an index, the engine retrieves the most relevant nodes (not necessarily whole documents) and can use them to generate answers or summaries.
+- The `source_nodes` field in query responses contains the actual nodes that matched your query, along with their scores and metadata.
+
+This design enables fine-grained retrieval and more accurate, context-aware responses from your LLM-powered applications.
+
+## Creating an Index With Documents
+
+To add documents to an index or create a new index, use the `/index` API route. This endpoint accepts a POST request with the index name and a list of documents to be indexed.
+
+### Create Index Request
+
+```json
+POST /index
+{
+  "index_name": "rag_index",
+  "documents": [
+    {
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ]
+}
+```
+
+- `index_name`: The name of the index to create or update.
+- `documents`: A list of documents, each with a `text` field and optional `metadata`.
+
+### Create Index Response
+
+```json
+[
+  {
+    "doc_id": "123456",
+    "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
+    "hash_value": "text_hash_value",
+    "metadata": {
+      "author": "kaito",
+    },
+    "is_truncated": false
+  }
+]
+```
+
+Each returned document includes its unique `doc_id`, the original text, metadata, and a flag indicating if the text was truncated. The `doc_id` will be important for document update/delete calls. 
+
+
+### Splitting Documents with CodeSplitter
+
+By default, RAGEngine splits documents into sentences. However, you can instruct the engine to split documents using the `CodeSplitter` (for code-aware chunking) by providing metadata in your API request.
+
+To use the `CodeSplitter`, set the `split_type` to `"code"` and specify the programming language in the `language` field of the document metadata. For example, when calling the RAGEngine API to index documents:
+
+```json
+{
+  "documents": [
+    {
+      "text": "def foo():\n    return 42\n\n# Another function\ndef bar():\n    pass",
+      "metadata": {
+        "split_type": "code",
+        "language": "python"
+      }
+    }
+  ]
+}
+```
+
+This instructs the RAGEngine to use code-aware splitting for the provided document. If `split_type` is not set or set to any other value, sentence splitting will be used by default.
+
+## List Documents
+
+To retrieve a paginated list of documents from a specific index, use the `/indexes/{index_name}/documents` API route. This endpoint accepts a GET request with optional query parameters for pagination, text truncation, and metadata filtering.
+
+### List Documents Request
+
+```
+GET /indexes/rag_index/documents?limit=5&offset=0&max_text_length=500
+```
+
+- `limit`: (optional) Maximum number of documents to return (default: 10, max: 100).
+- `offset`: (optional) Starting point for the document list (default: 0).
+- `max_text_length`: (optional) Maximum text length to return per document (default: 1000).
+- `metadata_filter`: (optional) A JSON string representing key-value pairs to filter documents by their metadata.
+
+### List Documents Response
+
+```json
+{
+  "documents": [
+    {
+      "doc_id": "123456",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
+      "hash_value": "text_hash_value",
+      "metadata": {
+        "author": "kaito",
+      },
+      "is_truncated": false
+    }
+  ],
+  "count": 1
+}
+```
+
+Each document in the response includes its unique `doc_id`, the (possibly truncated) text, metadata, and a flag indicating if the text was truncated.
+
+**Note:**  
+If you want to filter documents by metadata, provide the `metadata_filter` parameter as a JSON string. For example:
+
+```
+GET /indexes/rag_index/documents?metadata_filter={"author":"kaito"}
+```
+
+
+## Updating Documents
+
+To update existing documents in a specific index, use the `/indexes/{index_name}/documents` API route. This endpoint accepts a POST request with the index name in the URL and a list of documents to update in the request body.
+
+### Update Documents Request
+
+```json
+POST /indexes/rag_index/documents
+{
+  "documents": [
+    {
+      "doc_id": "123456",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data. Adding an information retrieval system gives you control over grounding data used by an LLM when it formulates a response.",
+      "hash_value": "text_hash_value",
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ]
+}
+```
+
+- `doc_id`: The unique identifier of the document to update.
+- `text`: The new or updated text for the document.
+- `metadata`: (Optional) Updated metadata for the document.
+
+### Update Documents Response
+
+```json
+{
+  "updated_documents": [
+    {
+      "doc_id": "123456",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data. Adding an information retrieval system gives you control over grounding data used by an LLM when it formulates a response.",
+      "hash_value": "text_hash_value",
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ],
+  "unchanged_documents": [],
+  "not_found_documents": []
+}
+```
+
+- `updated_documents`: Documents that were successfully updated.
+- `unchanged_documents`: Documents that were provided but did not require changes.
+- `not_found_documents`: Documents with IDs that were not found in the index.
+
+Use this endpoint to keep your indexed documents up to date with the latest content or metadata.
+
+## Delete Documents
+
+
+To delete one or more documents from a specific index, use the `/indexes/{index_name}/documents/delete` API route. This endpoint accepts a POST request with the index name in the URL and a list of document IDs to delete in the request body.
+
+### Delete Documents Request
+
+```json
+POST /indexes/rag_index/documents/delete
+{
+  "doc_ids": ["123456"]
+}
+```
+
+- `doc_ids`: A list of document IDs to delete from the specified index.
+
+### Delete Documents Response
+
+```json
+{
+  "deleted_doc_ids": ["123456"],
+  "not_found_doc_ids": []
+}
+```
+
+- `deleted_doc_ids`: Document IDs that were successfully deleted.
+- `not_found_doc_ids`: Document IDs that were not found in the index.
+
+Use this endpoint to remove documents that are no longer needed from your index.
+
+## Persist Index
+
+To save (persist) the data of an index to disk, use the `/persist/{index_name}` API route. This endpoint accepts a POST request with the index name in the URL and an optional `path` query parameter specifying where to save the index data.
+
+### Persist Index Request
+
+```
+POST /persist/rag_index?path=./custom_path
+```
+
+- `index_name`: The name of the index to persist.
+- `path`: (optional) The directory path where the index will be saved. If not provided, the default directory is used.
+
+### Persist Index Response
+
+```json
+{
+  "message": "Successfully persisted index rag_index to ./custom_path/rag_index."
+}
+```
+
+Use this endpoint to ensure your indexed data is safely stored on disk.
+
+---
+
+## Load Index
+
+To load an existing index from disk, use the `/load/{index_name}` API route. This endpoint accepts a POST request with the index name in the URL, an optional `path` query parameter specifying where to load the index from, and an optional `overwrite` flag.
+
+### Load Index Request
+
+```
+POST /load/rag_index?path=./custom_path/rag_index
+```
+
+- `index_name`: The name of the index to load.
+- `path`: (optional) The path to load the index from. If not provided, the default directory is used.
+- `overwrite`: (optional, default: false) If true, will overwrite the existing index if it already exists in memory.
+
+### Load Index Response
+
+```json
+{
+  "message": "Successfully loaded index rag_index from ./custom_path/rag_index."
+}
+```
+
+Use this endpoint to restore previously persisted indexes into memory for querying and updates.
+
+## Delete Index
+
+To delete an entire index and all of its documents, use the `/indexes/{index_name}` API route. This endpoint accepts a DELETE request with the index name in the URL. Deleting an index is irreversible and will remove all associated documents from memory.
+
+### Delete Index Request
+
+```
+DELETE /indexes/rag_index
+```
+
+- `index_name`: The name of the index to delete.
+
+### Delete Index Response
+
+```json
+{
+  "message": "Successfully deleted index rag_index."
+}
+```
+
+Use this endpoint to permanently remove an index and all its data when it is no longer needed.
+
+## Query Index
+
+To query a specific index for relevant documents, use the `/query` API route. This endpoint accepts a POST request with the index name, query string, and optional parameters for result count, and LLM generation.
+
+### Query Index Request
+
+```json
+POST /query
+{
+  "index_name": "rag_index",
+  "query": "What is RAG?",
+  "top_k": 5,
+  "llm_params": {
+    "temperature": 0.7,
+    "max_tokens": 2048
+  }
+}
+```
+
+- `index_name`: The name of the index to query.
+- `query`: The query string.
+- `top_k`: (optional) Number of top documents to retrieve (default: 5).
+- `llm_params`: (optional) Parameters for LLM-based generation (e.g., temperature, max_tokens).
+
+### Query Index Response
+
+```json
+{
+  "response": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
+  "source_nodes": [
+    {
+      "doc_id": "123456",
+      "node_id": "2853a565-8c1f-4982-acaa-a0ab52691435",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
+      "score": 0.95,
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ],
+  "metadata": {
+    "2853a565-8c1f-4982-acaa-a0ab52691435": {
+      "author": "kaito",
+    }
+  }
+}
+```
+
+- `response`: The generated answer or summary from the LLM (if enabled).
+- `source_nodes`: List of source nodes with their text, score, and metadata.
+- `metadata`: Additional metadata about the query or response.
+
+Use this endpoint to retrieve relevant information from your indexed documents and optionally generate answers using an LLM.
+
+## OpenAI-Compatible Chat Completions
+
+The RAGEngine provides an OpenAI-compatible chat completions endpoint at `/v1/chat/completions`. This endpoint allows you to use RAG capabilities with the familiar OpenAI API format, making it easy to integrate with existing applications that use OpenAI's chat completions.
+
+### RAG Bypass Conditions
+
+The chat completions endpoint automatically determines whether to use RAG or pass requests directly to the LLM based on the request content. Requests will **bypass RAG** and be sent directly to the LLM in the following cases:
+
+1. **No Index Specified**: When `index_name` is not provided in the request, the system treats it as a standard LLM request without retrieval.
+
+2. **Tool/Function Calls**: When the request contains `tools` or `functions` parameters, indicating the use of function calling capabilities.
+
+3. **Unsupported Message Roles**: When messages contain roles other than `user`, `system`, `assistant`, or `developer`. 
+
+4. **Non-Text Content**: When user messages contain non-text content types (such as images, audio, or other media formats).
+
+When any of these conditions are met, the request is passed through directly to the underlying LLM without document retrieval or augmentation, ensuring full compatibility with standard OpenAI API usage patterns.
+
+### Chat Completions Request
+
+```json
+POST /v1/chat/completions
+{
+  "index_name": "rag_index",
+  "model": "example_model",
+  "messages": [
+    {"role": "system", "content": "You are a knowledgeable assistant."},
+    {"role": "user", "content": "What is RAG?"}
+  ],
+  "temperature": 0.7,
+  "max_tokens": 2048,
+  "top_k": 5
+}
+```
+
+- `index_name`: (optional) The name of the index to query for relevant documents. If not included, the request will be sent to the LLM directly with no additional context.
+- `model`: The model identifier (for compatibility).
+- `messages`: Array of message objects with `role` and `content` fields.
+- `temperature`: (optional) Controls randomness in the response (0.0 to 1.0).
+- `max_tokens`: (optional) Maximum number of tokens to generate.
+- `top_k`: (optional) Number of top documents to retrieve from the index.
+
+### Chat Completions Response
+
+```json
+{
+  "id": "chatcmpl-123",
+  "object": "chat.completion",
+  "created": 1677652288,
+  "model": "example_model",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": "RAG stands for Retrieval-Augmented Generation. It's an architecture that augments the capabilities of a Large Language Model by adding an information retrieval system that provides grounding data."
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "source_nodes": [
+    {
+      "doc_id": "123456",
+      "node_id": "2853a565-8c1f-4982-acaa-a0ab52691435",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
+      "score": 0.95,
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ]
+}
+```
+
+- `id`: Unique identifier for the chat completion.
+- `object`: Type of object returned (`chat.completion`).
+- `created`: Unix timestamp of when the completion was created.
+- `model`: The model used for the completion.
+- `choices`: Array of completion choices (typically one).
+- `source_nodes`: (RAG-specific) List of source documents that informed the response.
+
+Use this endpoint to integrate RAG capabilities into applications that already use OpenAI's chat completions API format.
+
+## Example Python Client
+
+You can leverage the [example_rag_client.py](./example_rag_client.py) as a starting point for a rag client with inputs that match the route documentation above.
diff --git a/website/docs/rag.md b/website/docs/rag.md
index 97fa070..3671746 100644
--- a/website/docs/rag.md
+++ b/website/docs/rag.md
@@ -1,4 +1,6 @@
-# Retrieval-Augmented Generation (RAG)
+---
+title: Retrieval-Augmented Generation (RAG)
+---
 
 This document presents how to use the KAITO `ragengine` Custom Resource Definition (CRD) for retrieval-augumented generatoin workflow. By creating a RAGEngine resource, you can quickly stand up a service that indexes documents and queries them in conjunction with an existing LLM inference endpoint—no need to custom-build pipelines. This enables your large language model to answer questions based on your own private content.
 
@@ -78,351 +80,4 @@ spec:
 After you create your YAML configuration, run:
 ```sh
 kubectl apply -f examples/RAG/kaito_ragengine_phi_3.yaml
-```
-
-## API definitions and examples
-
-A **RAGEngine index** is a logical collection that organizes and stores your documents for retrieval-augmented generation workflows. The relationship between indexes, documents, and document nodes is as follows:
-
-- **Index**: An index is a named container that holds a set of documents. Each index is independent and can be created, updated, queried, persisted, loaded, or deleted via the API.
-
-- **Documents**: Documents are the primary units of content that you add to an index. Each document contains a `text` field (the content to be indexed) and optional `metadata` (such as author, source, or custom tags). When you add documents to an index, each is assigned a unique `doc_id`.
-
-- **Document Nodes**: When a document is ingested, it is automatically split into smaller chunks called *nodes*. The splitting strategy depends on the document type and metadata:
-  - By default, documents are split into sentences.
-  - If you specify code-aware splitting (using the `split_type` and `language` metadata), the document is split into code blocks or logical code units.
-  - Each node represents a chunk of text that is indexed and can be retrieved as part of a query.
-
-**How it works in practice:**
-- When you index a document, it is divided into nodes for efficient retrieval and semantic search.
-- When you query an index, the engine retrieves the most relevant nodes (not necessarily whole documents) and can use them to generate answers or summaries.
-- The `source_nodes` field in query responses contains the actual nodes that matched your query, along with their scores and metadata.
-
-This design enables fine-grained retrieval and more accurate, context-aware responses from your LLM-powered applications.
-
-### Creating an Index With Documents
-
-To add documents to an index or create a new index, use the `/index` API route. This endpoint accepts a POST request with the index name and a list of documents to be indexed.
-
-**Request Example:**
-
-```json
-POST /index
-{
-  "index_name": "rag_index",
-  "documents": [
-    {
-      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
-      "metadata": {
-        "author": "Microsoft",
-        "source": "https://learn.microsoft.com/en-us/azure/search/retrieval-augmented-generation-overview?tabs=docs"
-      }
-    }
-  ]
-}
-```
-
-- `index_name`: The name of the index to create or update.
-- `documents`: A list of documents, each with a `text` field and optional `metadata`.
-
-**Response Example:**
-
-```json
-[
-  {
-    "doc_id": "123456",
-    "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
-    "hash_value": "text_hash_value",
-    "metadata": {
-      "author": "Microsoft",
-      "source": "https://learn.microsoft.com/en-us/azure/search/retrieval-augmented-generation-overview?tabs=docs"
-    },
-    "is_truncated": false
-  }
-]
-```
-
-Each returned document includes its unique `doc_id`, the original text, metadata, and a flag indicating if the text was truncated. The `doc_id` will be important for document update/delete calls. 
-
-
-#### Splitting Documents with CodeSplitter
-
-By default, RAGEngine splits documents into sentences. However, you can instruct the engine to split documents using the `CodeSplitter` (for code-aware chunking) by providing metadata in your API request.
-
-To use the `CodeSplitter`, set the `split_type` to `"code"` and specify the programming language in the `language` field of the document metadata. For example, when calling the RAGEngine API to index documents:
-
-```json
-{
-  "documents": [
-    {
-      "text": "def foo():\n    return 42\n\n# Another function\ndef bar():\n    pass",
-      "metadata": {
-        "split_type": "code",
-        "language": "python"
-      }
-    }
-  ]
-}
-```
-
-This instructs the RAGEngine to use code-aware splitting for the provided document. If `split_type` is not set or set to any other value, sentence splitting will be used by default.
-
-### List Documents
-
-To retrieve a paginated list of documents from a specific index, use the `/indexes/{index_name}/documents` API route. This endpoint accepts a GET request with optional query parameters for pagination, text truncation, and metadata filtering.
-
-**Request Example:**
-
-```
-GET /indexes/rag_index/documents?limit=5&offset=0&max_text_length=500
-```
-
-- `limit`: (optional) Maximum number of documents to return (default: 10, max: 100).
-- `offset`: (optional) Starting point for the document list (default: 0).
-- `max_text_length`: (optional) Maximum text length to return per document (default: 1000).
-- `metadata_filter`: (optional) A JSON string representing key-value pairs to filter documents by their metadata.
-
-**Response Example:**
-
-```json
-{
-  "documents": [
-    {
-      "doc_id": "123456",
-      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
-      "hash_value": "text_hash_value",
-      "metadata": {
-        "author": "Microsoft",
-        "source": "https://learn.microsoft.com/en-us/azure/search/retrieval-augmented-generation-overview?tabs=docs"
-      },
-      "is_truncated": false
-    }
-  ],
-  "count": 1
-}
-```
-
-Each document in the response includes its unique `doc_id`, the (possibly truncated) text, metadata, and a flag indicating if the text was truncated.
-
-**Note:**  
-If you want to filter documents by metadata, provide the `metadata_filter` parameter as a JSON string. For example:
-
-```
-GET /indexes/rag_index/documents?metadata_filter={"author":"Microsoft"}
-```
-
-
-### Updating Documents
-
-To update existing documents in a specific index, use the `/indexes/{index_name}/documents` API route. This endpoint accepts a POST request with the index name in the URL and a list of documents to update in the request body.
-
-**Request Example:**
-
-```json
-POST /indexes/rag_index/documents
-{
-  "documents": [
-    {
-      "doc_id": "123456",
-      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data. Adding an information retrieval system gives you control over grounding data used by an LLM when it formulates a response.",
-      "hash_value": "text_hash_value",
-      "metadata": {
-        "author": "Microsoft",
-        "source": "https://learn.microsoft.com/en-us/azure/search/retrieval-augmented-generation-overview?tabs=docs"
-      },
-    }
-  ]
-}
-```
-
-- `doc_id`: The unique identifier of the document to update.
-- `text`: The new or updated text for the document.
-- `metadata`: (Optional) Updated metadata for the document.
-
-**Response Example:**
-
-```json
-{
-  "updated_documents": [
-    {
-      "doc_id": "123456",
-      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data. Adding an information retrieval system gives you control over grounding data used by an LLM when it formulates a response.",
-      "hash_value": "text_hash_value",
-      "metadata": {
-        "author": "Microsoft",
-        "source": "https://learn.microsoft.com/en-us/azure/search/retrieval-augmented-generation-overview?tabs=docs"
-      },
-    }
-  ],
-  "unchanged_documents": [],
-  "not_found_documents": []
-}
-```
-
-- `updated_documents`: Documents that were successfully updated.
-- `unchanged_documents`: Documents that were provided but did not require changes.
-- `not_found_documents`: Documents with IDs that were not found in the index.
-
-Use this endpoint to keep your indexed documents up to date with the latest content or metadata.
-
-### Delete Documents
-
-
-To delete one or more documents from a specific index, use the `/indexes/{index_name}/documents/delete` API route. This endpoint accepts a POST request with the index name in the URL and a list of document IDs to delete in the request body.
-
-**Request Example:**
-
-```json
-POST /indexes/rag_index/documents/delete
-{
-  "doc_ids": ["123456"]
-}
-```
-
-- `doc_ids`: A list of document IDs to delete from the specified index.
-
-**Response Example:**
-
-```json
-{
-  "deleted_doc_ids": ["123456"],
-  "not_found_doc_ids": []
-}
-```
-
-- `deleted_doc_ids`: Document IDs that were successfully deleted.
-- `not_found_doc_ids`: Document IDs that were not found in the index.
-
-Use this endpoint to remove documents that are no longer needed from your index.
-
-### Persist Index
-
-To save (persist) the data of an index to disk, use the `/persist/{index_name}` API route. This endpoint accepts a POST request with the index name in the URL and an optional `path` query parameter specifying where to save the index data.
-
-**Request Example:**
-
-```
-POST /persist/rag_index?path=./custom_path
-```
-
-- `index_name`: The name of the index to persist.
-- `path`: (optional) The directory path where the index will be saved. If not provided, the default directory is used.
-
-**Response Example:**
-
-```json
-{
-  "message": "Successfully persisted index rag_index to ./custom_path/rag_index."
-}
-```
-
-Use this endpoint to ensure your indexed data is safely stored on disk.
-
----
-
-### Load Index
-
-To load an existing index from disk, use the `/load/{index_name}` API route. This endpoint accepts a POST request with the index name in the URL, an optional `path` query parameter specifying where to load the index from, and an optional `overwrite` flag.
-
-**Request Example:**
-
-```
-POST /load/rag_index?path=./custom_path/rag_index
-```
-
-- `index_name`: The name of the index to load.
-- `path`: (optional) The path to load the index from. If not provided, the default directory is used.
-- `overwrite`: (optional, default: false) If true, will overwrite the existing index if it already exists in memory.
-
-**Response Example:**
-
-```json
-{
-  "message": "Successfully loaded index rag_index from ./custom_path/rag_index."
-}
-```
-
-Use this endpoint to restore previously persisted indexes into memory for querying and updates.
-
-### Delete Index
-
-To delete an entire index and all of its documents, use the `/indexes/{index_name}` API route. This endpoint accepts a DELETE request with the index name in the URL. Deleting an index is irreversible and will remove all associated documents from memory.
-
-**Request Example:**
-
-```
-DELETE /indexes/rag_index
-```
-
-- `index_name`: The name of the index to delete.
-
-**Response Example:**
-
-```json
-{
-  "message": "Successfully deleted index rag_index."
-}
-```
-
-Use this endpoint to permanently remove an index and all its data when it is no longer needed.
-
-### Query Index
-
-To query a specific index for relevant documents, use the `/query` API route. This endpoint accepts a POST request with the index name, query string, and optional parameters for result count, and LLM generation.
-
-**Request Example:**
-
-```json
-POST /query
-{
-  "index_name": "rag_index",
-  "query": "What is RAG?",
-  "top_k": 5,
-  "llm_params": {
-    "temperature": 0.7,
-    "max_tokens": 2048
-  }
-}
-```
-
-- `index_name`: The name of the index to query.
-- `query`: The query string.
-- `top_k`: (optional) Number of top documents to retrieve (default: 5).
-- `llm_params`: (optional) Parameters for LLM-based generation (e.g., temperature, max_tokens).
-
-**Response Example:**
-
-```json
-{
-  "response": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
-  "source_nodes": [
-    {
-      "doc_id": "123456",
-      "node_id": "2853a565-8c1f-4982-acaa-a0ab52691435",
-      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
-      "score": 0.95,
-      "metadata": {
-        "author": "Microsoft",
-        "source": "https://learn.microsoft.com/en-us/azure/search/retrieval-augmented-generation-overview?tabs=docs"
-      },
-    }
-  ],
-  "metadata": {
-    "2853a565-8c1f-4982-acaa-a0ab52691435": {
-      "author": "Microsoft",
-      "source": "https://learn.microsoft.com/en-us/azure/search/retrieval-augmented-generation-overview?tabs=docs"
-    }
-  }
-}
-```
-
-- `response`: The generated answer or summary from the LLM (if enabled).
-- `source_nodes`: List of source nodes with their text, score, and metadata.
-- `metadata`: Additional metadata about the query or response.
-
-Use this endpoint to retrieve relevant information from your indexed documents and optionally generate answers using an LLM.
-
-
-## Example Client
-
-You can leverage the [example_rag_client.py](./example_rag_client.py) as a starting point for a rag client with inputs that match the route documentation above.
+```
\ No newline at end of file
diff --git a/website/docs/tuning.md b/website/docs/tuning.md
index 9364646..121c119 100644
--- a/website/docs/tuning.md
+++ b/website/docs/tuning.md
@@ -130,7 +130,7 @@ Figure 1. KAITO tuning pod structure.
 All three containers use shared local volumes (by mounting the same `EmptyDir` volumes), hence file copies between containers are avoided.
 
 ## Pod structure when input and output are specified as Kubernetes volumes
-When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and ouput destination are already available in the specified volumes, the main container can directly access them.
+When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and output destination are already available in the specified volumes, the main container can directly access them.
 
 Other than the absence of the init and sidecar containers, the main container is the same as described in the previous section.
 
diff --git a/website/sidebars.js b/website/sidebars.js
index e798dbe..e87261e 100644
--- a/website/sidebars.js
+++ b/website/sidebars.js
@@ -44,7 +44,17 @@ const sidebars = {
                 'inference',
                 'multi-node-inference',
                 'tuning',
-                'rag',
+                {
+                    type: 'category',
+                    label: 'Retrieval-Augmented Generation (RAG)',
+                    link: {
+                        type: 'doc',
+                        id: 'rag',
+                    },
+                    items: [
+                        'rag-api',
+                    ],
+                },
                 'custom-model',
                 'tool-calling',
                 'model-as-oci-artifacts',
diff --git a/website/versioned_docs/version-v0.5.0/aws.md b/website/versioned_docs/version-v0.5.0/aws.md
index 14d56a5..216392d 100644
--- a/website/versioned_docs/version-v0.5.0/aws.md
+++ b/website/versioned_docs/version-v0.5.0/aws.md
@@ -85,7 +85,7 @@ inference:
 $ kubectl apply -f $kaito_workspace_aws
 ```
 
-The workspace status can be tracked by running the following command. When the WORKSPACEREADY column becomes `True`, the model has been deployed successfully.
+The workspace status can be tracked by running the following command. When the `WORKSPACESUCCEEDED` column becomes `True`, the model has been deployed successfully.
 
 ```sh
 $ kubectl get workspace workspace-falcon-7b
diff --git a/website/versioned_docs/version-v0.5.0/contributing.md b/website/versioned_docs/version-v0.5.0/contributing.md
index 2766aae..3aec963 100644
--- a/website/versioned_docs/version-v0.5.0/contributing.md
+++ b/website/versioned_docs/version-v0.5.0/contributing.md
@@ -40,7 +40,7 @@ EOF
 ```
 
 - `default_registry`: The container registry where Tilt pushes development images (e.g. `<acr_name>.azurecr.io`). Ensure you're logged in with `docker login` and have push/pull permissions.
-if you want to use Azure Container Registry (ACR), set up a new ACR instance and attach it to your AKS cluster with these commands:
+If you want to use Azure Container Registry (ACR), set up a new ACR instance and attach it to your AKS cluster with these commands:
 
 ```bash
 export ACR_NAME=<acr_name>
diff --git a/website/versioned_docs/version-v0.5.0/example_rag_client.py b/website/versioned_docs/version-v0.5.0/example_rag_client.py
index 60d5e83..e2a5760 100644
--- a/website/versioned_docs/version-v0.5.0/example_rag_client.py
+++ b/website/versioned_docs/version-v0.5.0/example_rag_client.py
@@ -11,9 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import requests
 import json
 
+import requests
+
+
 class KAITORAGClient:
     """
     A bare bones client for interacting with the RAGEngine API.
@@ -25,7 +27,7 @@ class KAITORAGClient:
         self.base_url = base_url
         self.headers = {
             "Content-Type": "application/json",
-            "Accept": "application/json"
+            "Accept": "application/json",
         }
 
     def index_documents(self, index_name, documents):
@@ -34,10 +36,7 @@ class KAITORAGClient:
         documents: list of dicts, each with 'text' and optional 'metadata'
         """
         url = f"{self.base_url}/index"
-        payload = {
-            "index_name": index_name,
-            "documents": documents
-        }
+        payload = {"index_name": index_name, "documents": documents}
         resp = requests.post(url, json=payload, headers=self.headers)
         resp.raise_for_status()
         return resp.json()
@@ -56,8 +55,8 @@ class KAITORAGClient:
             "top_k": top_k,
             "llm_params": {
                 "temperature": llm_temperature,
-                "max_tokens": llm_max_tokens
-            }
+                "max_tokens": llm_max_tokens,
+            },
         }
         resp = requests.post(url, json=payload, headers=self.headers)
         resp.raise_for_status()
@@ -89,7 +88,7 @@ class KAITORAGClient:
         List documents in the RAGEngine.
         """
         url = f"{self.base_url}/indexes/{index_name}/documents"
-        params = {"limit": limit, "offset": offset }
+        params = {"limit": limit, "offset": offset}
         if metadata_filter:
             params["metadata_filter"] = json.dumps(metadata_filter)
         resp = requests.get(url, headers=self.headers, params=params)
@@ -132,4 +131,4 @@ class KAITORAGClient:
         url = f"{self.base_url}/indexes/{index_name}"
         resp = requests.delete(url, headers=self.headers)
         resp.raise_for_status()
-        return resp.json()
\ No newline at end of file
+        return resp.json()
diff --git a/website/versioned_docs/version-v0.5.0/installation.md b/website/versioned_docs/version-v0.5.0/installation.md
index c96e135..eced392 100644
--- a/website/versioned_docs/version-v0.5.0/installation.md
+++ b/website/versioned_docs/version-v0.5.0/installation.md
@@ -2,7 +2,7 @@
 title: Installation
 ---
 
-The following guidance assumes **Azure Kubernetes Service(AKS)** is used to host the Kubernetes cluster. If you want to use Elastic Kubernetes Service (EKS) instead, please follow the installation guide [here](https://github.com/kaito-project/kaito/blob/main/docs/aws/aws_installation.md).
+The following guidance assumes **Azure Kubernetes Service (AKS)** is used to host the Kubernetes cluster. If you want to use Elastic Kubernetes Service (EKS) instead, please follow the installation guide [here](./aws.md).
 
 Before you begin, ensure you have the following tools installed:
 
diff --git a/website/versioned_docs/version-v0.5.0/intro.md b/website/versioned_docs/version-v0.5.0/intro.md
index b12321a..7c92b0b 100644
--- a/website/versioned_docs/version-v0.5.0/intro.md
+++ b/website/versioned_docs/version-v0.5.0/intro.md
@@ -5,7 +5,7 @@ slug: /
 
 :::info What's NEW!
 
-Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about from [here](https://kaito-project.github.io/kaito/docs/rag)!
+Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about it [here](https://kaito-project.github.io/kaito/docs/rag)!
 **Latest Release:** July 2nd, 2025. KAITO v0.5.0.
 
 **First Release:** Nov 15th, 2023. KAITO v0.1.0.
@@ -46,7 +46,7 @@ The [*gpu-provisioner*](https://github.com/Azure/gpu-provisioner) is an open sou
 
 
 As illustrated in the above figure, the **RAGEngine controller** reconciles the `ragengine` custom resource and creates a `RAGService` deployment. The `RAGService` provides the following capabilities:
-  - **Orchstration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
+  - **Orchestration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
   - **Embedding**: support both local and remote embedding services, to embed queries and documents in the vector database.
   - **Vector database**: support a built-in [faiss](https://github.com/facebookresearch/faiss) in-memory vector database. Remote vector database support will be added soon.
   - **Backend inference**: support any OAI compatible inference service.
diff --git a/website/versioned_docs/version-v0.5.0/kaito-on-byo-gpu-nodes.md b/website/versioned_docs/version-v0.5.0/kaito-on-byo-gpu-nodes.md
index ef3c664..054871a 100644
--- a/website/versioned_docs/version-v0.5.0/kaito-on-byo-gpu-nodes.md
+++ b/website/versioned_docs/version-v0.5.0/kaito-on-byo-gpu-nodes.md
@@ -210,7 +210,7 @@ EOF
 ```
 
 :::note
-In the above configuration you can see we have use a node labelSelector value as `apps: gpu`, this is the same label we have applied when we added the GPU node pool earlier.
+In the above configuration you can see we have used a node labelSelector value as `apps: gpu`, this is the same label we have applied when we added the GPU node pool earlier.
 :::
 
 Ensure that the workspace is ready by running the following command:
@@ -240,7 +240,7 @@ export CLUSTERIP=$(kubectl get \
 Let's send a request to the workspace to get an inference response. Modify the prompt as you see fit:
 
 ```bash
-export QUESTION="What's are LLMs?"
+export QUESTION="What are LLMs?"
 ```
 
 Run the following command to send the request:
diff --git a/website/versioned_docs/version-v0.5.0/monitoring.md b/website/versioned_docs/version-v0.5.0/monitoring.md
index 0f74ef6..2fc88ca 100644
--- a/website/versioned_docs/version-v0.5.0/monitoring.md
+++ b/website/versioned_docs/version-v0.5.0/monitoring.md
@@ -23,7 +23,7 @@ vLLM V1 engine is now enabled by default. Check [here](https://docs.vllm.ai/en/l
 | System Stats - Cache | `vllm:gpu_cache_usage_perc` | Gauge | GPU KV-cache usage. 1 means 100 percent usage |
 | System Stats - Cache | `vllm:cpu_cache_usage_perc` | Gauge | CPU KV-cache usage. 1 means 100 percent usage (DEPRECATED: KV cache offloading is not used in V1) |
 | System Stats - Cache | `vllm:cpu_prefix_cache_hit_rate` | Gauge | CPU prefix cache block hit rate (DEPRECATED: KV cache offloading is not used in V1) |
-| System Stats - Cache | `vllm:gpu_prefix_cache_hit_rate` | Gauge | GPU prefix cache block hit rat (DEPRECATED: use `vllm:gpu_prefix_cache_queries` and `vllm:gpu_prefix_cache_hits in V1`) |
+| System Stats - Cache | `vllm:gpu_prefix_cache_hit_rate` | Gauge | GPU prefix cache block hit rate (DEPRECATED: use `vllm:gpu_prefix_cache_queries` and `vllm:gpu_prefix_cache_hits in V1`) |
 | System Stats - Cache | `vllm:gpu_cache_usage_perc` | Counter | GPU prefix cache queries, in terms of number of queried blocks (V1 only) |
 | System Stats - Cache | `vllm:gpu_prefix_cache_hits` | Counter | GPU prefix cache hits, in terms of number of cached blocks (V1 only) |
 | Iteration Stats | `vllm:num_preemptions_total` | Counter | Cumulative number of preemption from the engine |
diff --git a/website/versioned_docs/version-v0.5.0/presets.md b/website/versioned_docs/version-v0.5.0/presets.md
index 5688c7f..97c39dd 100644
--- a/website/versioned_docs/version-v0.5.0/presets.md
+++ b/website/versioned_docs/version-v0.5.0/presets.md
@@ -21,7 +21,7 @@ Each preset model has its own hardware requirements in terms of GPU count and GP
 
 ## Distributed inference
 
-For models that support distributed inference, when the node count is larger than one, [torch distributed elastic](https://pytorch.org/docs/stable/distributed.elastic.html) is configured with master/worker pods running in multiple nodes and the service endpoint is the master pod.
+For models that support distributed inference, when the node count is larger than one, [Torch Distributed Elastic](https://pytorch.org/docs/stable/distributed.elastic.html) is configured with master/worker pods running in multiple nodes and the service endpoint is the master pod.
 
 The following preset models support multi-node distributed inference:
 
diff --git a/website/versioned_docs/version-v0.5.0/quick-start.md b/website/versioned_docs/version-v0.5.0/quick-start.md
index db63531..6159ae0 100644
--- a/website/versioned_docs/version-v0.5.0/quick-start.md
+++ b/website/versioned_docs/version-v0.5.0/quick-start.md
@@ -42,7 +42,7 @@ Track the workspace status to see when the model has been deployed successfully:
 kubectl get workspace workspace-phi-3-5-mini
 ```
 
-When the `WORKSPACEREADY` column becomes `True`, the model has been deployed successfully:
+When the `WORKSPACESUCCEEDED` column becomes `True`, the model has been deployed successfully:
 
 ```bash
 NAME                     INSTANCE                   RESOURCEREADY   INFERENCEREADY   JOBSTARTED   WORKSPACESUCCEEDED   AGE
diff --git a/website/versioned_docs/version-v0.5.0/tuning.md b/website/versioned_docs/version-v0.5.0/tuning.md
index 0b455e0..2f75fde 100644
--- a/website/versioned_docs/version-v0.5.0/tuning.md
+++ b/website/versioned_docs/version-v0.5.0/tuning.md
@@ -49,7 +49,7 @@ KAITO provides default tuning configurations for different tuning methods. They
 - [default QLoRA configmap](../../../charts/kaito/workspace/templates/qlora-params.yaml)
 
 ## Tuning configmaps
-User can specify a customized configmap via the `Config` field of the `TuningSpec`. The customized configmap should be structured based on the default configmaps provided by KAITO. Please read the following section carefully when attempting to change the default parameters used by KAITO.
+Users can specify a customized configmap via the `Config` field of the `TuningSpec`. The customized configmap should be structured based on the default configmaps provided by KAITO. Please read the following section carefully when attempting to change the default parameters used by KAITO.
 
 ### Categorized key parameters
 Note that changing these parameters may largely impact the tuning result. In addition, users can add extra parameters that are not presented in the default configmaps. For a complete list of supported parameters, please refer to the provided huggingface documentation.
@@ -130,7 +130,7 @@ Figure 1. KAITO tuning pod structure.
 All three containers use shared local volumes (by mounting the same `EmptyDir` volumes), hence file copies between containers are avoided.
 
 ## Pod structure when input and output are specified as Kubernetes volumes
-When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and ouput destination are already available in the specified volumes, the main container can directly access them.
+When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and output destination are already available in the specified volumes, the main container can directly access them.
 
 Other than the absence of the init and sidecar containers, the main container is the same as described in the previous section.
 
diff --git a/website/versioned_docs/version-v0.5.1/aws.md b/website/versioned_docs/version-v0.5.1/aws.md
index 14d56a5..216392d 100644
--- a/website/versioned_docs/version-v0.5.1/aws.md
+++ b/website/versioned_docs/version-v0.5.1/aws.md
@@ -85,7 +85,7 @@ inference:
 $ kubectl apply -f $kaito_workspace_aws
 ```
 
-The workspace status can be tracked by running the following command. When the WORKSPACEREADY column becomes `True`, the model has been deployed successfully.
+The workspace status can be tracked by running the following command. When the `WORKSPACESUCCEEDED` column becomes `True`, the model has been deployed successfully.
 
 ```sh
 $ kubectl get workspace workspace-falcon-7b
diff --git a/website/versioned_docs/version-v0.5.1/contributing.md b/website/versioned_docs/version-v0.5.1/contributing.md
index 2766aae..3aec963 100644
--- a/website/versioned_docs/version-v0.5.1/contributing.md
+++ b/website/versioned_docs/version-v0.5.1/contributing.md
@@ -40,7 +40,7 @@ EOF
 ```
 
 - `default_registry`: The container registry where Tilt pushes development images (e.g. `<acr_name>.azurecr.io`). Ensure you're logged in with `docker login` and have push/pull permissions.
-if you want to use Azure Container Registry (ACR), set up a new ACR instance and attach it to your AKS cluster with these commands:
+If you want to use Azure Container Registry (ACR), set up a new ACR instance and attach it to your AKS cluster with these commands:
 
 ```bash
 export ACR_NAME=<acr_name>
diff --git a/website/versioned_docs/version-v0.5.1/example_rag_client.py b/website/versioned_docs/version-v0.5.1/example_rag_client.py
index 60d5e83..e2a5760 100644
--- a/website/versioned_docs/version-v0.5.1/example_rag_client.py
+++ b/website/versioned_docs/version-v0.5.1/example_rag_client.py
@@ -11,9 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import requests
 import json
 
+import requests
+
+
 class KAITORAGClient:
     """
     A bare bones client for interacting with the RAGEngine API.
@@ -25,7 +27,7 @@ class KAITORAGClient:
         self.base_url = base_url
         self.headers = {
             "Content-Type": "application/json",
-            "Accept": "application/json"
+            "Accept": "application/json",
         }
 
     def index_documents(self, index_name, documents):
@@ -34,10 +36,7 @@ class KAITORAGClient:
         documents: list of dicts, each with 'text' and optional 'metadata'
         """
         url = f"{self.base_url}/index"
-        payload = {
-            "index_name": index_name,
-            "documents": documents
-        }
+        payload = {"index_name": index_name, "documents": documents}
         resp = requests.post(url, json=payload, headers=self.headers)
         resp.raise_for_status()
         return resp.json()
@@ -56,8 +55,8 @@ class KAITORAGClient:
             "top_k": top_k,
             "llm_params": {
                 "temperature": llm_temperature,
-                "max_tokens": llm_max_tokens
-            }
+                "max_tokens": llm_max_tokens,
+            },
         }
         resp = requests.post(url, json=payload, headers=self.headers)
         resp.raise_for_status()
@@ -89,7 +88,7 @@ class KAITORAGClient:
         List documents in the RAGEngine.
         """
         url = f"{self.base_url}/indexes/{index_name}/documents"
-        params = {"limit": limit, "offset": offset }
+        params = {"limit": limit, "offset": offset}
         if metadata_filter:
             params["metadata_filter"] = json.dumps(metadata_filter)
         resp = requests.get(url, headers=self.headers, params=params)
@@ -132,4 +131,4 @@ class KAITORAGClient:
         url = f"{self.base_url}/indexes/{index_name}"
         resp = requests.delete(url, headers=self.headers)
         resp.raise_for_status()
-        return resp.json()
\ No newline at end of file
+        return resp.json()
diff --git a/website/versioned_docs/version-v0.5.1/installation.md b/website/versioned_docs/version-v0.5.1/installation.md
index 1a69425..2f8810c 100644
--- a/website/versioned_docs/version-v0.5.1/installation.md
+++ b/website/versioned_docs/version-v0.5.1/installation.md
@@ -2,7 +2,7 @@
 title: Installation
 ---
 
-The following guidance assumes **Azure Kubernetes Service(AKS)** is used to host the Kubernetes cluster. If you want to use Elastic Kubernetes Service (EKS) instead, please follow the installation guide [here](https://github.com/kaito-project/kaito/blob/main/docs/aws/aws_installation.md).
+The following guidance assumes **Azure Kubernetes Service (AKS)** is used to host the Kubernetes cluster. If you want to use Elastic Kubernetes Service (EKS) instead, please follow the installation guide [here](./aws.md).
 
 Before you begin, ensure you have the following tools installed:
 
diff --git a/website/versioned_docs/version-v0.5.1/intro.md b/website/versioned_docs/version-v0.5.1/intro.md
index 9127b38..7b87767 100644
--- a/website/versioned_docs/version-v0.5.1/intro.md
+++ b/website/versioned_docs/version-v0.5.1/intro.md
@@ -5,7 +5,7 @@ slug: /
 
 :::info What's NEW!
 
-Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about from [here](https://kaito-project.github.io/kaito/docs/rag)!
+Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about it [here](https://kaito-project.github.io/kaito/docs/rag)!
 **Latest Release:** July 18th, 2025. KAITO v0.5.1.
 
 **First Release:** Nov 15th, 2023. KAITO v0.1.0.
@@ -46,7 +46,7 @@ The [*gpu-provisioner*](https://github.com/Azure/gpu-provisioner) is an open sou
 
 
 As illustrated in the above figure, the **RAGEngine controller** reconciles the `ragengine` custom resource and creates a `RAGService` deployment. The `RAGService` provides the following capabilities:
-  - **Orchstration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
+  - **Orchestration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
   - **Embedding**: support both local and remote embedding services, to embed queries and documents in the vector database.
   - **Vector database**: support a built-in [faiss](https://github.com/facebookresearch/faiss) in-memory vector database. Remote vector database support will be added soon.
   - **Backend inference**: support any OAI compatible inference service.
diff --git a/website/versioned_docs/version-v0.5.1/kaito-on-byo-gpu-nodes.md b/website/versioned_docs/version-v0.5.1/kaito-on-byo-gpu-nodes.md
index ef3c664..054871a 100644
--- a/website/versioned_docs/version-v0.5.1/kaito-on-byo-gpu-nodes.md
+++ b/website/versioned_docs/version-v0.5.1/kaito-on-byo-gpu-nodes.md
@@ -210,7 +210,7 @@ EOF
 ```
 
 :::note
-In the above configuration you can see we have use a node labelSelector value as `apps: gpu`, this is the same label we have applied when we added the GPU node pool earlier.
+In the above configuration you can see we have used a node labelSelector value as `apps: gpu`, this is the same label we have applied when we added the GPU node pool earlier.
 :::
 
 Ensure that the workspace is ready by running the following command:
@@ -240,7 +240,7 @@ export CLUSTERIP=$(kubectl get \
 Let's send a request to the workspace to get an inference response. Modify the prompt as you see fit:
 
 ```bash
-export QUESTION="What's are LLMs?"
+export QUESTION="What are LLMs?"
 ```
 
 Run the following command to send the request:
diff --git a/website/versioned_docs/version-v0.5.1/monitoring.md b/website/versioned_docs/version-v0.5.1/monitoring.md
index 0f74ef6..2fc88ca 100644
--- a/website/versioned_docs/version-v0.5.1/monitoring.md
+++ b/website/versioned_docs/version-v0.5.1/monitoring.md
@@ -23,7 +23,7 @@ vLLM V1 engine is now enabled by default. Check [here](https://docs.vllm.ai/en/l
 | System Stats - Cache | `vllm:gpu_cache_usage_perc` | Gauge | GPU KV-cache usage. 1 means 100 percent usage |
 | System Stats - Cache | `vllm:cpu_cache_usage_perc` | Gauge | CPU KV-cache usage. 1 means 100 percent usage (DEPRECATED: KV cache offloading is not used in V1) |
 | System Stats - Cache | `vllm:cpu_prefix_cache_hit_rate` | Gauge | CPU prefix cache block hit rate (DEPRECATED: KV cache offloading is not used in V1) |
-| System Stats - Cache | `vllm:gpu_prefix_cache_hit_rate` | Gauge | GPU prefix cache block hit rat (DEPRECATED: use `vllm:gpu_prefix_cache_queries` and `vllm:gpu_prefix_cache_hits in V1`) |
+| System Stats - Cache | `vllm:gpu_prefix_cache_hit_rate` | Gauge | GPU prefix cache block hit rate (DEPRECATED: use `vllm:gpu_prefix_cache_queries` and `vllm:gpu_prefix_cache_hits in V1`) |
 | System Stats - Cache | `vllm:gpu_cache_usage_perc` | Counter | GPU prefix cache queries, in terms of number of queried blocks (V1 only) |
 | System Stats - Cache | `vllm:gpu_prefix_cache_hits` | Counter | GPU prefix cache hits, in terms of number of cached blocks (V1 only) |
 | Iteration Stats | `vllm:num_preemptions_total` | Counter | Cumulative number of preemption from the engine |
diff --git a/website/versioned_docs/version-v0.5.1/presets.md b/website/versioned_docs/version-v0.5.1/presets.md
index 5688c7f..97c39dd 100644
--- a/website/versioned_docs/version-v0.5.1/presets.md
+++ b/website/versioned_docs/version-v0.5.1/presets.md
@@ -21,7 +21,7 @@ Each preset model has its own hardware requirements in terms of GPU count and GP
 
 ## Distributed inference
 
-For models that support distributed inference, when the node count is larger than one, [torch distributed elastic](https://pytorch.org/docs/stable/distributed.elastic.html) is configured with master/worker pods running in multiple nodes and the service endpoint is the master pod.
+For models that support distributed inference, when the node count is larger than one, [Torch Distributed Elastic](https://pytorch.org/docs/stable/distributed.elastic.html) is configured with master/worker pods running in multiple nodes and the service endpoint is the master pod.
 
 The following preset models support multi-node distributed inference:
 
diff --git a/website/versioned_docs/version-v0.5.1/quick-start.md b/website/versioned_docs/version-v0.5.1/quick-start.md
index db63531..6159ae0 100644
--- a/website/versioned_docs/version-v0.5.1/quick-start.md
+++ b/website/versioned_docs/version-v0.5.1/quick-start.md
@@ -42,7 +42,7 @@ Track the workspace status to see when the model has been deployed successfully:
 kubectl get workspace workspace-phi-3-5-mini
 ```
 
-When the `WORKSPACEREADY` column becomes `True`, the model has been deployed successfully:
+When the `WORKSPACESUCCEEDED` column becomes `True`, the model has been deployed successfully:
 
 ```bash
 NAME                     INSTANCE                   RESOURCEREADY   INFERENCEREADY   JOBSTARTED   WORKSPACESUCCEEDED   AGE
diff --git a/website/versioned_docs/version-v0.5.1/tuning.md b/website/versioned_docs/version-v0.5.1/tuning.md
index 0b455e0..2f75fde 100644
--- a/website/versioned_docs/version-v0.5.1/tuning.md
+++ b/website/versioned_docs/version-v0.5.1/tuning.md
@@ -49,7 +49,7 @@ KAITO provides default tuning configurations for different tuning methods. They
 - [default QLoRA configmap](../../../charts/kaito/workspace/templates/qlora-params.yaml)
 
 ## Tuning configmaps
-User can specify a customized configmap via the `Config` field of the `TuningSpec`. The customized configmap should be structured based on the default configmaps provided by KAITO. Please read the following section carefully when attempting to change the default parameters used by KAITO.
+Users can specify a customized configmap via the `Config` field of the `TuningSpec`. The customized configmap should be structured based on the default configmaps provided by KAITO. Please read the following section carefully when attempting to change the default parameters used by KAITO.
 
 ### Categorized key parameters
 Note that changing these parameters may largely impact the tuning result. In addition, users can add extra parameters that are not presented in the default configmaps. For a complete list of supported parameters, please refer to the provided huggingface documentation.
@@ -130,7 +130,7 @@ Figure 1. KAITO tuning pod structure.
 All three containers use shared local volumes (by mounting the same `EmptyDir` volumes), hence file copies between containers are avoided.
 
 ## Pod structure when input and output are specified as Kubernetes volumes
-When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and ouput destination are already available in the specified volumes, the main container can directly access them.
+When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and output destination are already available in the specified volumes, the main container can directly access them.
 
 Other than the absence of the init and sidecar containers, the main container is the same as described in the previous section.
 
diff --git a/website/versioned_docs/version-v0.6.x/aikit.md b/website/versioned_docs/version-v0.6.x/aikit.md
new file mode 100644
index 0000000..826e501
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/aikit.md
@@ -0,0 +1,133 @@
+# AIKit Integration with KAITO
+
+[AIKit](https://github.com/sozercan/aikit/) provides a streamlined way to package and deploy large language models (LLMs) as container images.
+
+This document demonstrates how to integrate AIKit-built models with KAITO workspaces for efficient AI model deployment on Kubernetes, including CPU-based inference and custom model creation with a variety of supported formats, such as GGUF, GPTQ, EXL2, and more.
+
+For more detailed information about AIKit, please refer to the [AIKit documentation](https://sozercan.github.io/aikit/docs/). For any AIKit-related issues, please open an issue in the [AIKit repository](https://github.com/sozercan/aikit/issues).
+
+## Overview
+
+AIKit enables you to:
+
+- 📦 [Package AI models](https://sozercan.github.io/aikit/docs/create-images) as OCI container images with minimal configuration
+- 🤏 Minimal image size, resulting in less vulnerabilities and smaller attack surface with a custom distroless-based image
+- 🏃 Run models with a variety of inference backends, such as text or image generation
+- 🖥️ Supports [AMD64 and ARM64 CPUs](https://sozercan.github.io/aikit/docs/create-images#multi-platform-support) and [GPU-accelerated inferencing with NVIDIA GPUs](https://sozercan.github.io/aikit/docs/gpu)
+- 🪄 Integrate seamlessly with KAITO's infrastructure management and deployment workflows
+
+:::note
+
+While AIKit and KAITO integrate well, they are separate projects. AIKit focuses on model packaging and deployment, while KAITO provides infrastructure management and Kubernetes deployment workflows via controllers. There may be differences in what model formats are supported by each project.
+
+:::
+
+## Deploying AIKit Models to KAITO
+
+### Cluster Setup
+
+This guide will provide instructions using a [kind](https://kind.sigs.k8s.io/) cluster for local development and testing so it's easy to get started.
+
+Please note that if you already have a Kubernetes cluster set up, you can skip the cluster setup section.
+
+- Download and install [kind](https://kind.sigs.k8s.io/docs/user/quick-start/)
+
+- Create a kind cluster:
+
+```bash
+kind create cluster --name kaito
+```
+
+- Install [KAITO workspace controller](installation.md#install-kaito-workspace-controller) on your cluster
+
+### KAITO Workspace Configuration
+
+Create a KAITO workspace configuration file to deploy your model. Here's a complete example:
+
+```yaml title="aikit-workspace.yaml"
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-llama-3point2-3b
+resource:
+  labelSelector:
+    matchLabels:
+      apps: llama-3point2-3b
+  preferredNodes:
+    - kaito-control-plane
+inference:
+  template:
+    spec:
+      containers:
+        - name: llama-3point2-3b
+          image: ghcr.io/sozercan/llama3.2:3b
+          args:
+            - "run"
+            - "--address=:5000"
+```
+
+:::info Memory Requirements
+Before deploying models, check the model's memory requirements to avoid Out of Memory (OOM) errors. Add appropriate `resources.requests.memory` and `resources.limits.memory` to your container spec based on the model requirements.
+
+For GGUF models:
+
+- 7B models generally require at least 8GB of RAM
+- 13B models generally require at least 16GB of RAM
+- 70B models generally require at least 64GB of RAM
+
+You can use [gguf-parser-go](https://github.com/gpustack/gguf-parser-go) to get a better estimate for the memory requirements for a given GGUF model, and quantization.
+:::
+
+Label the nodes with the applicable label to ensure the workspace can schedule pods on them.
+
+```bash
+kubectl label nodes kaito-control-plane apps=llama-3point2-3b
+```
+
+Deploy the workspace using:
+
+```bash
+kubectl apply -f aikit-workspace.yaml
+```
+
+AIKit provides a number of pre-built and curated models that can be used directly. Please refer to [Pre-made Models](https://sozercan.github.io/aikit/docs/premade-models) for available options.
+
+:::tip
+
+Alternatively, if you are on a supported cloud provider and want the cloud provider to auto-provision the nodes for you, you can define an `instanceType` for KAITO to autoprovision nodes, including CPU and GPU nodes.
+
+You can specify the instance type based on your cloud provider's offerings. For example, for Azure, you can specify a `Standard_D2ads_v5`, which is a CPU SKU like this:
+
+```yaml
+resource:
+  instanceType: "Standard_D2ads_v5"
+  labelSelector:
+    matchLabels:
+      apps: llama-3point2-3b
+```
+
+:::
+
+After workspace deployment succeeds, please refer to [Quick Start](quick-start.mdx#monitor-deployment) for monitoring the workspace and testing model inference.
+
+#### Custom Model Creation and Integration
+
+AIKit provides a simple way to create custom models without additional tools except for [Docker](https://docs.docker.com/desktop/install/linux-install/)!
+
+Here's an example on how to create a custom model and integrate it with KAITO:
+
+```bash
+export IMAGE_NAME="your-registry/your-model:latest"
+
+docker buildx build -t $IMAGE_NAME --push \
+    --build-arg="model=huggingface://TheBloke/Llama-2-7B-Chat-GGUF/llama-2-7b-chat.Q4_K_M.gguf" \
+    "https://raw.githubusercontent.com/sozercan/aikit/main/models/aikitfile.yaml"
+```
+
+After building the image, you can use it in your KAITO workspace configuration by updating the `image` field.
+
+For more information on creating custom models, refer to the [AIKit documentation](https://sozercan.github.io/aikit/docs/create-images).
+
+:::info
+AIKit supports a subset of backends, (such as [`llama.cpp`](https://sozercan.github.io/aikit/docs/llama-cpp), [`diffusers`](https://sozercan.github.io/aikit/docs/diffusion), [`exllamav2`](https://sozercan.github.io/aikit/docs/exllama2), and others) from [LocalAI](https://localai.io/) at this time. Please see [Inference Supported Backends](https://sozercan.github.io/aikit/docs/) section for more details, and updates.
+:::
diff --git a/website/versioned_docs/version-v0.6.x/aws.md b/website/versioned_docs/version-v0.6.x/aws.md
new file mode 100644
index 0000000..820a516
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/aws.md
@@ -0,0 +1,72 @@
+---
+title: AWS Setup
+---
+
+This guide covers setting up auto-provisioning capabilities for KAITO on Amazon Elastic Kubernetes Service (EKS). Auto-provisioning allows KAITO to automatically create GPU nodes when needed for your AI workloads.
+
+## Prerequisites
+
+- An EKS cluster with KAITO workspace controller installed (see [Installation](installation))
+- [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) for managing AWS resources
+- [eksctl](https://eksctl.io/installation/) for EKS cluster management
+- [kubectl](https://kubernetes.io/docs/tasks/tools/) configured to access your EKS cluster
+
+## Understanding Auto-Provisioning on AWS
+
+KAITO can use [Karpenter](https://karpenter.sh/) to automatically provision GPU nodes. This controller:
+
+- Creates new GPU nodes when workspaces require specific instance types
+- Supports various AWS GPU instances (g4, g5, p3, p4 series, etc.)
+- Manages node lifecycle based on workload demands
+- Integrates with AWS IAM for secure access
+
+:::note
+Alternative: If you already have GPU nodes or manage them separately, use the preferred nodes approach in the [Quick Start](quick-start) instead.
+:::
+
+## Set Up Auto-Provisioning
+
+### Create EKS Cluster and install Karpenter
+
+Follow the instructions [here](https://karpenter.sh/docs/getting-started/getting-started-with-karpenter/) to create an EKS cluster and install Karpenter.
+
+Then update the KAITO workspace controller Helm chart values for AWS:
+
+```bash
+helm update kaito-workspace --namespace kaito-workspace --set cloudProviderName=aws
+```
+
+### Using Auto-Provisioning
+
+Once Karpenter is set up, you can create workspaces that automatically provision GPU nodes:
+
+```yaml title="phi-4-workspace.yaml"
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-phi-4-mini
+resource:
+  instanceType: "g5.4xlarge" # Will trigger node creation
+  labelSelector:
+    matchLabels:
+      apps: phi-4-mini
+inference:
+  preset:
+    name: phi-4-mini-instruct
+```
+
+Apply the workspace:
+
+```bash
+kubectl apply -f phi-4-workspace.yaml
+```
+
+## Supported AWS GPU Instance Types
+
+The GPU provisioner supports various AWS GPU SKUs, see [supported options here](https://github.com/kaito-project/kaito/blob/main/pkg/sku/aws_sku_handler.go).
+
+For the complete list and specifications, see the [AWS GPU instance documentation](https://docs.aws.amazon.com/dlami/latest/devguide/gpu.html).
+
+## Clean Up
+
+See the Karpenter documentation for instructions on how to clean up your EKS cluster and remove Karpenter [here](https://karpenter.sh/docs/getting-started/getting-started-with-karpenter/#9-delete-the-cluster).
diff --git a/website/versioned_docs/version-v0.6.x/azure.md b/website/versioned_docs/version-v0.6.x/azure.md
new file mode 100644
index 0000000..4197a67
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/azure.md
@@ -0,0 +1,174 @@
+---
+title: Azure Setup
+---
+
+This guide covers setting up auto-provisioning capabilities for KAITO on Azure Kubernetes Service (AKS). Auto-provisioning allows KAITO to automatically create GPU nodes when needed for your AI workloads.
+
+## Prerequisites
+
+- An AKS cluster with KAITO workspace controller installed
+  - See [Step 1](#step-1-create-and-configure-an-aks-cluster) to create an AKS cluster
+  - See [Installation](installation) to install the KAITO workspace controller
+- [Azure CLI](https://learn.microsoft.com/cli/azure/install-azure-cli) for managing Azure resources
+- [kubectl](https://kubernetes.io/docs/tasks/tools/) configured to access your AKS cluster
+
+## Understanding Auto-Provisioning on Azure
+
+KAITO can use the [Azure GPU Provisioner](https://github.com/Azure/gpu-provisioner) to automatically provision GPU nodes. This controller:
+
+- Creates new GPU nodes when workspaces require specific instance types
+- Supports various Azure GPU SKUs (Standard_NC series, etc.)
+- Manages node lifecycle based on workload demands
+- Integrates with Azure's managed identity system for secure access
+
+### When to Use Auto-Provisioning
+
+Choose auto-provisioning when:
+- You want KAITO to manage GPU node creation automatically
+- Your workloads have varying GPU requirements
+- You prefer to specify exact Azure instance types in your workspaces
+
+:::note
+Alternative: If you already have GPU nodes or manage them separately, use the [preferred nodes approach](quick-start#option-1-using-preferred-nodes-existing-gpu-nodes) instead.
+:::
+
+## Set Up Auto-Provisioning
+
+### Step 1: Create and configure an AKS Cluster
+
+If you don't have an AKS cluster yet, you can create one using the Azure CLI:
+
+```bash
+export RESOURCE_GROUP="kaito-rg"
+export CLUSTER_NAME="kaito-cluster"
+export LOCATION="eastus"
+az group create --name $RESOURCE_GROUP --location $LOCATION
+az aks create --resource-group $RESOURCE_GROUP --name $CLUSTER_NAME --enable-oidc-issuer --enable-workload-identity --enable-managed-identity --generate-ssh-keys
+```
+
+Connect to the cluster:
+```bash
+az aks get-credentials --resource-group $RESOURCE_GROUP --name $CLUSTER_NAME
+```
+
+### Step 2: Create Managed Identity
+
+Create a managed identity for the GPU provisioner with the necessary permissions:
+
+```bash
+export SUBSCRIPTION=$(az account show --query id -o tsv)
+export IDENTITY_NAME="kaitoprovisioner"
+
+# Create the managed identity
+az identity create --name $IDENTITY_NAME -g $RESOURCE_GROUP
+
+# Get the principal ID for role assignment
+export IDENTITY_PRINCIPAL_ID=$(az identity show --name $IDENTITY_NAME -g $RESOURCE_GROUP --subscription $SUBSCRIPTION --query 'principalId' -o tsv)
+
+# Assign Contributor role to the cluster
+az role assignment create \
+  --assignee $IDENTITY_PRINCIPAL_ID \
+  --scope /subscriptions/$SUBSCRIPTION/resourceGroups/$RESOURCE_GROUP/providers/Microsoft.ContainerService/managedClusters/$CLUSTER_NAME \
+  --role "Contributor"
+```
+
+### Step 3: Install GPU Provisioner
+
+Install the Azure GPU Provisioner using Helm:
+
+```bash
+export GPU_PROVISIONER_VERSION=0.3.5
+
+# Download and configure Helm values
+curl -sO https://raw.githubusercontent.com/Azure/gpu-provisioner/main/hack/deploy/configure-helm-values.sh
+chmod +x ./configure-helm-values.sh && ./configure-helm-values.sh $CLUSTER_NAME $RESOURCE_GROUP $IDENTITY_NAME
+
+# Install GPU provisioner
+helm install gpu-provisioner \
+  --values gpu-provisioner-values.yaml \
+  --set settings.azure.clusterName=$CLUSTER_NAME \
+  --wait \
+  https://github.com/Azure/gpu-provisioner/raw/gh-pages/charts/gpu-provisioner-$GPU_PROVISIONER_VERSION.tgz \
+  --namespace gpu-provisioner \
+  --create-namespace
+```
+
+### Step 4: Create Federated Credential
+
+Create the federated identity credential to allow the GPU provisioner to access Azure resources:
+
+```bash
+export AKS_OIDC_ISSUER=$(az aks show -n $CLUSTER_NAME -g $RESOURCE_GROUP --subscription $SUBSCRIPTION --query "oidcIssuerProfile.issuerUrl" -o tsv)
+
+az identity federated-credential create \
+  --name kaito-federatedcredential \
+  --identity-name $IDENTITY_NAME \
+  -g $RESOURCE_GROUP \
+  --issuer $AKS_OIDC_ISSUER \
+  --subject system:serviceaccount:"gpu-provisioner:gpu-provisioner" \
+  --audience api://AzureADTokenExchange \
+  --subscription $SUBSCRIPTION
+```
+
+## Verify Setup
+
+Check that the GPU provisioner is running correctly:
+
+```bash
+# Check Helm installations
+helm list -n gpu-provisioner
+helm list -n kaito-workspace
+
+# Check GPU provisioner status
+kubectl describe deploy gpu-provisioner -n gpu-provisioner
+kubectl get pods -n gpu-provisioner
+```
+
+The GPU provisioner pod should be in a `Running` state. If it's failing, check the logs:
+
+```bash
+kubectl logs --selector=app.kubernetes.io/name=gpu-provisioner -n gpu-provisioner
+```
+
+## Using Auto-Provisioning
+
+Once set up, you can create workspaces that automatically provision GPU nodes:
+
+```yaml title="phi-4-workspace.yaml"
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-phi-4-mini
+resource:
+  instanceType: "Standard_NC6s_v3"  # Will trigger node creation
+  labelSelector:
+    matchLabels:
+      apps: phi-4-mini
+inference:
+  preset:
+    name: phi-4-mini-instruct
+```
+
+Then apply the workspace:
+
+```bash
+kubectl apply -f phi-4-workspace.yaml
+```
+
+## Supported Azure GPU Instance Types
+
+The GPU provisioner supports various Azure GPU SKUs, see [supported options here](https://github.com/kaito-project/kaito/blob/main/pkg/sku/azure_sku_handler.go).
+
+For the complete list and specifications, see the [Azure GPU-optimized VM sizes documentation](https://learn.microsoft.com/en-us/azure/virtual-machines/sizes-gpu).
+
+## Clean Up
+
+To remove the auto-provisioning setup:
+
+```bash
+# Uninstall GPU provisioner
+helm uninstall gpu-provisioner -n gpu-provisioner
+
+# Delete the managed identity (optional)
+az identity delete --name $IDENTITY_NAME -g $RESOURCE_GROUP
+```
diff --git a/website/versioned_docs/version-v0.6.x/contributing.md b/website/versioned_docs/version-v0.6.x/contributing.md
new file mode 100644
index 0000000..3aec963
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/contributing.md
@@ -0,0 +1,100 @@
+---
+title: Contributing
+---
+
+This project welcomes contributions and suggestions.
+
+KAITO has adopted the [Cloud Native Compute Foundation Code of Conduct](https://github.com/cncf/foundation/blob/main/code-of-conduct.md). For more information see the [KAITO Code of Conduct](https://github.com/kaito-project/kaito/blob/main/CODE_OF_CONDUCT.md).
+
+## Development
+
+### Prerequisites
+
+- [Go 1.24](https://go.dev/doc/install) or later
+- [Tilt](https://tilt.dev/) for rapid, local development
+- [Docker](https://docs.docker.com/get-started/get-docker/) for building container images
+- Kubernetes cluster with Karpenter supported (e.g. AKS, EKS)
+
+### Getting Started
+
+#### 1. Fork and Clone the Repository
+
+```bash
+git clone git@github.com:<username>/kaito.git
+cd kaito
+```
+
+#### 2. Create `tilt-settings.yaml`
+
+Tilt streamlines Kubernetes development by automating image builds, deployments, and live updates. Instead of manually building and pushing images to a registry, Tilt handles this process and hot-reloads changes (Go binaries and CRD YAML files) in real-time.
+
+To configure Tilt, create a tilt-settings.yaml file in the repository root:
+
+```bash
+cat <<EOF > tilt-settings.yaml
+default_registry: <acr_name>.azurecr.io
+allowed_contexts:
+  - <context_name>
+cluster_name: <cluster_name>
+EOF
+```
+
+- `default_registry`: The container registry where Tilt pushes development images (e.g. `<acr_name>.azurecr.io`). Ensure you're logged in with `docker login` and have push/pull permissions.
+If you want to use Azure Container Registry (ACR), set up a new ACR instance and attach it to your AKS cluster with these commands:
+
+```bash
+export ACR_NAME=<acr_name>
+export RESOURCE_GROUP=<resource_group>
+export CLUSTER_NAME=<cluster_name>
+
+az acr -n $ACR_NAME create --resource-group $RESOURCE_GROUP --sku Basic
+az acr login -n $ACR_NAME
+az aks update -n $CLUSTER_NAME -g $RESOURCE_GROUP --attach-acr $ACR_NAME
+```
+
+- `allowed_contexts`: The Kubernetes contexts Tilt is permitted to use. List your available contexts with `kubectl config get-contexts`.
+
+- `cluster_name`: The name of your Kubernetes cluster. If omitted, Tilt defaults to your current kubectl context.
+
+#### 3. Start Tilt
+
+To start Tilt, run the following command in the root of the repository:
+
+```bash
+tilt up
+```
+
+<details>
+<summary>Example Output</summary>
+
+<!-- markdown-link-check-disable -->
+
+```bash
+Tilt started on http://localhost:10350/
+v0.34.0, built 2025-03-11
+
+(space) to open the browser
+(s) to stream logs (--stream=true)
+(t) to open legacy terminal mode (--legacy=true)
+(ctrl-c) to exit
+Opening browser: http://localhost:10350/
+```
+
+<!-- markdown-link-check-enable -->
+
+</details>
+
+<br/>
+
+![Tilt UI](/img/tilt.png)
+
+Tilt automates the following steps during `tilt up`:
+
+<!-- markdown-link-check-disable -->
+1. **Start the Tilt Server**: Runs locally at http://localhost:10350/ and opens the Tilt UI in your browser (press space).
+<!-- markdown-link-check-enable -->
+1. **Generate and Build**: Creates CRD YAML files (`make generate manifests`) and compiles the controller manager Go binary (`go build`), saving it to `./tilt_bin`.
+2. **Build Development Image**: Packages the binary into a container image and tag it with `<default_registry>/kaito/workspace:tilt-<hash>`), then pushes it to your container registry.
+3. **Apply CRDs**: Deploys CRDs to the cluster with `kubectl apply --server-side -f charts/kaito/crds/`.
+4. **Deploy Resources**: Renders the Helm chart (`helm template charts/kaito/workspace`) and applies it to the cluster, including the controller manager, webhooks, RBAC, and NVIDIA device plugin resources.
+5. **Live Updates**: Detects changes to the controller manager source code, rebuilds the binary, updates the container, and restarts the controller manager within the container without re-deploying the pod - all in ~30 seconds.
diff --git a/website/versioned_docs/version-v0.6.x/custom-model.md b/website/versioned_docs/version-v0.6.x/custom-model.md
new file mode 100644
index 0000000..f6955b8
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/custom-model.md
@@ -0,0 +1,83 @@
+---
+title: Custom Model Integration
+---
+
+## Option 1: Use Pre-Built Docker Image Without Model Weights
+If you want to avoid building a Docker image with model weights, use our pre-built reference image (`ghcr.io/kaito-project/kaito/llm-reference-preset:latest`). This image, built with [Dockerfile.reference](../../../examples/custom-model-integration/Dockerfile.reference), dynamically downloads model weights from HuggingFace at runtime, reducing the need to create and maintain custom images.
+
+
+- **[Sample Deployment YAML](../../../examples/custom-model-integration/reference-image-deployment.yaml)**
+
+
+## Option 2: Build a Custom Docker Image with Model Weights
+
+### Step 1: Clone the Repository
+
+```sh
+git clone https://github.com/kaito-project/kaito.git
+```
+
+### Step 2: Download Your Private/Custom Model Weights
+
+For example, assuming HuggingFace weights:
+```sh
+git lfs install
+git clone git@hf.co:<MODEL_ID>  # Example: git clone git@hf.co:bigscience/bloom
+# OR
+git clone https://huggingface.co/bigscience/bloom
+```
+
+Alternatively, use curl:
+```
+curl -sSL https://huggingface.co/bigscience/bloom/resolve/main/config.json?download=true -o config.json
+```
+
+More information on downloading models from HuggingFace can be found [here](https://huggingface.co/docs/hub/en/models-downloading).
+
+
+### Step 3: Log In to Your Container Registry
+
+Before pushing the Docker image, ensure you’re logged into the appropriate container registry. Here are general login methods depending on the registry you use:
+
+1. GitHub Container Registry (ghcr.io):
+```sh
+echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdin
+```
+Replace CR_PAT with your GitHub Personal Access Token and USERNAME with your GitHub username. This token should have the write:packages and read:packages permissions.
+
+2. Azure Container Registry (ACR): If you're using ACR:
+
+```sh
+az acr login --name <REGISTRY_NAME>
+```
+Replace `<REGISTRY_NAME>` with your Azure Container Registry name.
+
+3. Docker Hub or Other Container Registries:
+```sh
+docker login <REGISTRY_URL>
+```
+Enter your username and password when prompted. Replace `<REGISTRY_URL>` with your registry URL, such as `docker.io` for Docker Hub.
+
+
+### Step 4: Build Docker Image with Private/Custom Weights
+
+1. Set Environment Variables
+
+Before building the Docker image, set the relevant environment variables for the image name, version, and weights path:
+```sh
+export IMAGE_NAME="modelsregistry.azurecr.io/phi-3-mini-4k-instruct:0.0.1"
+export VERSION="0.0.1"
+export WEIGHTS_PATH="kaito/phi-3-mini-4k-instruct/weights"
+```
+
+2. Build and Push the Docker Image
+
+Navigate to the KAITO base directory and build the Docker image, ensuring the weights directory is included in the build context:
+```sh
+docker build -t <IMAGE_NAME> --file docker/presets/workspace/models/tfs/Dockerfile --build-arg WEIGHTS_PATH=<WEIGHTS_PATH> --build-arg MODEL_TYPE=text-generation --build-arg VERSION=<VERSION> .
+
+docker push <IMAGE_NAME>
+```
+
+### Step 5: Deploy
+Follow the [Custom Template](../../../examples/custom-model-integration/custom-deployment-template.yaml)
\ No newline at end of file
diff --git a/website/versioned_docs/version-v0.6.x/example_rag_client.py b/website/versioned_docs/version-v0.6.x/example_rag_client.py
new file mode 100644
index 0000000..e2a5760
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/example_rag_client.py
@@ -0,0 +1,134 @@
+# Copyright (c) KAITO authors.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import requests
+
+
+class KAITORAGClient:
+    """
+    A bare bones client for interacting with the RAGEngine API.
+    This client provides methods to index documents, query the engine,
+    update documents, delete documents, list documents, and manage indexes.
+    """
+
+    def __init__(self, base_url):
+        self.base_url = base_url
+        self.headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+
+    def index_documents(self, index_name, documents):
+        """
+        Index documents in the RAGEngine.
+        documents: list of dicts, each with 'text' and optional 'metadata'
+        """
+        url = f"{self.base_url}/index"
+        payload = {"index_name": index_name, "documents": documents}
+        resp = requests.post(url, json=payload, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
+
+    def query(self, index_name, query, llm_temperature, llm_max_tokens, top_k=5):
+        """
+        Query the RAGEngine.
+        query: str, the query text
+        top_k: int, number of results to return
+        metadata: optional dict, additional metadata for the query
+        """
+        url = f"{self.base_url}/query"
+        payload = {
+            "index_name": index_name,
+            "query": query,
+            "top_k": top_k,
+            "llm_params": {
+                "temperature": llm_temperature,
+                "max_tokens": llm_max_tokens,
+            },
+        }
+        resp = requests.post(url, json=payload, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
+
+    def update_documents(self, index_name, documents):
+        """
+        Update a document in the RAGEngine.
+        documents: list of dicts, each with 'id', 'text', and optional 'metadata'
+        """
+        url = f"{self.base_url}/indexes/{index_name}/documents"
+        payload = {"documents": documents}
+        resp = requests.post(url, json=payload, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
+
+    def delete_documents(self, index_name, document_ids):
+        """
+        Delete a document from the RAGEngine by its ID.
+        """
+        url = f"{self.base_url}/indexes/{index_name}/documents/delete"
+        payload = {"doc_ids": document_ids}
+        resp = requests.post(url, json=payload, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
+
+    def list_documents(self, index_name, metadata_filter, limit=10, offset=0):
+        """
+        List documents in the RAGEngine.
+        """
+        url = f"{self.base_url}/indexes/{index_name}/documents"
+        params = {"limit": limit, "offset": offset}
+        if metadata_filter:
+            params["metadata_filter"] = json.dumps(metadata_filter)
+        resp = requests.get(url, headers=self.headers, params=params)
+        resp.raise_for_status()
+        return resp.json()
+
+    def list_indexes(self):
+        """
+        List all indexes in the RAGEngine.
+        """
+        url = f"{self.base_url}/indexes"
+        resp = requests.get(url, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
+
+    def persist_index(self, index_name, path="/tmp"):
+        """
+        Persist an index in the RAGEngine.
+        """
+        url = f"{self.base_url}/persist/{index_name}"
+        params = {"path": path}
+        resp = requests.post(url, params=params, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
+
+    def load_index(self, index_name, path="/tmp", overwrite=True):
+        """
+        Load an index in the RAGEngine.
+        """
+        url = f"{self.base_url}/load/{index_name}"
+        params = {"path": path, "overwrite": overwrite}
+        resp = requests.post(url, params=params, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
+
+    def delete_index(self, index_name):
+        """
+        Delete an index from the RAGEngine.
+        """
+        url = f"{self.base_url}/indexes/{index_name}"
+        resp = requests.delete(url, headers=self.headers)
+        resp.raise_for_status()
+        return resp.json()
diff --git a/website/versioned_docs/version-v0.6.x/faq.md b/website/versioned_docs/version-v0.6.x/faq.md
new file mode 100644
index 0000000..ba41b80
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/faq.md
@@ -0,0 +1,59 @@
+---
+title: FAQ
+---
+
+### How do I ensure preferred nodes are correctly labeled for use in my workspace?
+
+For using preferred nodes, make sure the node has the label specified in the labelSelector under matchLabels. For example, if your labelSelector is:
+
+```yaml
+labelSelector:
+  matchLabels:
+    apps: falcon-7b
+```
+
+Then the node should have the label: `apps=falcon-7b`.
+
+### How to upgrade the existing deployment to use the latest model configuration?
+
+When using hosted public models, you can delete the existing inference workload (`Deployment` or `StatefulSet`) manually, and the workspace controller will create a new one with the latest preset configuration (e.g., the image version) defined in the current release.
+
+For private models, it is recommended to create a new workspace with a new image version in the Spec.
+
+### How to update model/inference parameters to override the KAITO Preset Configuration?
+
+KAITO provides a limited capability to override preset configurations for models that use `transformer` runtime manually.
+
+To update parameters for a deployed model, perform `kubectl edit` against the workload, which could be either a `StatefulSet` or `Deployment`.
+
+For example, to enable 4-bit quantization on a `falcon-7b-instruct` deployment:
+
+```bash
+kubectl edit deployment workspace-falcon-7b-instruct
+```
+
+Within the deployment specification, locate and modify the command field.
+
+**Original:**
+```bash
+accelerate launch --num_processes 1 --num_machines 1 --machine_rank 0 --gpu_ids all inference_api.py --pipeline text-generation --torch_dtype bfloat16
+```
+
+**Modified to enable 4-bit Quantization:**
+```bash
+accelerate launch --num_processes 1 --num_machines 1 --machine_rank 0 --gpu_ids all inference_api.py --pipeline text-generation --torch_dtype bfloat16 --load_in_4bit
+```
+
+Currently, we allow users to change the following parameters manually:
+
+- `pipeline`: For text-generation models this can be either `text-generation` or `conversational`.
+- `load_in_4bit` or `load_in_8bit`: Model quantization resolution.
+
+Should you need to customize other parameters, kindly file an issue for potential future inclusion.
+
+### What is the difference between instruct and non-instruct models?
+
+The main distinction lies in their intended use cases:
+
+- **Instruct models**: Fine-tuned versions optimized for interactive chat applications. They are typically the preferred choice for most implementations due to their enhanced performance in conversational contexts.
+- **Non-instruct (raw) models**: Designed for further fine-tuning with your own data.
diff --git a/website/versioned_docs/version-v0.6.x/headlamp-kaito.md b/website/versioned_docs/version-v0.6.x/headlamp-kaito.md
new file mode 100644
index 0000000..d93d300
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/headlamp-kaito.md
@@ -0,0 +1,36 @@
+---
+title: Headlamp KAITO
+---
+
+![Headlamp](/img/Headlamp.png)
+
+Headlamp plugin for KAITO is a web-based dashboard extension for managing and monitoring KAITO Workspaces through the Headlamp Kubernetes dashboard. This extension provides an intuitive interface for operators to deploy, monitor, and manage AI workloads.
+
+## Overview
+
+Headlamp plugin for KAITO is a plugin for [Headlamp](https://headlamp.dev/), the Kubernetes web UI, that adds specialized views and management capabilities for KAITO resources. Key features include:
+
+- **Visual Workspace Management**: Create, view, and manage KAITO workspaces through a user-friendly interface
+- **Resource Monitoring**: Real-time visibility into GPU utilization, inference metrics, and system health
+- **Model Deployment**: Simplified deployment of pre-configured AI models from KAITO's preset library
+
+![Headlamp-KAITO plugin installation](/img/headlamp-kaito-plugin-installation.png)
+
+## Getting Started
+
+For detailed installation instructions, usage guides, and configuration options, please refer to the official documentation:
+
+![Headlamp-KAITO Documentation](/img/headlamp-kaito-docs.png)
+
+**[Headlamp-KAITO Documentation](https://kaito-project.github.io/headlamp-kaito/)**
+
+## Resources
+
+- **[GitHub Repository](https://github.com/kaito-project/headlamp-kaito)** - Source code, issues, and contributions
+- **[Documentation](https://kaito-project.github.io/headlamp-kaito/)** - Complete setup and usage guide
+- **[Headlamp](https://headlamp.dev/)** - The underlying Kubernetes dashboard
+
+## Support
+
+For issues:
+- Report bugs and request features on [GitHub Issues](https://github.com/kaito-project/headlamp-kaito/issues)
diff --git a/website/versioned_docs/version-v0.6.x/inference.md b/website/versioned_docs/version-v0.6.x/inference.md
new file mode 100644
index 0000000..b1694e0
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/inference.md
@@ -0,0 +1,258 @@
+---
+title: Inference
+---
+
+This document presents how to use the KAITO `workspace` Custom Resource Definition (CRD) for model serving and serving with LoRA adapters.
+
+:::tip Multi-Node Inference
+For large models requiring multiple nodes, see the [Multi-Node Inference](./multi-node-inference.md) documentation.
+:::
+
+## Usage
+
+The basic usage for inference is simple. Users just need to specify the GPU SKU used for inference in the `resource` spec and one of the KAITO supported model name in the `inference` spec in the `workspace` custom resource. For example,
+
+```yaml
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-falcon-7b
+resource:
+  instanceType: "Standard_NC24ads_A100_v4"
+  labelSelector:
+    matchLabels:
+      apps: falcon-7b
+inference:
+  preset:
+    name: "falcon-7b"
+```
+
+If a user runs KAITO in an on-premise Kubernetes cluster where GPU SKUs are unavailable, the GPU nodes can be pre-configured. The user should ensure that the corresponding vendor-specific GPU plugin is installed successfully in every prepared node, i.e. the node status should report a non-zero GPU resource in the allocatable field. For example:
+
+```
+$ kubectl get node $NODE_NAME -o json | jq .status.allocatable
+{
+  "cpu": "XXXX",
+  "ephemeral-storage": "YYYY",
+  "hugepages-1Gi": "0",
+  "hugepages-2Mi": "0",
+  "memory": "ZZZZ",
+  "nvidia.com/gpu": "1",
+  "pods": "100"
+}
+```
+
+Next, the user needs to add the node names in the `preferredNodes` field in the `resource` spec. As a result, the KAITO controller will skip the steps for GPU node provisioning and use the prepared nodes to run the inference workload.
+
+:::warning
+The node objects of the preferred nodes need to contain the same matching labels as specified in the `resource` spec. Otherwise, the KAITO controller would not recognize them.
+:::
+
+### Inference runtime selection
+
+KAITO now supports both [vLLM](https://github.com/vllm-project/vllm) and [transformers](https://github.com/huggingface/transformers) runtime. `vLLM` provides better serving latency and throughput. `transformers` provides more compatibility with models in the Huggingface model hub.
+
+From KAITO v0.4.0, the default runtime is switched to `vLLM`. If you want to use `transformers` runtime, you can specify the runtime in the `inference` spec using an annotation. For example,
+
+```yaml
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-falcon-7b
+  annotations:
+    kaito.sh/runtime: "transformers"
+resource:
+  instanceType: "Standard_NC24ads_A100_v4"
+  labelSelector:
+    matchLabels:
+      apps: falcon-7b
+inference:
+  preset:
+    name: "falcon-7b"
+```
+
+:::note Multi-Node Support
+Multi-node distributed inference is currently supported only with the vLLM runtime. For details on configuring multi-node deployments, see [Multi-Node Inference](./multi-node-inference.md).
+:::
+
+### Inference with custom parameters
+
+Users can customize vLLM runtime parameters by creating a ConfigMap containing an `inference_config.yaml` file and referencing it in the workspace spec. For example:
+
+```yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  namespace: myns
+  name: my-inference-params
+data:
+  inference_config.yaml: |
+    max_probe_steps: 6
+    vllm:
+      gpu-memory-utilization: 0.95  # Controls GPU memory usage (0.0-1.0)
+      tensor-parallel-size: 2        # Number of GPUs for tensor parallelism
+      max-model-len: 131072         # Maximum sequence length
+      swap-space: 4                 # CPU swap space in GB
+      cpu-offload-gb: 0             # Amount of GPU memory to offload to CPU
+---
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  namespace: myns
+  name: workspace-example
+resource:
+  instanceType: "Standard_NC24ads_A100_v4"
+  labelSelector:
+    matchLabels:
+      apps: example
+inference:
+  preset:
+    name: "example-model"
+  config: "my-inference-params"  # Reference to ConfigMap name
+```
+
+Key vLLM parameters include:
+- `gpu-memory-utilization`: Controls fraction of GPU memory allocated (between 0.0 and 1.0)
+- `tensor-parallel-size`: Number of GPUs to use for tensor parallelism
+- `max-model-len`: Maximum sequence length the model can handle
+
+For the complete list of vLLM parameters, refer to the [vLLM documentation](https://docs.vllm.ai/en/latest/serving/engine_args.html).
+
+### Inference with LoRA adapters
+
+KAITO also supports running the inference workload with LoRA adapters produced by [model fine-tuning jobs](./tuning.md). Users can specify one or more adapters in the `adapters` field of the `inference` spec. For example,
+
+```yaml
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-falcon-7b
+resource:
+  instanceType: "Standard_NC24ads_A100_v4"
+  labelSelector:
+    matchLabels:
+      apps: falcon-7b
+inference:
+  preset:
+    name: "falcon-7b"
+  adapters:
+    - source:
+        name: "falcon-7b-adapter"
+        image:  "<YOUR_IMAGE>"
+      strength: "0.2"
+```
+Currently, only images are supported as adapter sources. The `strength` field specifies the multiplier applied to the adapter weights relative to the raw model weights.
+
+**Note:** When building a container image for an existing adapter, ensure all adapter files are copied to the **/data** directory inside the container.
+
+For detailed `InferenceSpec` API definitions, refer to the [documentation](https://github.com/kaito-project/kaito/blob/2ccc93daf9d5385649f3f219ff131ee7c9c47f3e/api/v1alpha1/workspace_types.go#L75).
+
+### Inference API
+
+The OpenAPI specification for the inference API is available at [vLLM API](../../../presets/workspace/inference/vllm/api_spec.json), [transformers API](../../../presets/workspace/inference/text-generation/api_spec.json).
+
+#### vLLM inference API
+
+vLLM supports OpenAI-compatible inference APIs. Check [here](https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html) for more details.
+
+#### Transformers inference API
+
+The inference service endpoint is `/chat`.
+
+**basic example**
+```
+curl -X POST "http://<SERVICE>:80/chat" -H "accept: application/json" -H "Content-Type: application/json" -d '{"prompt":"YOUR_PROMPT_HERE"}'
+```
+
+**example with full configurable parameters**
+```
+curl -X POST \
+    -H "accept: application/json" \
+    -H "Content-Type: application/json" \
+    -d '{
+        "prompt":"YOUR_PROMPT_HERE",
+        "return_full_text": false,
+        "clean_up_tokenization_spaces": false,
+        "prefix": null,
+        "handle_long_generation": null,
+        "generate_kwargs": {
+                "max_length":200,
+                "min_length":0,
+                "do_sample":true,
+                "early_stopping":false,
+                "num_beams":1,
+                "num_beam_groups":1,
+                "diversity_penalty":0.0,
+                "temperature":1.0,
+                "top_k":10,
+                "top_p":1,
+                "typical_p":1,
+                "repetition_penalty":1,
+                "length_penalty":1,
+                "no_repeat_ngram_size":0,
+                "encoder_no_repeat_ngram_size":0,
+                "bad_words_ids":null,
+                "num_return_sequences":1,
+                "output_scores":false,
+                "return_dict_in_generate":false,
+                "forced_bos_token_id":null,
+                "forced_eos_token_id":null,
+                "remove_invalid_values":null
+            }
+        }' \
+        "http://<SERVICE>:80/chat"
+```
+
+**parameters**
+- `prompt`: The initial text provided by the user, from which the model will continue generating text.
+- `return_full_text`: If False only generated text is returned, else full text is returned.
+- `clean_up_tokenization_spaces`: True/False, determines whether to remove potential extra spaces in the text output.
+- `prefix`: Prefix added to the prompt.
+- `handle_long_generation`: Provides strategies to address generations beyond the model's maximum length capacity.
+- `max_length`: The maximum total number of tokens in the generated text.
+- `min_length`: The minimum total number of tokens that should be generated.
+- `do_sample`: If True, sampling methods will be used for text generation, which can introduce randomness and variation.
+- `early_stopping`: If True, the generation will stop early if certain conditions are met, for example, when a satisfactory number of candidates have been found in beam search.
+- `num_beams`: The number of beams to be used in beam search. More beams can lead to better results but are more computationally expensive.
+- `num_beam_groups`: Divides the number of beams into groups to promote diversity in the generated results.
+- `diversity_penalty`: Penalizes the score of tokens that make the current generation too similar to other groups, encouraging diverse outputs.
+- `temperature`: Controls the randomness of the output by scaling the logits before sampling.
+- `top_k`: Restricts sampling to the k most likely next tokens.
+- `top_p`: Uses nucleus sampling to restrict the sampling pool to tokens comprising the top p probability mass.
+- `typical_p`: Adjusts the probability distribution to favor tokens that are "typically" likely, given the context.
+- `repetition_penalty`: Penalizes tokens that have been generated previously, aiming to reduce repetition.
+- `length_penalty`: Modifies scores based on sequence length to encourage shorter or longer outputs.
+- `no_repeat_ngram_size`: Prevents the generation of any n-gram more than once.
+- `encoder_no_repeat_ngram_size`: Similar to `no_repeat_ngram_size` but applies to the encoder part of encoder-decoder models.
+- `bad_words_ids`: A list of token ids that should not be generated.
+- `num_return_sequences`: The number of different sequences to generate.
+- `output_scores`: Whether to output the prediction scores.
+- `return_dict_in_generate`: If True, the method will return a dictionary containing additional information.
+- `pad_token_id`: The token ID used for padding sequences to the same length.
+- `eos_token_id`: The token ID that signifies the end of a sequence.
+- `forced_bos_token_id`: The token ID that is forcibly used as the beginning of a sequence token.
+- `forced_eos_token_id`: The token ID that is forcibly used as the end of a sequence when max_length is reached.
+- `remove_invalid_values`: If True, filters out invalid values like NaNs or infs from model outputs to prevent crashes.
+
+# Inference workload
+
+Depending on whether the specified model supports distributed inference or not, the KAITO controller will choose to use either Kubernetes **apps.deployment** workload (by default) or Kubernetes **apps.statefulset** workload (if the model supports distributed inference) to manage the inference service, which is exposed using a Cluster-IP type of Kubernetes `service`.
+
+For multi-node distributed inference, KAITO uses StatefulSets to ensure stable pod identity and coordination between leader and worker pods. See [Multi-Node Inference](./multi-node-inference.md) for detailed architecture information.
+
+When adapters are specified in the `inference` spec, the KAITO controller adds an initcontainer for each adapter in addition to the main container. The pod structure is shown in Figure 1.
+
+![KAITO inference service pod structure](/img/kaito-inference-adapter.png)
+
+If an image is specified as the adapter source, the corresponding initcontainer uses that image as its container image. These initcontainers ensure all adapter data is available locally before the inference service starts. The main container uses a supported model image, launching the [inference_api.py](../../../presets/workspace/inference/text-generation/inference_api.py) script.
+
+All containers share local volumes by mounting the same `EmptyDir` volumes, avoiding file copies between containers.
+
+## Workload update
+
+To update the `adapters` field in the `inference` spec, users can modify the `workspace` custom resource. The KAITO controller will apply the changes, triggering a workload deployment update. This will recreate the inference service pod, resulting in a brief service downtime. Once the new adapters are merged with the raw model weights and loaded into GPU memory, the service will resume.
+
+
+# Troubleshooting
+
+TBD
\ No newline at end of file
diff --git a/website/versioned_docs/version-v0.6.x/installation.md b/website/versioned_docs/version-v0.6.x/installation.md
new file mode 100644
index 0000000..d76f8c0
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/installation.md
@@ -0,0 +1,91 @@
+---
+title: Installation
+---
+
+KAITO (Kubernetes AI Toolchain Operator) can be installed on any Kubernetes cluster using Helm. This guide covers the basic installation of the KAITO workspace controller.
+
+## Prerequisites
+
+Before you begin, ensure you have the following tools installed:
+
+- An existing Kubernetes cluster (can be hosted on any cloud provider) with NVIDIA GPU nodes
+  - For cloud provider-specific guides, see:
+    - [Azure Setup using AKS](azure)
+    - [AWS Setup using EKS](aws)
+- [Helm](https://helm.sh) to install the operator
+- [kubectl](https://kubernetes.io/docs/tasks/tools/) to interact with your Kubernetes cluster
+
+## Install KAITO Workspace Controller
+
+Install the KAITO workspace controller using Helm:
+
+```bash
+export KAITO_WORKSPACE_VERSION=0.6.0
+export CLUSTER_NAME=kaito
+
+helm install kaito-workspace \
+  https://github.com/kaito-project/kaito/raw/gh-pages/charts/kaito/workspace-$KAITO_WORKSPACE_VERSION.tgz \
+  --namespace kaito-workspace \
+  --create-namespace \
+  --set clusterName="$CLUSTER_NAME" \
+  --wait
+```
+
+### Verify KAITO Installation
+
+Check that the KAITO workspace controller is running:
+
+```bash
+kubectl get pods -n kaito-workspace
+kubectl describe deploy kaito-workspace -n kaito-workspace
+```
+
+You should see the workspace controller pod in a `Running` state.
+
+## Setup GPU Nodes
+
+You need to create GPU nodes in order to run a Workspace with KAITO. There are two options:
+
+- **Bring your own GPU (BYO) nodes**: Create your own GPU nodes to run KAITO deployments on. 
+- **Auto-provisioning**: Set up automatic GPU node provisioning for your cloud provider. 
+
+### Option 1: Bring your own GPU nodes
+
+:::tip
+This is recommended if you just want to try KAITO with the least amount of setup required or if you only plan to use 1-2 GPU nodes. If your cloud provider is not listed in the [Auto-provisioning section](#option-2-auto-provision-gpu-nodes), you must create your own GPU nodes. If it is, you can still manually create the nodes if you don't want to set up auto-provisioning.
+:::
+
+Create your GPU nodes and label them for quick access. In these docs, we will use the label `accelerator=nvidia` but any label can work.
+
+List your GPU nodes and verify that they are all present and ready.
+
+```bash
+kubectl get nodes -l accelerator=nvidia
+```
+
+The output should look similar to this. Note the names of the nodes as they will be specified when deploying a Workspace.
+
+```
+NAME                                  STATUS   ROLES    AGE     VERSION
+gpunp-26695285-vmss000000             Ready    <none>   2d21h   v1.31.9
+gpunp-26695285-vmss000001             Ready    <none>   2d21h   v1.31.9
+```
+
+### Option 2: Auto-provision GPU nodes
+
+:::tip
+This is recommended for production environments and any non-trivial use cases, or if you plan to use a number of nodes.
+:::
+
+The following cloud providers support auto-provisioning GPU nodes in addition to BYO nodes.
+
+- [Azure (AKS)](azure#setup-auto-provisioning) - Set up GPU auto-provisioning with Azure GPU Provisioner
+- [AWS (EKS)](aws#setup-auto-provisioning) - Set up GPU auto-provisioning with Karpenter
+
+## Next Steps
+
+Once KAITO is installed, you can:
+
+- Follow the [Quick Start](quick-start) guide to deploy your first model
+- Try auto-provisioning with cloud provider if you have not done so
+- Explore the available [model presets](presets)
diff --git a/website/versioned_docs/version-v0.6.x/intro.md b/website/versioned_docs/version-v0.6.x/intro.md
new file mode 100644
index 0000000..7b87767
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/intro.md
@@ -0,0 +1,67 @@
+---
+title: Introduction
+slug: /
+---
+
+:::info What's NEW!
+
+Retrieval Augmented Generation (RAG) support is live! - KAITO RagEngine uses LlamaIndex and FAISS, learn about it [here](https://kaito-project.github.io/kaito/docs/rag)!
+**Latest Release:** July 18th, 2025. KAITO v0.5.1.
+
+**First Release:** Nov 15th, 2023. KAITO v0.1.0.
+:::
+
+KAITO is an operator that automates the AI/ML model inference or tuning workload in a Kubernetes cluster.
+The target models are popular open-sourced large models such as [falcon](https://huggingface.co/tiiuae) and [phi-3](https://huggingface.co/docs/transformers/main/en/model_doc/phi3).
+
+## Key Features
+
+KAITO has the following key differentiations compared to most of the mainstream model deployment methodologies built on top of virtual machine infrastructures:
+
+- **Container-based Model Management**: Manage large model files using container images with an OpenAI-compatible server for inference calls
+- **Preset Configurations**: Avoid adjusting workload parameters based on GPU hardware with built-in configurations
+- **Multiple Runtime Support**: Support for popular inference runtimes including [vLLM](https://github.com/vllm-project/vllm) and [transformers](https://github.com/huggingface/transformers)
+- **Auto-provisioning**: Automatically provision GPU nodes based on model requirements
+- **Public Registry**: Host large model images in the public Microsoft Container Registry (MCR) when licenses allow
+
+Using KAITO, the workflow of onboarding large AI inference models in Kubernetes is largely simplified.
+
+## Architecture
+
+KAITO follows the classic Kubernetes Custom Resource Definition(CRD)/controller design pattern. Users manage a `workspace` custom resource which describes the GPU requirements and the inference or tuning specification. KAITO controllers automate the deployment by reconciling the `workspace` custom resource.
+
+![KAITO architecture](/img/arch.png)
+
+The above figure presents the KAITO architecture overview. Its major components consist of:
+
+- **Workspace controller**: Reconciles the `workspace` custom resource, creates `machine` custom resources to trigger node auto provisioning, and creates the inference or tuning workload (`deployment`, `statefulset` or `job`) based on the model preset configurations.
+- **Node provisioner controller**: The controller's name is *gpu-provisioner* in [gpu-provisioner helm chart](https://github.com/Azure/gpu-provisioner/tree/main/charts/gpu-provisioner). It uses the `machine` CRD originated from [Karpenter](https://sigs.k8s.io/karpenter) to interact with the workspace controller. It integrates with Azure Resource Manager REST APIs to add new GPU nodes to the AKS or AKS Arc cluster.
+
+:::note
+The [*gpu-provisioner*](https://github.com/Azure/gpu-provisioner) is an open sourced component. It can be replaced by other controllers if they support [Karpenter-core](https://sigs.k8s.io/karpenter) APIs.
+:::
+
+**NEW!** Starting with version v0.5.0, KAITO releases a new operator, **RAGEngine**, which is used to streamline the process of managing a Retrieval Augmented Generation(RAG) service.
+![KAITO RAGEngine architecture](/img/ragarch.png)
+
+
+As illustrated in the above figure, the **RAGEngine controller** reconciles the `ragengine` custom resource and creates a `RAGService` deployment. The `RAGService` provides the following capabilities:
+  - **Orchestration**: use [LlamaIndex](https://github.com/run-llama/llama_index) orchestrator.
+  - **Embedding**: support both local and remote embedding services, to embed queries and documents in the vector database.
+  - **Vector database**: support a built-in [faiss](https://github.com/facebookresearch/faiss) in-memory vector database. Remote vector database support will be added soon.
+  - **Backend inference**: support any OAI compatible inference service.
+
+The details of the service APIs can be found in this [document](./rag.md).
+
+
+## Getting Started
+
+👉 To get started, please see the [Workspace Installation Guide](installation) and the [RAGEngine Installation Guide](./rag.md)!
+
+👉 For a quick start tutorial, check out [Quick Start](quick-start)!
+
+## Community
+
+- **GitHub**: [kaito-project/kaito](https://github.com/kaito-project/kaito)
+- **Slack**: [Join our community](https://join.slack.com/t/kaito-z6a6575/shared_invite/zt-37gh89vw7-odHfqmPRc5oRnDG99SBJNA)
+- **Email**: [kaito-dev@microsoft.com](mailto:kaito-dev@microsoft.com)
diff --git a/website/versioned_docs/version-v0.6.x/kaito-on-byo-gpu-nodes.md b/website/versioned_docs/version-v0.6.x/kaito-on-byo-gpu-nodes.md
new file mode 100644
index 0000000..054871a
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/kaito-on-byo-gpu-nodes.md
@@ -0,0 +1,256 @@
+---
+title: Bring Your Own GPU Nodes
+---
+
+This guide walks you through deploying KAITO on a Kubernetes cluster with self-provisioned GPU nodes.
+
+## Prerequisites
+
+If you are following the demo as is then you would need access to an Azure account.
+
+### Tools
+
+- [Azure CLI](https://learn.microsoft.com/en-us/cli/azure/install-azure-cli)
+- [kubectl](https://kubernetes.io/docs/tasks/tools/#kubectl)
+- [Helm](https://helm.sh/docs/intro/install/)
+- [jq](https://jqlang.github.io/jq/download/)
+
+## Set up a Kubernetes cluster with GPU nodes
+
+:::note
+If you already have a Kubernetes cluster, you can skip this section.
+:::
+
+For the sake of this guide, we will create an Azure Kubernetes Service (AKS) cluster.
+
+### Environment variables
+
+Make necessary changes to the following environment variables and copy paste them into your terminal:
+
+```bash
+export LOCATION="southcentralus"
+export RESOURCE_GROUP="kaito-rg"
+export AKS_RG="${RESOURCE_GROUP}-aks"
+export CLUSTER_NAME="kaito"
+export AKS_WORKER_USER_NAME="azuser"
+export SSH_KEY=~/.ssh/id_rsa.pub
+export GPU_NODE_SIZE="Standard_NC24ads_A100_v4"
+export GPU_NODE_COUNT=1
+export GPU_NODE_POOL_NAME="gpunodes"
+```
+
+### Create a resource group
+
+Run the following command to create a resource group:
+
+```bash
+az group create \
+    --name "${RESOURCE_GROUP}" \
+    --location "${LOCATION}"
+```
+
+### Create an Azure Kubernetes Service (AKS) cluster
+
+Run the following command to create an AKS cluster:
+
+```bash
+az aks create \
+    --resource-group "${RESOURCE_GROUP}" \
+    --node-resource-group "${AKS_RG}" \
+    --name "${CLUSTER_NAME}" \
+    --enable-oidc-issuer \
+    --enable-workload-identity \
+    --enable-managed-identity \
+    --node-count 1 \
+    --location "${LOCATION}" \
+    --ssh-key-value "${SSH_KEY}" \
+    --admin-username "${AKS_WORKER_USER_NAME}" \
+    --os-sku Ubuntu
+```
+
+### Add GPU nodes
+
+Your GPU nodes should be prepared for workload deployment by installing the GPU driver and k8s device plugin, specific to your NVIDIA or AMD GPU. The [NVIDIA GPU Operator](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/overview.html) and [AMD GPU Operator](https://instinct.docs.amd.com/projects/gpu-operator/en/latest/overview.html) are useful to automate the management and installation of both GPU software components. In the below example, the `gpu-driver` API field is set to `none` at create time of the GPU node pool on Azure Kubernetes Service (AKS) to allow for installation of the NVIDIA GPU Operator:
+
+:::warning
+On AKS GPU-enabled node pools, the `gpu-driver` API field should be set to `none`, to avoid duplicate installation of the GPU driver and/or unexpected conflicts on your GPU nodes. You can learn more [here](https://learn.microsoft.com/azure/aks/gpu-cluster?tabs=add-ubuntu-gpu-node-pool#skip-gpu-driver-installation).
+:::
+
+Run the following command to add GPU node to the AKS cluster:
+
+```bash
+az aks nodepool add \
+    --name "${GPU_NODE_POOL_NAME}" \
+    --resource-group "${RESOURCE_GROUP}" \
+    --cluster-name "${CLUSTER_NAME}" \
+    --node-count "${GPU_NODE_COUNT}" \
+    --node-vm-size "${GPU_NODE_SIZE}" \
+    --gpu-driver none
+```
+
+### Download kubeconfig
+
+Run the following command to download the kubeconfig file:
+
+```bash
+az aks get-credentials \
+    --resource-group "${RESOURCE_GROUP}" \
+    --name "${CLUSTER_NAME}"
+```
+
+## Prepare the Kubernetes cluster for GPU workloads
+
+### Install the NVIDIA GPU operator
+
+:::note
+If you have already set up your Kubernetes cluster with the NVIDIA GPU operator, you can skip the following steps for installation.
+:::
+
+Run the following commands to create a namespace for the GPU operator:
+
+```bash
+kubectl create ns gpu-operator
+kubectl label --overwrite ns gpu-operator pod-security.kubernetes.io/enforce=privileged
+```
+
+Run the following commands to install the GPU operator:
+
+```bash
+helm repo add nvidia https://helm.ngc.nvidia.com/nvidia
+helm repo update
+
+helm install \
+    --wait \
+    --generate-name \
+    -n gpu-operator \
+    --create-namespace \
+    nvidia/gpu-operator
+```
+
+Ensure that the GPU operator is installed by running the following command:
+
+```bash
+kubectl -n gpu-operator wait pod \
+    --for=condition=Ready \
+    -l app.kubernetes.io/component=gpu-operator \
+    --timeout=300s
+```
+
+Finally ensure that the `nvidia` runtime class is created by running the following command:
+
+```bash
+kubectl get runtimeclass nvidia
+```
+
+A typical output would look like this:
+
+```bash
+$ kubectl get runtimeclass nvidia
+NAME     HANDLER   AGE
+nvidia   nvidia    16m
+```
+
+### Label the GPU nodes
+
+We need to label the GPU nodes `apps=gpu`, so that the KAITO workspace controller can schedule the inference workloads on these nodes. If you are following along the guide, you can run the following command to label the GPU nodes:
+
+```bash
+kubectl get nodes \
+    -l agentpool="${GPU_NODE_POOL_NAME}" \
+    -o name | \
+    xargs -I {} \
+    kubectl label --overwrite {} apps=gpu
+```
+
+:::tip
+If you have used a different set up to create the GPU nodes, you can label the nodes manually by running the following command: `kubectl label node <node-name> apps=gpu`.
+:::
+
+## Install KAITO on the Kubernetes cluster
+
+Run the following command to install KAITO:
+
+```bash
+helm install workspace \
+    ./charts/kaito/workspace \
+    --namespace kaito-workspace \
+    --create-namespace
+```
+
+Ensure that kaito is installed by running the following command:
+
+```bash
+kubectl -n kaito-workspace wait pod \
+    --for=condition=Ready \
+    -l app.kubernetes.io/instance=workspace \
+    --timeout=300s
+```
+
+## Deploying a model
+
+### Deploy a workspace with a GPU model
+
+To deploy a workspace with a GPU model, run the following command:
+
+```yaml
+cat <<EOF | kubectl apply -f -
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-falcon-7b
+resource:
+  instanceType: "${GPU_NODE_SIZE}"
+  labelSelector:
+    matchLabels:
+      apps: gpu
+inference:
+  preset:
+    name: "falcon-7b"
+EOF
+```
+
+:::note
+In the above configuration you can see we have used a node labelSelector value as `apps: gpu`, this is the same label we have applied when we added the GPU node pool earlier.
+:::
+
+Ensure that the workspace is ready by running the following command:
+
+```bash
+kubectl get workspace workspace-falcon-7b
+```
+
+A typical output would look like this:
+
+```bash
+$ kubectl get workspace workspace-falcon-7b
+NAME                  INSTANCE                   RESOURCEREADY   INFERENCEREADY   JOBSTARTED   WORKSPACESUCCEEDED   AGE
+workspace-falcon-7b   Standard_NC24ads_A100_v4   True            True                          True                 16m
+```
+
+### Use the workspace
+
+Run the following command to find the cluster IP to send the request to:
+
+```bash
+export CLUSTERIP=$(kubectl get \
+    svc workspace-falcon-7b \
+    -o jsonpath="{.spec.clusterIPs[0]}")
+```
+
+Let's send a request to the workspace to get an inference response. Modify the prompt as you see fit:
+
+```bash
+export QUESTION="What are LLMs?"
+```
+
+Run the following command to send the request:
+
+```bash
+kubectl run -it --rm \
+    --restart=Never \
+    curl --image=curlimages/curl \
+    -- curl -X POST http://$CLUSTERIP/chat \
+    -H "accept: application/json" \
+    -H "Content-Type: application/json" \
+    -d "{\"prompt\":\"${QUESTION}\"}" | jq
+```
diff --git a/website/versioned_docs/version-v0.6.x/kaito-oom-prevention.md b/website/versioned_docs/version-v0.6.x/kaito-oom-prevention.md
new file mode 100644
index 0000000..c2e6237
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/kaito-oom-prevention.md
@@ -0,0 +1,49 @@
+---
+title: OOM Prevention
+---
+
+Large Language Model (LLM) inference typically requires significant GPU memory. Unlike traditional virtual memory management in operating systems, CUDA does not support GPU memory swapping. This limitation often leads to CUDA Out-Of-Memory (OOM) errors when GPU memory is insufficient for inference requests. Addressing GPU OOM issues is one of the most challenging aspects of LLM deployment and tuning.
+
+This document explains the mechanisms KAITO implements to prevent CUDA OOM errors in inference workspaces. Since KAITO uses vLLM as its default inference engine, most of the discussion focuses on vLLM use cases.
+
+## Background
+
+The following diagram illustrates the GPU memory layout when serving inference requests using the vLLM engine:
+
+![GPU memory layout](/img/gpumemlayout.png)
+
+The GPU memory is divided into three main components:
+
+1. **Model Parameters**: These are the model weights (loaded from weight files, with or without quantization) that remain in GPU memory throughout the server's lifecycle.
+
+2. **vLLM Managed KV Cache**: The vLLM cache manager attempts to allocate all remaining GPU memory for KV cache blocks to maximize memory utilization. To prevent the cache manager from allocating too much memory, vLLM provides a `gpu-memory-utilization` configuration parameter. This parameter specifies the fraction of GPU memory to be used for the model inference, ranging from 0 to 1. For example, a value of 0.9 reserves 10% of GPU memory for other operations.
+
+3. **Other Memory**: A small portion of memory is reserved for ephemeral operations such as model activation.
+
+The KV cache usage is dynamic and roughly proportional to the context length (the total length of input and output tokens). vLLM requires the KV cache to be sufficient for the full context when serving a request. To enforce this constraint, vLLM uses a `max-model-len` parameter, which defaults to the model's maximum context length (e.g., 128K for Phi-3 models) if not specified. During startup, vLLM performs an internal probing test using a fake request with the specified `max-model-len`. If the configured KV cache cannot handle this request, the vLLM engine fails to start. After startup, any inference request exceeding the `max-model-len` will fail.
+
+## OOM Scenarios and KAITO's Solutions
+
+Based on the above understanding, we identify and address the following potential OOM scenarios:
+
+### Scenario 1: Insufficient GPU Memory for Model Parameters
+
+Loading model parameters into GPU memory is a fundamental requirement for any LLM engine. KAITO records the required memory size for each supported model based on 16-bit floating-point precision, as documented in the model's HuggingFace documentation. The validation webhook checks if the specified GPU SKU has sufficient aggregated memory to meet these requirements. However, this check is skipped for SKUs not in KAITO's built-in SKU list.
+
+> **Note**: KAITO allows users to modify the default vLLM configuration through a custom ConfigMap, which may include model quantization options. In such cases, the webhook's resource check might be overly restrictive because the actual memory footprint will be smaller if 4-bit or 8-bit quantization is enabled. Users can bypass this check by adding the annotation `kaito.sh/bypass-resource-checks: "true"` to the workspace custom resource when using quantized models on GPUs with limited memory.
+
+### Scenario 2: Insufficient Memory for Other Operations
+
+Setting a high `gpu-memory-utilization` value maximizes GPU memory utilization but can lead to OOM errors during operations like model activation if insufficient memory remains. Since different GPU SKUs have different GPU memory sizes, using a fixed percentage threshold for `gpu-memory-utilization` is suboptimal. KAITO dynamically adjusts the `gpu-memory-utilization` based on the GPU SKU's memory size, ensuring a smaller threshold for GPU SKUs with less than 20GB memory and a larger threshold (capped at 95%) for GPU SKUs with larger memory.
+
+### Scenario 3: KV Cache Size Insufficient for `max-model-len`
+
+This issue commonly occurs with models supporting large context windows (e.g., 128K). When users set a high `max-model-len` or use models with large default context windows, vLLM's internal probing test may fail on GPUs with limited memory. Typically, finding the optimal `max-model-len` to maximize KV cache utilization without causing OOM errors requires tedious manual searches.
+
+To simplify this process, KAITO implements an external probing test using a binary search algorithm to find a near-optimal `max-model-len` based on available GPU memory. While this test slightly increases service startup time, it eliminates the need for manual tuning. Users can find the calculated `max-model-len` in the inference service pod logs.
+
+> **Note**: The external probing test currently only supports SKUs with a single GPU instance. For SKUs with multiple GPU instances and small per-instance memory (e.g., `Standard_NC12s_v3` with two 16GB v100 instances), users must manually specify the `max-model-len` in the workspace's ConfigMap. The validation webhook enforces this requirement.
+
+## Generic Memory Usage Reduction
+
+Benchmark tests demonstrate that enabling `expandable_segments=true` in PYTORCH_CUDA_ALLOC_CONF helps reduce memory fragmentation and peak memory consumption, particularly for tuning jobs. KAITO enables this flag by default in workspace pods.
diff --git a/website/versioned_docs/version-v0.6.x/model-as-oci-artifacts.md b/website/versioned_docs/version-v0.6.x/model-as-oci-artifacts.md
new file mode 100644
index 0000000..73a3b26
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/model-as-oci-artifacts.md
@@ -0,0 +1,189 @@
+---
+title: Model As OCI Artifacts
+description: Efficiently distribute Large Language Models using Open Container Initiative (OCI) Artifacts
+---
+
+# Model As OCI Artifacts
+
+KAITO efficiently distributes Large Language Models (LLMs) by leveraging Open Container Initiative (OCI) Artifacts to package and deploy model weights separately from the inference runtime. This architecture provides superior performance, scalability, and maintainability for large language model deployments.
+
+## Overview
+
+KAITO uses a split architecture approach for model distribution:
+
+- **Base Images**: Contain the inference runtime and dependencies
+- **OCI Artifacts**: Contain the model weights and configuration files stored separately in OCI registries
+
+This design choice enables KAITO to optimize both build and deployment performance while maintaining compatibility with standard container infrastructure. The model weights are downloaded as OCI artifacts during pod initialization, providing faster deployment times and more efficient resource utilization compared to traditional monolithic container images.
+
+## Design Rationale
+
+KAITO chose the OCI Artifacts approach to address fundamental challenges in large language model distribution and deployment:
+
+### Build Efficiency
+
+By separating model weights from base images, KAITO eliminates redundant rebuilds:
+
+- **Reduced Build Time**: Model images only need to be built once and can be reused across base image updates
+- **Efficient Build Context**: Large model files are not included in Docker build context, avoiding unnecessary transfers
+- **Resource Optimization**: Build processes for large models like Falcon-40B are reduced from nearly 2 hours to minutes
+
+### Deployment Performance
+
+KAITO's OCI Artifacts approach optimizes the deployment pipeline:
+
+- **Concurrent Downloads**: Model weights can be downloaded in parallel, improving bandwidth utilization
+- **Advanced Compression**: Zstd compression provides better decompression performance than traditional gzip
+- **Optimized Unpacking**: Separate artifact handling reduces serialization bottlenecks during pod startup
+
+#### Understanding Container Performance
+
+KAITO's design optimizes all phases of the container pull process:
+
+1. **Downloading** layer data from the registry
+2. **Decompressing** the layer if necessary  
+3. **Checking** sha256 digest
+4. **Unpacking** the layer, applying changes to the snapshot
+
+KAITO addresses these performance characteristics by optimizing not just the download phase (30% of total time) but also the decompression, validation, and unpacking phases through its OCI Artifacts approach.
+
+## How KAITO Uses OCI Artifacts
+
+### 1. Efficient Build Process
+
+KAITO uses [ORAS](https://github.com/oras-project/oras) to package model weights and configuration files directly into OCI artifacts without requiring large build contexts. This approach achieves the same distribution goals as traditional Docker builds but with significantly improved efficiency.
+
+### 2. Optimized Compression
+
+KAITO employs zstd compression instead of gzip, providing superior decompression performance specifically optimized for large model weight files.
+
+### 3. Split Architecture Design
+
+KAITO implements a two-component architecture:
+
+- **Base Image**: Contains the inference runtime and dependencies
+- **OCI Artifacts**: Contains the model weights and configuration files
+
+This architectural separation enables model images to be built once and reused across multiple base image updates, significantly reducing maintenance overhead.
+
+## When to Use OCI Artifacts
+
+KAITO automatically uses OCI Artifacts for supported models, providing benefits in several scenarios:
+
+### Recommended Use Cases
+
+- **Large Language Models**: Models with weights exceeding 1GB benefit significantly from the optimized distribution
+- **Frequent Updates**: Environments with regular base image updates for security patches or feature additions
+- **Multi-Model Deployments**: When deploying multiple model variants that share similar runtime requirements
+- **Resource-Constrained Environments**: Clusters where build time and bandwidth optimization are critical
+
+### Performance Considerations
+
+The OCI Artifacts approach provides the most benefit when:
+- Model weights are large relative to the runtime container
+- Network bandwidth is limited or expensive
+- Build infrastructure resources are constrained
+- Deployment frequency is high
+
+## Architecture
+
+### Model Weights Download Process
+
+The system uses an [initContainer](https://kubernetes.io/docs/concepts/workloads/pods/init-containers/) to download model weights as OCI artifacts using ORAS:
+
+```mermaid
+sequenceDiagram
+    participant CR as Container Runtime
+    participant IC as InitContainer
+    participant PV as Persistent Volume
+    participant IS as Inference Server
+
+    CR->>CR: Pull Base Image
+    CR->>IC: Start InitContainer
+    IC->>IC: Pull Model Weights as OCI Artifacts using ORAS
+    IC->>PV: Store Model Weights
+    IC->>CR: Complete
+    CR->>IS: Start Main Inference Container
+    IS->>PV: Load Model Weights
+    IS->>IS: Inference Server Ready
+```
+
+
+
+### OCI Artifacts vs OCI Images
+
+The Open Container Initiative (OCI) defines specifications and standards for container technologies, including the OCI Distribution Specification. OCI Image Manifests have a required field `config.mediaType` that differentiates between various types of artifacts.
+
+- **OCI Image**: A subset of OCI artifacts, accepting only specific mediatypes
+- **OCI Artifacts**: More general format that can contain various types of content including model weights
+
+ORAS (OCI Registry As Storage) is the tool used for managing OCI artifacts, including pushing, pulling, and handling metadata in OCI registries.
+
+### Alternative Approaches for Large Models
+
+For extremely large models that may not be practical to package as OCI artifacts due to size constraints, KAITO also supports:
+
+- **Direct Download from Hugging Face**: Models can be downloaded directly from Hugging Face Hub during pod initialization for models that exceed registry limitations
+- **External Model Caching**: Integration with external model caching solutions for improved performance with very large models
+
+## Compatibility
+
+### Container Runtimes
+
+- **CRI-O**: Supports general OCI artifacts
+- **Containerd**: Limited support for OCI artifacts (see [containerd issue](https://github.com/containerd/containerd/issues/11381#issuecomment-2917050414))
+
+### OCI Registries
+
+Most OCI registries support OCI artifacts. For a complete list of compatible registries, see the [ORAS compatibility documentation](https://oras.land/docs/compatible_oci_registries#registries-supporting-oci-artifacts).
+
+## Benefits
+
+### Performance Improvements
+
+- **Reduced Build Time**: Model images only need to be built once
+- **Faster Pulls**: Improved download concurrency and compression
+- **Better Resource Usage**: Optimized bandwidth utilization
+
+### Operational Benefits
+
+- **Simplified Maintenance**: Base image updates don't require rebuilding all model images
+- **Storage Efficiency**: Better compression ratios with zstd
+- **Scalability**: More efficient handling of large model files
+
+## Performance Results
+
+Testing on Standard_NC24s_v3 with phi4 model shows significant improvements:
+
+![Model as OCI Artifacts Evaluation](/img/model-as-oci-artifacts-evaluation.png)
+
+The evaluation compared different configurations:
+
+| Configuration | Description | Benefits |
+|---------------|-------------|----------|
+| Baseline (single-layer-tar-gz) | Current approach with all files in one layer | Simple, self-contained |
+| OCI Artifacts | Base image + separate model artifacts | Reduced build time, better performance |
+
+## Getting Started
+
+### Automatic Usage
+
+KAITO automatically uses OCI Artifacts for supported models when the feature is available. No additional configuration is required - simply deploy your model using standard KAITO workflows and the system will optimize the distribution automatically.
+
+### Verifying OCI Artifacts Usage
+
+To confirm that your model is using OCI Artifacts:
+
+1. Check the pod specification for initContainers that download model artifacts
+2. Monitor pod startup times for improved performance compared to traditional container pulls
+3. Review registry storage to see separate base images and model artifact entries
+
+### Troubleshooting
+
+If you encounter issues with OCI Artifacts:
+
+- Ensure your container runtime supports OCI Artifacts (see Compatibility section)
+- Verify registry compatibility with OCI Artifacts specification
+- Check initContainer logs for artifact download status
+- Refer to the [technical specification](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20250609-model-as-oci-artifacts.md) for advanced configuration options
+
diff --git a/website/versioned_docs/version-v0.6.x/monitoring.md b/website/versioned_docs/version-v0.6.x/monitoring.md
new file mode 100644
index 0000000..2fc88ca
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/monitoring.md
@@ -0,0 +1,50 @@
+---
+title: Monitoring
+---
+
+Inference service exposes Prometheus metrics for monitoring system stats.
+
+## vLLM runtime
+
+vLLM exposes Prometheus metrics at the `/metrics` endpoint. These metrics provide detailed insights into the system's performance, resource utilization, and request processing statistics. The following table lists all available metrics:
+
+vLLM version: 0.8.2
+
+:::note
+vLLM V1 engine is now enabled by default. Check [here](https://docs.vllm.ai/en/latest/getting_started/v1_user_guide.html) for details.
+:::
+
+| Category | Metric Name | Type | Description |
+|----------|-------------|------|-------------|
+| System Stats - Scheduler | `vllm:num_requests_running` | Gauge | Number of requests currently running on GPU |
+| System Stats - Scheduler | `vllm:num_requests_waiting` | Gauge | Number of requests waiting to be processed |
+| System Stats - Scheduler | `vllm:lora_requests_info`   | Gauge | Running stats on lora requests |
+| System Stats - Scheduler | `vllm:num_requests_swapped` | Gauge | Number of requests swapped to CPU (DEPRECATED: KV cache offloading is not used in V1) |
+| System Stats - Cache | `vllm:gpu_cache_usage_perc` | Gauge | GPU KV-cache usage. 1 means 100 percent usage |
+| System Stats - Cache | `vllm:cpu_cache_usage_perc` | Gauge | CPU KV-cache usage. 1 means 100 percent usage (DEPRECATED: KV cache offloading is not used in V1) |
+| System Stats - Cache | `vllm:cpu_prefix_cache_hit_rate` | Gauge | CPU prefix cache block hit rate (DEPRECATED: KV cache offloading is not used in V1) |
+| System Stats - Cache | `vllm:gpu_prefix_cache_hit_rate` | Gauge | GPU prefix cache block hit rate (DEPRECATED: use `vllm:gpu_prefix_cache_queries` and `vllm:gpu_prefix_cache_hits in V1`) |
+| System Stats - Cache | `vllm:gpu_cache_usage_perc` | Counter | GPU prefix cache queries, in terms of number of queried blocks (V1 only) |
+| System Stats - Cache | `vllm:gpu_prefix_cache_hits` | Counter | GPU prefix cache hits, in terms of number of cached blocks (V1 only) |
+| Iteration Stats | `vllm:num_preemptions_total` | Counter | Cumulative number of preemption from the engine |
+| Iteration Stats | `vllm:prompt_tokens_total` | Counter | Number of prefill tokens processed |
+| Iteration Stats | `vllm:generation_tokens_total` | Counter | Number of generation tokens processed |
+| Iteration Stats | `vllm:iteration_tokens_total`  | Histogram | Number of tokens per engine_step |
+| Iteration Stats | `vllm:time_to_first_token_seconds` | Histogram | Time to first token in seconds |
+| Iteration Stats | `vllm:time_per_output_token_seconds` | Histogram | Time per output token in seconds |
+| Request Stats | `vllm:e2e_request_latency_seconds` | Histogram | End to end request latency in seconds |
+| Request Stats | `vllm:request_queue_time_seconds` | Histogram | Time spent in WAITING phase for request |
+| Request Stats | `vllm:request_inference_time_seconds` | Histogram | Time spent in RUNNING phase for request |
+| Request Stats | `vllm:request_prefill_time_seconds` | Histogram | Time spent in PREFILL phase for request |
+| Request Stats | `vllm:request_decode_time_seconds` | Histogram | Time spent in DECODE phase for request |
+| Request Stats | `vllm:request_prompt_tokens` | Histogram | Number of prefill tokens processed per request |
+| Request Stats | `vllm:request_generation_tokens` | Histogram | Number of generation tokens processed per request |
+| Request Stats | `vllm:request_max_num_generation_tokens` | Histogram | Maximum number of requested generation tokens |
+| Request Stats | `vllm:request_params_n` | Histogram | The 'n' request parameter |
+| Request Stats | `vllm:request_params_max_tokens` | Histogram | The 'max_tokens' request parameter |
+| Request Stats | `vllm:request_success_total` | Counter | Count of successfully processed requests |
+| Speculative Decoding | `vllm:spec_decode_draft_acceptance_rate` | Gauge | Speculative token acceptance rate (DEPRECATED: Unused in V1) |
+| Speculative Decoding | `vllm:spec_decode_efficiency` | Gauge | Speculative decoding system efficiency (DEPRECATED: Unused in V1) |
+| Speculative Decoding | `vllm:spec_decode_num_accepted_tokens_total` | Counter | Number of accepted tokens |
+| Speculative Decoding | `vllm:spec_decode_num_draft_tokens_total` | Counter | Number of draft tokens |
+| Speculative Decoding | `vllm:spec_decode_num_emitted_tokens_total` | Counter | Number of emitted tokens (DEPRECATED: Unused in V1) |
diff --git a/website/versioned_docs/version-v0.6.x/multi-node-inference.md b/website/versioned_docs/version-v0.6.x/multi-node-inference.md
new file mode 100644
index 0000000..9c40caf
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/multi-node-inference.md
@@ -0,0 +1,264 @@
+---
+title: Multi-Node Inference
+---
+
+This document explains how to configure and use multi-node distributed inference in KAITO for large models that require more GPU resources than a single node can provide.
+
+## Overview
+
+Multi-node inference allows you to deploy large AI models across multiple nodes (servers) when the model is too large to fit on a single node. KAITO supports different parallelism strategies depending on your model's requirements:
+
+| Strategy | Use Case | Supported |
+|----------|----------|-----------|
+| **Single GPU** | Small models that fit on one GPU | ✅ |
+| **Single-Node Multi-GPU** | Models that need multiple GPUs but fit on one node | ✅ |
+| **Multi-Node Multi-GPU** | Very large models (400B+ parameters) requiring multiple nodes | ✅ |
+
+## When to Use Multi-Node Inference
+
+Consider multi-node inference when:
+
+- Your model has 400B+ parameters and cannot fit on a single node
+- You need to serve models like very large language models that exceed single-node memory capacity
+- You have specific performance requirements that benefit from distributed processing
+
+:::note
+Multi-node inference introduces additional complexity and network overhead. Only use it when your model truly requires more resources than a single node can provide.
+:::
+
+## Supported Models
+
+The following preset models support multi-node distributed inference:
+
+- **Llama3**: `llama-3.3-70b-instruct`
+- **DeepSeek**: `deepseek-r1-0528`, `deepseek-v3-0324` 
+
+Check the [presets documentation](./presets.md) for the complete list and their specific requirements.
+
+## Configuration
+
+### Basic Multi-Node Setup
+
+To deploy a model across multiple nodes, specify the node count in your Workspace configuration:
+
+```yaml
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-large-model
+resource:
+  count: 2                    # Number of nodes to use
+  instanceType: "Standard_NC80adis_H100_v5"
+  labelSelector:
+    matchLabels:
+      apps: large-model
+inference:
+  preset:
+    name: "llama-3.3-70b-instruct"
+```
+
+### Pre-Provisioned Nodes
+
+If you're using pre-provisioned GPU nodes, specify them explicitly:
+
+```yaml
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-large-model
+resource:
+  count: 2
+  preferredNodes:
+    - gpu-node-1
+    - gpu-node-2
+  labelSelector:
+    matchLabels:
+      apps: large-model
+inference:
+  preset:
+    name: "llama-3.3-70b-instruct"
+```
+
+:::warning
+Pre-provisioned nodes must have the same matching labels as specified in the `resource` spec, and each node must report available GPU resources.
+:::
+
+### Custom vLLM Parameters for Multi-Node
+
+You can customize vLLM runtime parameters for distributed inference using a ConfigMap:
+
+```yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: distributed-inference-config
+data:
+  inference_config.yaml: |
+    vllm:
+      gpu-memory-utilization: 0.95
+      max-model-len: 131072
+---
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-large-model
+resource:
+  count: 2
+  instanceType: "Standard_NC80adis_H100_v5"
+  labelSelector:
+    matchLabels:
+      apps: large-model
+inference:
+  preset:
+    name: "llama-3.3-70b-instruct"
+  config: "distributed-inference-config"
+```
+
+Key parameters for multi-node inference:
+- `tensor-parallel-size`: Automatically set by KAITO based on the number of GPUs per node
+- `pipeline-parallel-size`: Automatically set by KAITO based on the number of nodes
+- `gpu-memory-utilization`: Fraction of GPU memory to use (0.0-1.0) - user configurable
+- `max-model-len`: Maximum sequence length - user configurable
+
+:::note
+The `tensor-parallel-size` and `pipeline-parallel-size` parameters are automatically managed by KAITO based on your cluster configuration and do not need to be specified in the ConfigMap.
+:::
+
+## Architecture
+
+### Single-Node Multi-GPU
+
+When using multiple GPUs on a single node, KAITO uses **tensor parallelism** to split the model across GPUs within that node:
+
+```
+┌─────────────────────────────────────┐
+│             Node 1                  │
+│  ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐    │
+│  │GPU 1│ │GPU 2│ │GPU 3│ │GPU 4│    │
+│  └─────┘ └─────┘ └─────┘ └─────┘    │
+│           Model Split               │
+└─────────────────────────────────────┘
+```
+
+### Multi-Node Multi-GPU
+
+For multi-node deployments, KAITO combines **pipeline parallelism** between nodes and **tensor parallelism** within each node:
+
+```
+┌─────────────────────┐    ┌─────────────────────┐
+│       Node 1        │    │       Node 2        │
+│  ┌─────┐┌─────┐     │    │  ┌─────┐ ┌─────┐    │
+│  │GPU 1││GPU 2│     │◄──►│  │GPU 3│ │GPU 4│    │
+│  └─────┘└─────┘     │    │  └─────┘ └─────┘    │
+│   Layers 1-N/2      │    │   Layers N/2+1-N    │
+└─────────────────────┘    └─────────────────────┘
+```
+
+## Resource Validation
+
+KAITO automatically validates that your configuration provides sufficient resources:
+
+- **GPU Count**: `(GPUs per instance) × (workspace.resource.count) ≥ (Required GPUs for model)`
+- **Memory**: `(GPU memory) × (Total GPUs) ≥ (Required model memory)`
+
+If validation fails, you'll receive an error message when creating or updating the workspace.
+
+:::tip Resource Optimization
+KAITO may use fewer nodes than specified in `workspace.resource.count` if the model can fit efficiently on fewer nodes. This optimizes GPU utilization and reduces network overhead, but be mindful of the costs when provisioning many nodes.
+:::
+
+## Service Architecture
+
+Multi-node inference uses Kubernetes StatefulSets to ensure stable pod identity and coordination:
+
+- **Leader Pod** (index 0): Coordinates the distributed inference and serves the API
+- **Worker Pods** (index 1+): Join the Ray cluster and participate in model serving
+
+The service endpoint points to the leader pod, which handles all incoming requests and coordinates with worker pods.
+
+:::note Future Enhancement
+KAITO will support [LeaderWorkerSet](https://lws.sigs.k8s.io/docs/overview/) in the future to provide better management of leader-worker topologies and improved fault tolerance for multi-node deployments.
+:::
+
+## Health Monitoring
+
+Multi-node deployments use specialized health checks:
+
+- **Liveness Probes**: Monitor Ray cluster health and detect dead actors
+- **Readiness Probes**: Check service availability via the leader pod's `/health` endpoint
+
+If worker pods fail, the leader will restart to reinitialize the entire cluster, ensuring all pods are synchronized.
+
+## Best Practices
+
+1. **Resource Planning**: Carefully plan your GPU and memory requirements before deployment
+2. **Network Bandwidth**: Ensure sufficient network bandwidth between nodes for optimal performance
+3. **Monitoring**: Monitor both individual node health and overall cluster performance
+4. **Cost Management**: Be aware that multi-node deployments can be expensive; only use when necessary
+
+## Troubleshooting
+
+### Service Unavailable After Deployment
+
+If the service becomes unavailable:
+
+1. Check if all pods are running: `kubectl get pods -l app=<your-app-label>`
+2. Verify Ray cluster health in leader pod logs
+3. Ensure network connectivity between nodes
+4. Check resource allocation and GPU availability
+
+### Worker Pod Failures
+
+Worker pod failures will trigger leader pod restart to reinitialize the cluster:
+
+1. Monitor pod restart events
+2. Check for resource constraints (memory, GPU)
+3. Verify node-to-node network connectivity
+4. Review pod logs for Ray cluster connection issues
+
+### Performance Issues
+
+If you experience poor performance:
+
+1. Monitor network latency between nodes
+2. Check GPU utilization across all nodes
+3. Review memory usage and potential bottlenecks
+4. Consider adjusting parallelism parameters
+
+## API Usage
+
+Multi-node inference services expose the same API as single-node deployments. The vLLM runtime provides OpenAI-compatible endpoints:
+
+```bash
+# Get the cluster IP of your service
+kubectl get services
+
+# Check service health
+kubectl run -it --rm --restart=Never curl --image=curlimages/curl -- \
+  curl http://<CLUSTER-IP>:80/health
+
+# Generate text using chat completions API
+kubectl run -it --rm --restart=Never curl --image=curlimages/curl -- \
+  curl -X POST http://<CLUSTER-IP>:80/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "llama-3.3-70b-instruct",
+    "messages": [{"role": "user", "content": "Your prompt here"}],
+    "max_tokens": 100
+  }'
+```
+
+For detailed API specifications, see the [inference documentation](./inference.md#inference-api).
+
+## Limitations
+
+- **Custom Models**: Multi-node inference is currently only supported for preset models
+- **Fault Tolerance**: The system requires leader restart when worker pods fail
+- **Network Dependency**: Performance heavily depends on inter-node network quality
+- **Complexity**: Debugging and monitoring are more complex than single-node deployments
+
+## Related Documentation
+
+- [Inference](./inference.md) - General inference documentation
+- [Presets](./presets.md) - Supported models and their requirements
+- [Custom Model](./custom-model.md) - Using custom models (single-node only)
\ No newline at end of file
diff --git a/website/versioned_docs/version-v0.6.x/preset-onboarding.md b/website/versioned_docs/version-v0.6.x/preset-onboarding.md
new file mode 100644
index 0000000..7ba3054
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/preset-onboarding.md
@@ -0,0 +1,30 @@
+---
+title: Preset onboarding
+---
+
+This document describes how to add a new supported OSS model in KAITO. The process is designed to allow community users to initiate the request. KAITO maintainers will follow up and deal with managing the model images and guiding the code changes to set up the model preset configurations.
+
+## Step 1: Make a proposal
+
+This step is done by the requestor. The requestor should make a PR to describe the target OSS model following this [template](https://github.com/kaito-project/kaito/blob/main/docs/proposals/YYYYMMDD-model-template.md). The proposal status should be `provisional` in the beginning. KAITO maintainers will review the PR and decide to accept or reject the PR. The PR could be rejected if the target OSS model has low usage, or it has strict license limitations, or it is a relatively small model with limited capabilities.
+
+
+## Step 2: Validate and test the model
+
+This step is done by KAITO maintainers. Based on the information provided in the proposal, KAITO maintainers will download the model and test it using the specified runtime. The entire process is automated via GitHub actions when KAITO maintainers file a PR to add the model to the [supported\_models.yaml](https://github.com/kaito-project/kaito/blob/main/presets/workspace/models/supported_models.yaml).
+
+
+## Step 3: Push model image to MCR
+
+This step is done by KAITO maintainers. If the model license allows, KAITO maintainers will push the model image to MCR, making the image publicly available. This step is skipped if only private access is allowed for the model image. Once this step is done, KAITO maintainers will update the status of the proposal submitted in Step 1 to `ready to integrate`.
+
+## Step 4: Add preset configurations
+
+This step is done by the requestor. The requestor will work on a PR to register the model with preset configurations. The PR will contain code changes to implement a simple inference interface. [Here](https://github.com/kaito-project/kaito/blob/main/presets/workspace/models/falcon/model.go) is an existing example. In the same PR, or a separate PR, the status of the proposal status should be updated to `integrated`.
+
+## Step 5: Add an E2E test
+
+This step is done by the requestor. A new e2e test should be added to [here](https://github.com/kaito-project/kaito/blob/main/test/e2e/preset_test.go) which ensures the inference service is up and running with preset configurations.
+
+
+After all the above are done, a new model becomes available in KAITO.
diff --git a/website/versioned_docs/version-v0.6.x/presets.md b/website/versioned_docs/version-v0.6.x/presets.md
new file mode 100644
index 0000000..97c39dd
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/presets.md
@@ -0,0 +1,32 @@
+---
+title: Presets
+---
+
+The current supported model families with preset configurations are listed below.
+
+| Model Family                                | Compatible KAITO Versions |
+|---------------------------------------------|---------------------------|
+| [falcon](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/falcon)         | v0.0.1+                   |
+| [mistral](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/mistral)       | v0.2.0+                   |
+| [phi2](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/phi2)             | v0.2.0+                   |
+| [phi3](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/phi3)             | v0.3.0+                   |
+| [phi4](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/phi4)             | v0.4.5+                   |
+| [qwen7b](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/qwen)           | v0.4.1+                   |
+| [qwen32b](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/qwen)          | v0.4.5+                   |
+| [llama3](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/llama3)         | v0.4.6+                   |
+
+## Validation
+
+Each preset model has its own hardware requirements in terms of GPU count and GPU memory defined in the respective `model.go` file. KAITO controller performs a validation check of whether the specified SKU and node count are sufficient to run the model or not. In case the provided SKU is not in the known list, the controller bypasses the validation check which means users need to ensure the model can run with the provided SKU.
+
+## Distributed inference
+
+For models that support distributed inference, when the node count is larger than one, [Torch Distributed Elastic](https://pytorch.org/docs/stable/distributed.elastic.html) is configured with master/worker pods running in multiple nodes and the service endpoint is the master pod.
+
+The following preset models support multi-node distributed inference:
+
+| Model Family | Models | Multi-Node Support |
+|--------------|--------|-------------------|
+| [llama3](https://github.com/kaito-project/kaito/tree/main/presets/workspace/models/llama3) | `llama-3.3-70b-instruct` | ✅ |
+
+For detailed information on configuring and using multi-node inference, see the [Multi-Node Inference](./multi-node-inference.md) documentation.
diff --git a/website/versioned_docs/version-v0.6.x/proposals.md b/website/versioned_docs/version-v0.6.x/proposals.md
new file mode 100644
index 0000000..c521b45
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/proposals.md
@@ -0,0 +1,56 @@
+---
+title: Proposals
+---
+
+This section contains proposals for adding new models to KAITO. Each proposal describes the process of evaluating and integrating new OSS models into the KAITO ecosystem.
+
+## Proposal Template
+
+Before creating a new model proposal, please use the following template: [Model Proposal Template](https://github.com/kaito-project/kaito/blob/main/docs/proposals/YYYYMMDD-model-template.md)
+
+## Current Proposals
+
+Below are the current model proposals in various stages of integration:
+
+### Provisional Status
+- [Llama 3.3 70B Instruct](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20250529-llama-3.3-70b-instruct.md) - Meta's multilingual instruction-tuned 70B model
+- [Qwen2.5 Coder](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20250103-qwen2.5-coder.md) - Qwen2.5 series for code generation
+- [Phi-4 Instruct](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20241212-phi4-instruct.md) - Microsoft's latest Phi-4 instruction-tuned model
+- [Distributed Inference](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20250325-distributed-inference.md) - Support for distributed inference across multiple GPUs
+- [Model as OCI Artifacts](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20250609-model-as-oci-artifacts.md) - Packaging models as OCI artifacts
+
+### Integrated Status
+- [Mistral Instruct](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20240205-mistral-instruct.md) - Mistral AI's instruction-tuned model
+- [Mistral](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20240205-mistral.md) - Base Mistral model
+- [Phi-2](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20240206-phi-2.md) - Microsoft's Phi-2 small language model
+- [Phi-3 Instruct](https://github.com/kaito-project/kaito/blob/main/docs/proposals/20240527-phi3-instruct.md) - Microsoft's Phi-3 instruction-tuned model
+
+## Proposal Process
+
+For detailed information about the model onboarding process, see the [Model Onboarding Guide](./preset-onboarding.md).
+
+### Step 1: Create a Proposal
+Use the model proposal template to describe the target OSS model, including licensing, usage statistics, and technical requirements.
+
+### Step 2: Model Validation
+KAITO maintainers validate and test the proposed model using the specified runtime.
+
+### Step 3: Image Publishing
+If licensing allows, model images are published to Microsoft Container Registry (MCR).
+
+### Step 4: Integration
+Implement preset configurations and inference interfaces for the model.
+
+### Step 5: Testing
+Add comprehensive E2E tests to ensure the model works correctly with KAITO.
+
+## Contributing a Proposal
+
+To contribute a new model proposal:
+
+1. Fork the KAITO repository
+2. Copy the [model template](https://github.com/kaito-project/kaito/blob/main/docs/proposals/YYYYMMDD-model-template.md) to `docs/proposals/YYYYMMDD-<model-name>.md`
+3. Fill out all required sections
+4. Submit a pull request for review
+
+The proposal status will be updated as it progresses through the integration pipeline.
diff --git a/website/versioned_docs/version-v0.6.x/quick-start.mdx b/website/versioned_docs/version-v0.6.x/quick-start.mdx
new file mode 100644
index 0000000..978abc8
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/quick-start.mdx
@@ -0,0 +1,217 @@
+---
+title: Quick Start
+---
+
+After installing KAITO, you can quickly deploy a phi-4-mini-instruct inference service to get started.
+
+## Prerequisites
+
+- A Kubernetes cluster with KAITO installed (see [Installation](installation))
+- `kubectl` configured to access your cluster
+- **GPU nodes available in your cluster** - You have two options:
+  - **Bring your own GPU (BYO) nodes**: If you already have GPU nodes in your cluster, you can use the preferred nodes approach (shown below)
+  - **Auto-provisioning**: Set up automatic GPU node provisioning for your cloud provider (see [these steps](installation#option-2-auto-provision-gpu-nodes))
+
+## Deploy Your First Model
+
+Your model can be deployed either on your own GPU nodes, or KAITO can auto-provision nodes for it to run on.
+
+### Option 1: Bring your own GPU nodes
+
+Let's start by deploying a phi-4-mini-instruct model on your existing GPU nodes.
+
+First get the nodes.
+
+```bash
+kubectl get nodes -l accelerator=nvidia
+```
+
+The output should look similar to this, showing all your GPU nodes.
+
+```
+NAME                                  STATUS   ROLES    AGE     VERSION
+gpunp-26695285-vmss000000             Ready    <none>   2d21h   v1.31.9
+gpunp-26695285-vmss000001             Ready    <none>   2d21h   v1.31.9
+```
+
+The GPU nodes will need a label in order for a KAITO Workspace to select it. We'll use the label `apps=llm-inference` for this example. Label the nodes you want to use. 
+
+```bash
+kubectl label node gpunp-26695285-vmss000001 apps=llm-inference
+kubectl label node gpunp-26695285-vmss000000 apps=llm-inference
+```
+
+Create a YAML file named `phi-4-workspace.yaml` with the following content. The label selector here must match the labels you set on the nodes.
+
+```yaml title="phi-4-workspace.yaml"
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-phi-4-mini
+resource:
+  preferredNodes:
+    - gpunp-26695285-vmss000000
+    - gpunp-26695285-vmss000001
+  labelSelector:
+    matchLabels:
+      apps: llm-inference
+inference:
+  preset:
+    name: phi-4-mini-instruct
+```
+
+Apply your configuration to your cluster:
+
+```bash
+kubectl apply -f phi-4-workspace.yaml
+```
+
+### Option 2: Auto-provision GPU nodes
+
+The following cloud providers support auto-provisioning GPU nodes in addition to BYO nodes.
+
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+<Tabs groupId="provider">
+<TabItem value="azure" label="Azure" default>
+
+:::info
+If you have not already, follow the steps [here](azure#setup-auto-provisioning) to install the `gpu-provisioner` Helm chart.
+
+Ensure that the chart is ready and pods are running with [these steps](azure#verify-setup) otherwise the GPU nodes will not be provisioned.
+:::
+
+Create a YAML file named `phi-4-workspace.yaml` with the following content. The `instanceType` field will specify what nodes will be auto-provisioned instead of only matching existing nodes in the BYO case.
+
+```yaml title="phi-4-workspace.yaml"
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-phi-4-mini
+resource:
+  instanceType: "Standard_NC6s_v3" # Will trigger node creation
+  labelSelector:
+    matchLabels:
+      apps: phi-4-mini
+inference:
+  preset:
+    name: phi-4-mini-instruct
+```
+
+</TabItem>
+<TabItem value="aws" label="AWS" default>
+
+:::info
+If you have not already, follow the steps [here](aws#setup-auto-provisioning) to install the `Karpenter` Helm chart.
+
+Ensure that the chart is ready and pods are running, otherwise the GPU nodes will not be provisioned.
+:::
+
+Create a YAML file named `phi-4-workspace.yaml` with the following content. The `instanceType` field will specify what nodes will be auto-provisioned instead of only matching existing nodes in the BYO case.
+
+```yaml title="phi-4-workspace.yaml"
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-phi-4-mini
+resource:
+  instanceType: "g5.4xlarge" # Will trigger node creation
+  labelSelector:
+    matchLabels:
+      apps: phi-4-mini
+inference:
+  preset:
+    name: phi-4-mini-instruct
+```
+
+</TabItem>
+</Tabs>
+
+Apply your configuration to your cluster:
+
+```bash
+kubectl apply -f phi-4-workspace.yaml
+```
+
+## Monitor Deployment
+
+Track the workspace status to see when the model has been deployed successfully:
+
+```bash
+kubectl get workspace workspace-phi-4-mini
+```
+
+When the `WORKSPACESUCCEEDED` column becomes `True`, the model has been deployed successfully:
+
+```bash
+NAME                   INSTANCE                   RESOURCEREADY   INFERENCEREADY   JOBSTARTED   WORKSPACESUCCEEDED   AGE
+workspace-phi-4-mini   Standard_NC24ads_A100_v4   True            True                          True                 4h15m
+```
+
+:::note
+The `INSTANCE` column will default to `Standard_NC24ads_A100_v4` if you have not set up auto-provisioning. If you have auto-provisioning configured, it will show the specific instance type used.
+:::
+
+## Test the Model
+
+Find the inference service's cluster IP and test it using a temporary curl pod:
+
+```bash
+# Get the service endpoint
+kubectl get svc workspace-phi-4-mini
+export CLUSTERIP=$(kubectl get svc workspace-phi-4-mini -o jsonpath="{.spec.clusterIPs[0]}")
+
+# List available models
+kubectl run -it --rm --restart=Never curl --image=curlimages/curl -- curl -s http://$CLUSTERIP/v1/models | jq
+```
+
+You should see output similar to:
+
+```json
+{
+  "object": "list",
+  "data": [
+    {
+      "id": "phi-4-mini-instruct",
+      "object": "model",
+      "created": 1733370094,
+      "owned_by": "vllm",
+      "root": "/workspace/vllm/weights",
+      "parent": null,
+      "max_model_len": 16384
+    }
+  ]
+}
+```
+
+## Make an Inference Call
+
+Now make an inference call using the model:
+
+```bash
+kubectl run -it --rm --restart=Never curl --image=curlimages/curl -- curl -X POST http://$CLUSTERIP/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "phi-4-mini-instruct",
+    "messages": [{"role": "user", "content": "What is kubernetes?"}],
+    "max_tokens": 50,
+    "temperature": 0
+  }'
+```
+
+## Next Steps
+
+🎉 Congratulations! You've successfully deployed and tested your first model with KAITO.
+
+**What's Next:**
+
+- **Explore More Models**: Check out the full range of supported models in the [presets documentation](https://github.com/kaito-project/kaito/tree/main/presets)
+- **Advanced Configurations**: Learn about [workspace configurations](https://github.com/kaito-project/kaito/blob/main/api/v1alpha1/workspace_types.go)
+- **Custom Models**: See how to [contribute new models](https://github.com/kaito-project/kaito/blob/main/docs/How-to-add-new-models.md)
+
+**Additional Resources:**
+
+- [Fine-tuning Guide](https://github.com/kaito-project/kaito/tree/main/examples/fine-tuning) - Customize models for your specific use cases
+- [RAG Examples](https://github.com/kaito-project/kaito/tree/main/examples/RAG) - Retrieval-Augmented Generation patterns
+- [Troubleshooting Guide](https://github.com/kaito-project/kaito/blob/main/docs/troubleshooting.md) - Common issues and solutions
diff --git a/website/versioned_docs/version-v0.6.x/rag-api.md b/website/versioned_docs/version-v0.6.x/rag-api.md
new file mode 100644
index 0000000..5550caf
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/rag-api.md
@@ -0,0 +1,423 @@
+---
+title: API Definitions and Examples
+---
+
+A **RAGEngine index** is a logical collection that organizes and stores your documents for retrieval-augmented generation workflows. The relationship between indexes, documents, and document nodes is as follows:
+
+- **Index**: An index is a named container that holds a set of documents. Each index is independent and can be created, updated, queried, persisted, loaded, or deleted via the API.
+
+- **Documents**: Documents are the primary units of content that you add to an index. Each document contains a `text` field (the content to be indexed) and optional `metadata` (such as author, source, or custom tags). When you add documents to an index, each is assigned a unique `doc_id`.
+
+- **Document Nodes**: When a document is ingested, it is automatically split into smaller chunks called *nodes*. The splitting strategy depends on the document type and metadata:
+  - By default, documents are split into sentences.
+  - If you specify code-aware splitting (using the `split_type` and `language` metadata), the document is split into code blocks or logical code units.
+  - Each node represents a chunk of text that is indexed and can be retrieved as part of a query.
+
+**How it works in practice:**
+- When you index a document, it is divided into nodes for efficient retrieval and semantic search.
+- When you query an index, the engine retrieves the most relevant nodes (not necessarily whole documents) and can use them to generate answers or summaries.
+- The `source_nodes` field in query responses contains the actual nodes that matched your query, along with their scores and metadata.
+
+This design enables fine-grained retrieval and more accurate, context-aware responses from your LLM-powered applications.
+
+## Creating an Index With Documents
+
+To add documents to an index or create a new index, use the `/index` API route. This endpoint accepts a POST request with the index name and a list of documents to be indexed.
+
+### Create Index Request
+
+```json
+POST /index
+{
+  "index_name": "rag_index",
+  "documents": [
+    {
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ]
+}
+```
+
+- `index_name`: The name of the index to create or update.
+- `documents`: A list of documents, each with a `text` field and optional `metadata`.
+
+### Create Index Response
+
+```json
+[
+  {
+    "doc_id": "123456",
+    "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
+    "hash_value": "text_hash_value",
+    "metadata": {
+      "author": "kaito",
+    },
+    "is_truncated": false
+  }
+]
+```
+
+Each returned document includes its unique `doc_id`, the original text, metadata, and a flag indicating if the text was truncated. The `doc_id` will be important for document update/delete calls. 
+
+
+### Splitting Documents with CodeSplitter
+
+By default, RAGEngine splits documents into sentences. However, you can instruct the engine to split documents using the `CodeSplitter` (for code-aware chunking) by providing metadata in your API request.
+
+To use the `CodeSplitter`, set the `split_type` to `"code"` and specify the programming language in the `language` field of the document metadata. For example, when calling the RAGEngine API to index documents:
+
+```json
+{
+  "documents": [
+    {
+      "text": "def foo():\n    return 42\n\n# Another function\ndef bar():\n    pass",
+      "metadata": {
+        "split_type": "code",
+        "language": "python"
+      }
+    }
+  ]
+}
+```
+
+This instructs the RAGEngine to use code-aware splitting for the provided document. If `split_type` is not set or set to any other value, sentence splitting will be used by default.
+
+## List Documents
+
+To retrieve a paginated list of documents from a specific index, use the `/indexes/{index_name}/documents` API route. This endpoint accepts a GET request with optional query parameters for pagination, text truncation, and metadata filtering.
+
+### List Documents Request
+
+```
+GET /indexes/rag_index/documents?limit=5&offset=0&max_text_length=500
+```
+
+- `limit`: (optional) Maximum number of documents to return (default: 10, max: 100).
+- `offset`: (optional) Starting point for the document list (default: 0).
+- `max_text_length`: (optional) Maximum text length to return per document (default: 1000).
+- `metadata_filter`: (optional) A JSON string representing key-value pairs to filter documents by their metadata.
+
+### List Documents Response
+
+```json
+{
+  "documents": [
+    {
+      "doc_id": "123456",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data.",
+      "hash_value": "text_hash_value",
+      "metadata": {
+        "author": "kaito",
+      },
+      "is_truncated": false
+    }
+  ],
+  "count": 1
+}
+```
+
+Each document in the response includes its unique `doc_id`, the (possibly truncated) text, metadata, and a flag indicating if the text was truncated.
+
+**Note:**  
+If you want to filter documents by metadata, provide the `metadata_filter` parameter as a JSON string. For example:
+
+```
+GET /indexes/rag_index/documents?metadata_filter={"author":"kaito"}
+```
+
+
+## Updating Documents
+
+To update existing documents in a specific index, use the `/indexes/{index_name}/documents` API route. This endpoint accepts a POST request with the index name in the URL and a list of documents to update in the request body.
+
+### Update Documents Request
+
+```json
+POST /indexes/rag_index/documents
+{
+  "documents": [
+    {
+      "doc_id": "123456",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data. Adding an information retrieval system gives you control over grounding data used by an LLM when it formulates a response.",
+      "hash_value": "text_hash_value",
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ]
+}
+```
+
+- `doc_id`: The unique identifier of the document to update.
+- `text`: The new or updated text for the document.
+- `metadata`: (Optional) Updated metadata for the document.
+
+### Update Documents Response
+
+```json
+{
+  "updated_documents": [
+    {
+      "doc_id": "123456",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model (LLM) like ChatGPT by adding an information retrieval system that provides grounding data. Adding an information retrieval system gives you control over grounding data used by an LLM when it formulates a response.",
+      "hash_value": "text_hash_value",
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ],
+  "unchanged_documents": [],
+  "not_found_documents": []
+}
+```
+
+- `updated_documents`: Documents that were successfully updated.
+- `unchanged_documents`: Documents that were provided but did not require changes.
+- `not_found_documents`: Documents with IDs that were not found in the index.
+
+Use this endpoint to keep your indexed documents up to date with the latest content or metadata.
+
+## Delete Documents
+
+
+To delete one or more documents from a specific index, use the `/indexes/{index_name}/documents/delete` API route. This endpoint accepts a POST request with the index name in the URL and a list of document IDs to delete in the request body.
+
+### Delete Documents Request
+
+```json
+POST /indexes/rag_index/documents/delete
+{
+  "doc_ids": ["123456"]
+}
+```
+
+- `doc_ids`: A list of document IDs to delete from the specified index.
+
+### Delete Documents Response
+
+```json
+{
+  "deleted_doc_ids": ["123456"],
+  "not_found_doc_ids": []
+}
+```
+
+- `deleted_doc_ids`: Document IDs that were successfully deleted.
+- `not_found_doc_ids`: Document IDs that were not found in the index.
+
+Use this endpoint to remove documents that are no longer needed from your index.
+
+## Persist Index
+
+To save (persist) the data of an index to disk, use the `/persist/{index_name}` API route. This endpoint accepts a POST request with the index name in the URL and an optional `path` query parameter specifying where to save the index data.
+
+### Persist Index Request
+
+```
+POST /persist/rag_index?path=./custom_path
+```
+
+- `index_name`: The name of the index to persist.
+- `path`: (optional) The directory path where the index will be saved. If not provided, the default directory is used.
+
+### Persist Index Response
+
+```json
+{
+  "message": "Successfully persisted index rag_index to ./custom_path/rag_index."
+}
+```
+
+Use this endpoint to ensure your indexed data is safely stored on disk.
+
+---
+
+## Load Index
+
+To load an existing index from disk, use the `/load/{index_name}` API route. This endpoint accepts a POST request with the index name in the URL, an optional `path` query parameter specifying where to load the index from, and an optional `overwrite` flag.
+
+### Load Index Request
+
+```
+POST /load/rag_index?path=./custom_path/rag_index
+```
+
+- `index_name`: The name of the index to load.
+- `path`: (optional) The path to load the index from. If not provided, the default directory is used.
+- `overwrite`: (optional, default: false) If true, will overwrite the existing index if it already exists in memory.
+
+### Load Index Response
+
+```json
+{
+  "message": "Successfully loaded index rag_index from ./custom_path/rag_index."
+}
+```
+
+Use this endpoint to restore previously persisted indexes into memory for querying and updates.
+
+## Delete Index
+
+To delete an entire index and all of its documents, use the `/indexes/{index_name}` API route. This endpoint accepts a DELETE request with the index name in the URL. Deleting an index is irreversible and will remove all associated documents from memory.
+
+### Delete Index Request
+
+```
+DELETE /indexes/rag_index
+```
+
+- `index_name`: The name of the index to delete.
+
+### Delete Index Response
+
+```json
+{
+  "message": "Successfully deleted index rag_index."
+}
+```
+
+Use this endpoint to permanently remove an index and all its data when it is no longer needed.
+
+## Query Index
+
+To query a specific index for relevant documents, use the `/query` API route. This endpoint accepts a POST request with the index name, query string, and optional parameters for result count, and LLM generation.
+
+### Query Index Request
+
+```json
+POST /query
+{
+  "index_name": "rag_index",
+  "query": "What is RAG?",
+  "top_k": 5,
+  "llm_params": {
+    "temperature": 0.7,
+    "max_tokens": 2048
+  }
+}
+```
+
+- `index_name`: The name of the index to query.
+- `query`: The query string.
+- `top_k`: (optional) Number of top documents to retrieve (default: 5).
+- `llm_params`: (optional) Parameters for LLM-based generation (e.g., temperature, max_tokens).
+
+### Query Index Response
+
+```json
+{
+  "response": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
+  "source_nodes": [
+    {
+      "doc_id": "123456",
+      "node_id": "2853a565-8c1f-4982-acaa-a0ab52691435",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
+      "score": 0.95,
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ],
+  "metadata": {
+    "2853a565-8c1f-4982-acaa-a0ab52691435": {
+      "author": "kaito",
+    }
+  }
+}
+```
+
+- `response`: The generated answer or summary from the LLM (if enabled).
+- `source_nodes`: List of source nodes with their text, score, and metadata.
+- `metadata`: Additional metadata about the query or response.
+
+Use this endpoint to retrieve relevant information from your indexed documents and optionally generate answers using an LLM.
+
+## OpenAI-Compatible Chat Completions
+
+The RAGEngine provides an OpenAI-compatible chat completions endpoint at `/v1/chat/completions`. This endpoint allows you to use RAG capabilities with the familiar OpenAI API format, making it easy to integrate with existing applications that use OpenAI's chat completions.
+
+### RAG Bypass Conditions
+
+The chat completions endpoint automatically determines whether to use RAG or pass requests directly to the LLM based on the request content. Requests will **bypass RAG** and be sent directly to the LLM in the following cases:
+
+1. **No Index Specified**: When `index_name` is not provided in the request, the system treats it as a standard LLM request without retrieval.
+
+2. **Tool/Function Calls**: When the request contains `tools` or `functions` parameters, indicating the use of function calling capabilities.
+
+3. **Unsupported Message Roles**: When messages contain roles other than `user`, `system`, `assistant`, or `developer`. 
+
+4. **Non-Text Content**: When user messages contain non-text content types (such as images, audio, or other media formats).
+
+When any of these conditions are met, the request is passed through directly to the underlying LLM without document retrieval or augmentation, ensuring full compatibility with standard OpenAI API usage patterns.
+
+### Chat Completions Request
+
+```json
+POST /v1/chat/completions
+{
+  "index_name": "rag_index",
+  "model": "example_model",
+  "messages": [
+    {"role": "system", "content": "You are a knowledgeable assistant."},
+    {"role": "user", "content": "What is RAG?"}
+  ],
+  "temperature": 0.7,
+  "max_tokens": 2048,
+  "top_k": 5
+}
+```
+
+- `index_name`: (optional) The name of the index to query for relevant documents. If not included, the request will be sent to the LLM directly with no additional context.
+- `model`: The model identifier (for compatibility).
+- `messages`: Array of message objects with `role` and `content` fields.
+- `temperature`: (optional) Controls randomness in the response (0.0 to 1.0).
+- `max_tokens`: (optional) Maximum number of tokens to generate.
+- `top_k`: (optional) Number of top documents to retrieve from the index.
+
+### Chat Completions Response
+
+```json
+{
+  "id": "chatcmpl-123",
+  "object": "chat.completion",
+  "created": 1677652288,
+  "model": "example_model",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": "RAG stands for Retrieval-Augmented Generation. It's an architecture that augments the capabilities of a Large Language Model by adding an information retrieval system that provides grounding data."
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "source_nodes": [
+    {
+      "doc_id": "123456",
+      "node_id": "2853a565-8c1f-4982-acaa-a0ab52691435",
+      "text": "Retrieval Augmented Generation (RAG) is an architecture that augments the capabilities of a Large Language Model...",
+      "score": 0.95,
+      "metadata": {
+        "author": "kaito",
+      }
+    }
+  ]
+}
+```
+
+- `id`: Unique identifier for the chat completion.
+- `object`: Type of object returned (`chat.completion`).
+- `created`: Unix timestamp of when the completion was created.
+- `model`: The model used for the completion.
+- `choices`: Array of completion choices (typically one).
+- `source_nodes`: (RAG-specific) List of source documents that informed the response.
+
+Use this endpoint to integrate RAG capabilities into applications that already use OpenAI's chat completions API format.
+
+## Example Python Client
+
+You can leverage the [example_rag_client.py](./example_rag_client.py) as a starting point for a rag client with inputs that match the route documentation above.
diff --git a/website/versioned_docs/version-v0.6.x/rag.md b/website/versioned_docs/version-v0.6.x/rag.md
new file mode 100644
index 0000000..3671746
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/rag.md
@@ -0,0 +1,83 @@
+---
+title: Retrieval-Augmented Generation (RAG)
+---
+
+This document presents how to use the KAITO `ragengine` Custom Resource Definition (CRD) for retrieval-augumented generatoin workflow. By creating a RAGEngine resource, you can quickly stand up a service that indexes documents and queries them in conjunction with an existing LLM inference endpoint—no need to custom-build pipelines. This enables your large language model to answer questions based on your own private content.
+
+## Installation
+
+> Be sure you've cloned this repo and followed [kaito workspace installation](./installation.md) if you plan to use local embedding model. RAGEngine needs the gpu-provisioner component to provision GPU nodes.
+
+```bash
+helm install ragengine ./charts/kaito/ragengine --namespace kaito-ragengine --create-namespace
+```
+
+## Verify installation
+You can run the following commands to verify the installation of the controllers were successful.
+
+Check status of the Helm chart installations.
+
+```bash
+helm list -n kaito-ragengine
+```
+
+Check status of the `ragengine`.
+
+```bash
+kubectl describe deploy ragengine -n kaito-ragengine
+```
+
+## Clean up
+
+```bash
+helm uninstall kaito-ragengine
+```
+
+## Usage
+
+### Prerequisite
+Before creating a RAGEngine, ensure you have an accessible model inference endpoint. This endpoint can be:
+
+1.	A model deployed through KAITO Workspace CRD (e.g., a local Hugging Face model, a vLLM instance, etc.).
+2.	An external API (e.g., Huggingface service or other REST-based LLM providers).
+
+### Define the RAGEngine
+Create a YAML manifest defining your RAGEngine. Key fields under spec include:
+
+Embedding: how to generate vector embeddings for your documents. You may choose remote or local (one must be left unset if you pick the other):
+
+```yaml
+embedding:
+    local:
+      modelID: "BAAI/bge-small-en-v1.5"
+```
+InferenceService: points to the LLM endpoint that RAGEngine will call for final text generation.
+```yaml
+inferenceService:
+  url: "<inference-url>/v1/completions"
+```
+Users also need to specify the GPU SKU used for inference in the `compute` spec. For example,
+
+```yaml
+apiVersion: kaito.sh/v1alpha1
+kind: RAGEngine
+metadata:
+  name: ragengine-start
+spec:
+  compute:
+    instanceType: "Standard_NC4as_T4_v3"
+    labelSelector:
+      matchLabels:
+        apps: ragengine-example
+  embedding:
+    local:
+      modelID: "BAAI/bge-small-en-v1.5"
+  inferenceService:
+    url: "<inference-url>/v1/completions"
+```
+
+### Apply the manifest
+After you create your YAML configuration, run:
+```sh
+kubectl apply -f examples/RAG/kaito_ragengine_phi_3.yaml
+```
\ No newline at end of file
diff --git a/website/versioned_docs/version-v0.6.x/tool-calling.md b/website/versioned_docs/version-v0.6.x/tool-calling.md
new file mode 100644
index 0000000..586ed0f
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/tool-calling.md
@@ -0,0 +1,184 @@
+---
+title: Tool Calling
+---
+
+KAITO supports tool calling, allowing you to integrate external tools into your inference service. This feature enables the model to call APIs or execute functions based on the input it receives, enhancing its capabilities beyond text generation.
+
+## Requirements
+
+### Supported Inference Runtimes
+
+Currently, tool calling is only supported with the vLLM inference runtime.
+
+### Supported Models
+
+The following preset models are configured to support tool calling:
+- `phi-4-mini-instruct`
+- `phi-4`
+- `llama-3.1-8b-instruct`
+- `llama-3.3-70b-instruct`
+- `mistral-7b`
+- `mistral-7b-instruct`
+- `qwen2.5-coder-7b-instruct`
+- `qwen2.5-coder-32b-instruct`
+- `deepseek-r1-0528`
+- `deepseek-v3-0324`
+
+
+For more details on the inference configuration, refer to [vLLM tool calling documentation](https://docs.vllm.ai/en/latest/features/tool_calling.html).
+
+## Examples
+
+Assuming you have a running Workspace instance with the `phi-4-mini-tool-call` preset model, you can use the following examples to test tool calling. First, ensure that the inference service is running and accessible:
+
+```bash
+kubectl port-forward svc/workspace-phi-4-mini-tool-call 8000
+```
+
+### Named Function Calling
+
+![Named Function Calling](/img/function-calling.gif)
+*Source: Daily Dose of Data Science [^1]*
+
+```python
+from openai import OpenAI
+import json
+
+client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
+
+def get_weather(location: str, unit: str):
+    return f"Getting the weather for {location} in {unit}..."
+tool_functions = {"get_weather": get_weather}
+
+tools = [{
+    "type": "function",
+    "function": {
+        "name": "get_weather",
+        "description": "Get the current weather in a given location",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "location": {"type": "string", "description": "City and state, e.g., 'San Francisco, CA'"},
+                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
+            },
+            "required": ["location", "unit"]
+        }
+    }
+}]
+
+response = client.chat.completions.create(
+    model=client.models.list().data[0].id,
+    messages=[{"role": "user", "content": "What's the weather like in San Francisco?"}],
+    tools=tools,
+    tool_choice="auto"
+)
+
+tool_call = response.choices[0].message.tool_calls[0].function
+print(f"Function called: {tool_call.name}")
+print(f"Arguments: {tool_call.arguments}")
+print(f"Result: {tool_functions[tool_call.name](**json.loads(tool_call.arguments))}")
+```
+
+Expected output:
+
+```
+Function called: get_weather
+Arguments: {"location": "San Francisco, CA", "unit": "fahrenheit"}
+Result: Getting the weather for San Francisco, CA in fahrenheit...
+```
+
+### Model Context Protocol (MCP)
+
+With the right client framework, inference workload provisioned by KAITO can also call external tools using the [Model Context Protocol (MCP)](https://modelcontextprotocol.io/). This allows the model to integrate and share data with external tools, systems, and data sources.
+
+![MCP Overview](/img/mcp.gif)
+*Source: Daily Dose of Data Science [^1]*
+
+In the following example, we will use [uv](https://docs.astral.sh/uv/) to create a Python virtual environment and install the necessary dependencies for [AutoGen](https://microsoft.github.io/autogen/stable//index.html) to call the [DeepWiki](https://deepwiki.com/) MCP service and ask questions about the KAITO project.
+
+```bash
+mkdir kaito-mcp
+cd kaito-mcp
+# Create and activate a virtual environment
+uv venv && source .venv/bin/activate
+# Install dependencies
+uv pip install "autogen-ext[openai]" "autogen-agentchat" "autogen-ext[mcp]"
+```
+
+Create a Python script `test.py` with the following content:
+
+```python
+import asyncio
+
+from autogen_agentchat.agents import AssistantAgent
+from autogen_agentchat.ui import Console
+from autogen_core import CancellationToken
+from autogen_core.models import ModelFamily, ModelInfo
+from autogen_ext.models.openai import OpenAIChatCompletionClient
+from autogen_ext.tools.mcp import (StreamableHttpMcpToolAdapter,
+                                   StreamableHttpServerParams)
+from openai import OpenAI
+
+
+async def main() -> None:
+    # Create server params for the remote MCP service
+    server_params = StreamableHttpServerParams(
+        url="https://mcp.deepwiki.com/mcp",
+        timeout=30.0,
+        terminate_on_close=True,
+    )
+
+    # Get the ask_question tool from the server
+    adapter = await StreamableHttpMcpToolAdapter.from_server_params(server_params, "ask_question")
+
+    model = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy").models.list().data[0].id
+    model_info: ModelInfo = {
+        "vision": False,
+        "function_calling": True,
+        "json_output": True,
+        "family": ModelFamily.UNKNOWN,
+        "structured_output": True,
+        "multiple_system_messages": True,
+    }
+
+    # Create an agent that can use the ask_question tool
+    model_client = OpenAIChatCompletionClient(base_url="http://localhost:8000/v1", api_key="dummy", model=model, model_info=model_info)
+    agent = AssistantAgent(
+        name="deepwiki",
+        model_client=model_client,
+        tools=[adapter],
+        system_message="You are a helpful assistant.",
+    )
+
+    await Console(
+        agent.run_stream(task="In the GitHub repository 'kaito-project/kaito', how many preset models are there?", cancellation_token=CancellationToken())
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+To run the script, execute the following command in your terminal:
+
+```bash
+uv run test.py
+```
+
+Expected output:
+
+```
+---------- TextMessage (user) ----------
+In the GitHub repository 'kaito-project/kaito', how many preset models are there?
+
+---------- ToolCallRequestEvent (deepwiki) ----------
+[FunctionCall(id='chatcmpl-tool-4e22b15c32d34430b80078a3acc41f0d', arguments='{"repoName": "kaito-project/kaito", "question": "How many preset models are there?"}', name='ask_question')]
+
+---------- ToolCallExecutionEvent (deepwiki) ----------
+[FunctionExecutionResult(content='[{"type": "text", "text": "There are 16 preset models in the Kaito project.  These models are defined in the `supported_models.yaml` file  and registered programmatically within the codebase. ...", "annotations": null, "meta": null}]', name='ask_question', call_id='chatcmpl-tool-4e22b15c32d34430b80078a3acc41f0d', is_error=False)]
+
+---------- ToolCallSummaryMessage (deepwiki) ----------
+[{"type": "text", "text": "There are 16 preset models in the Kaito project.  These models are defined in the `supported_models.yaml` file  and registered programmatically within the codebase. ...", "annotations": null, "meta": null}]
+```
+
+[^1]: https://www.dailydoseofds.com/p/function-calling-mcp-for-llms/
diff --git a/website/versioned_docs/version-v0.6.x/tuning.md b/website/versioned_docs/version-v0.6.x/tuning.md
new file mode 100644
index 0000000..2f75fde
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/tuning.md
@@ -0,0 +1,168 @@
+---
+title: Fine Tuning
+---
+
+This document presents how to use the KAITO `workspace` Custom Resource Definition (CRD) for parameter-efficient fine-tuning (PEFT) of models, how a Kubernetes job is designed to automate the tuning workflow, and several best practices for troubleshooting.
+
+## Usage
+KAITO tuning APIs allow users to specify supported tuning methods like [LoRA or QLoRA](https://huggingface.co/docs/peft/main/en/conceptual_guides/lora), the input dataset and configuration settings, and the output destination for saving the tuning results. Currently, KAITO supports URL, image
+and Kubernetes volume as the types of tuning input sources, and image, Kubernetes volume as the types of tuning output destination.
+
+
+### Tuning workspace
+Here are two examples of using KAITO workspace CRD to define workspaces for tuning different models:
+
+Example 1: Tuning [`phi-3-mini`](../../../examples/fine-tuning/kaito_workspace_tuning_phi_3.yaml). This example uses a public dataset specified by a URL in the input.
+
+Example 2: Tuning `falcon-7b`. This example shows how to use an image as the source of input data.
+```yaml
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-tuning-falcon
+resource:
+  instanceType: "Standard_NC24ads_A100_v4"
+  labelSelector:
+    matchLabels:
+      app: tuning-falcon
+tuning:
+  preset:
+    name: falcon-7b
+  method: qlora
+  input:
+    image: PULLREGISTRY/DATA_NAME_HERE:0.0.1
+    imagePullSecrets:
+      - IMAGE_PULL_SECRETS_HERE
+  output:
+    image: PUSHREGISTRY/ADAPTER_NAME_HERE:0.0.1  # Tuning Output
+    imagePushSecret: IMAGE_PUSH_SECRET_HERE
+
+```
+
+Example 3: Tuning [`phi-3-mini`](../../../examples/fine-tuning/kaito_workspace_tuning_phi_3_with_pvc_volume.yaml). This example shows how to use a Kubernetes volume as the source of input dataset and output destination. We use AzureFile as an example, but any other supported volume type can be used. You should save your input dataset in the volume before creating the workspace, and the output adapter will be saved in the output volume after the tuning job is completed.
+
+The detailed `TuningSpec` API definitions can be found [here](https://github.com/kaito-project/kaito/blob/2ccc93daf9d5385649f3f219ff131ee7c9c47f3e/api/v1alpha1/workspace_types.go#L145).
+
+### Tuning configurations
+KAITO provides default tuning configurations for different tuning methods. They are managed by Kubernetes configmaps.
+- [default LoRA configmap](../../../charts/kaito/workspace/templates/lora-params.yaml)
+- [default QLoRA configmap](../../../charts/kaito/workspace/templates/qlora-params.yaml)
+
+## Tuning configmaps
+Users can specify a customized configmap via the `Config` field of the `TuningSpec`. The customized configmap should be structured based on the default configmaps provided by KAITO. Please read the following section carefully when attempting to change the default parameters used by KAITO.
+
+### Categorized key parameters
+Note that changing these parameters may largely impact the tuning result. In addition, users can add extra parameters that are not presented in the default configmaps. For a complete list of supported parameters, please refer to the provided huggingface documentation.
+
+ModelConfig([full list](https://huggingface.co/docs/transformers/v4.40.2/en/model_doc/auto#transformers.AutoModelForCausalLM.from_pretrained))
+- torch_dtype: Specifies the data type for PyTorch tensors, e.g., "bfloat16".
+- local_files_only: Indicates whether to only use local files.
+- device_map: Configures device mapping for the model, typically "auto".
+
+QuantizationConfig([full list](https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/quantization#transformers.BitsAndBytesConfig))
+- load_in_4bit: Enables loading the model with 4-bit precision.
+- bnb_4bit_quant_type: Specifies the type of 4-bit quantization, e.g., "nf4".
+- bnb_4bit_compute_dtype: Data type for computation, e.g., "bfloat16".
+- bnb_4bit_use_double_quant: Enables double quantization.
+
+LoraConfig([full list](https://huggingface.co/docs/peft/v0.8.2/en/package_reference/lora#peft.LoraConfig))
+- **r**: Rank of the low-rank matrices used in LoRA.
+- lora_alpha: Scaling factor for LoRA.
+- lora_dropout: Dropout rate for LoRA layers.
+
+TrainingArguments([full list](https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/trainer#transformers.TrainingArguments))
+- ddp_find_unused_parameters: Flag to handle unused parameters during distributed training.
+- save_strategy: Strategy for saving checkpoints, e.g., "epoch".
+- per_device_train_batch_size: Batch size per device during training.
+- num_train_epochs: Total number of training epochs to perform, defaults to 3.0.
+
+DataCollator([full list](https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/data_collator#transformers.DataCollatorForLanguageModeling))
+- mlm: Masked language modeling flag.
+
+DatasetConfig([full list](https://github.com/kaito-project/kaito/blob/main/presets/workspace/tuning/text-generation/cli.py#L44))
+- shuffle_dataset: Whether to shuffle the dataset.
+- train_test_split: Proportion of data used for training, typically set to 1 for using all data.
+
+## Input dataset format
+The input dataset for fine-tuning should follow specific formats defined in the HuggingFace trainer library. Supported formats include conversational and instruction formats.
+
+For example, [HuggingFace Dolly 15k OAI-style dataset](https://huggingface.co/datasets/philschmid/dolly-15k-oai-style/tree/main)
+
+- Conversational format
+  ```json
+  {
+    "messages": [
+      {"role": "system", "content": "Marv is a factual chatbot that is also sarcastic."},
+      {"role": "user", "content": "What's the capital of France?"},
+      {"role": "assistant", "content": "Paris, as if everyone doesn't know that already."}
+    ]
+  }
+  ```
+
+For example, [HuggingFace Instruction Dataset](https://huggingface.co/datasets/HuggingFaceH4/instruction-dataset/tree/main)
+
+- Instruction format
+  ```json
+  {"prompt": "<prompt text>", "completion": "<ideal generated text>"}
+  ```
+
+If your dataset is not in one of these formats, it will be passed directly to the training library ([SFTTrainer](https://huggingface.co/docs/trl/en/sft_trainer)) without any preprocessing. This may result in undefined behavior if the dataset does not align with the trainer's expected input structure. To ensure proper functionality, you may need to preprocess the dataset to match one of the supported formats. For more details, please refer to this [documentation](https://huggingface.co/docs/trl/v0.9.4/sft_trainer#dataset-format-support).
+
+
+Note: if you build a container image for the input dataset, please copy the dataset to the **`/data`** directory inside the container.
+
+# Tuning Job
+KAITO uses the Kubernetes **batchv1.job** workload to manage the tuning Pod. When a tuning workspace custom resource is created, the KAITO controller will create a job with the same name as the workspace in the same namespace. The pod structure will be different depending on whether the input and output are specified as URLs, images, or Kubernetes volumes.
+
+## Pod structure when input and output are specified as URLs or images
+To streamline the tuning workflow, KAITO adds two containers in addition to the main container that runs the tuning process. The pod structure is illustrated in Figure 1.
+
+![KAITO fine tuning](/img/kaito-fine-tuning.png)
+
+Figure 1. KAITO tuning pod structure.
+
+- Initcontainer `data-downloader`: It downloads the training input dataset from the URLs specified in the tuning spec if needed. If an image is specified in the input, the `data-downloader` container uses the specified image as the container image. This initcontainer ensures the training data is available locally before the training process starts.
+
+- Sidecar container: It is introduced to support automatically pushing the tuning results to a container registry. This container, with `docker` installed, runs a script to periodically check the training progress. Once the training is done, indicated by a sentinel file created by the training process, the script builds a container image containing the training results and pushes the image to the specified container registry.
+
+- Main container: It uses one of the supported model images. The image entry launches the [fine\_tuning.py](https://github.com/kaito-project/kaito/blob/main/presets/workspace/tuning/text-generation/fine_tuning.py) script.
+
+All three containers use shared local volumes (by mounting the same `EmptyDir` volumes), hence file copies between containers are avoided.
+
+## Pod structure when input and output are specified as Kubernetes volumes
+When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and output destination are already available in the specified volumes, the main container can directly access them.
+
+Other than the absence of the init and sidecar containers, the main container is the same as described in the previous section.
+
+# Troubleshooting
+
+### Job pod failures
+When the tuning job reaches the failed state, at least one of the above three containers has encountered errors. Users can check the logs of these containers using the `kubectl logs PODNAME -n NAMESPACE -c CONTAINERNAME` command.
+
+For the initcontainer and sidecar container, possible errors include invalid input/output URLs or invalid image pull secrets. Users can fix these problems by updating the workspace custom resource with corrections. The KAITO controller will create a new job using the updated spec.
+
+For the main container, errors may occur when CUDA reports out of GPU memory. Users should reduce the batch size (the default is 1) if it has been customized to a value larger than 1. If the batch size is already 1, the workspace must be recreated using a different GPU SKU with larger GPU memory. Note that KAITO has optimized the training memory usage by dropping the preallocated memory cache. Our internal tests show that the performance impact due to this change is negligible.
+
+### LoraConfig target modules errors
+If you encounter the error: `ValueError: Target modules {'target_module_here'} not found in the base model. Please check the target modules and try again.`, you need to manually specify the target_modules parameter in your KAITO configmap. This is a huggingface requirement.
+
+This error occurs because the automatic module detection failed for your model. You must identify the specific target modules for your model and explicitly list them in the LoraConfig section. For example, a valid configuration for phi-4-mini-instruct would look like:
+```yaml
+LoraConfig: # Configurable Parameters: https://huggingface.co/docs/peft/v0.8.2/en/package_reference/lora#peft.LoraConfig
+    r: 8
+    lora_alpha: 8
+    lora_dropout: 0.0
+    target_modules: ["qkv_proj", "o_proj"]
+```
+You can find valid module names by checking the "Files info" tab on the model card on Huggingface and examining the model.safetensors.index.json file.
+
+For more information on target_modules configuration, see the official [PEFT documentation](https://huggingface.co/docs/peft/en/package_reference/lora#peft.LoraConfig.target_modules).
+
+### Time for job completion
+The training job can take a long time depending on the size of the input dataset and training pipeline configurations. The total training time is largely determined by the total number of training steps, calculated as:
+```
+total steps = number of epochs * (number of samples in dataset / batch size)
+```
+where `number of epochs` and `batch size` can be customized in the tuning configmap. However, if the `max_steps` parameter is also specified in the configmap, training will stop after reaching the max steps, even if the specified epochs have not been completed. Users can track the tuning progress in the job pod's log, reported by the number of steps completed out of the total.
+
+Please file issues if you experience abnormal slowness of the training job.
diff --git a/website/versioned_docs/version-v0.6.x/usage.md b/website/versioned_docs/version-v0.6.x/usage.md
new file mode 100644
index 0000000..0857d36
--- /dev/null
+++ b/website/versioned_docs/version-v0.6.x/usage.md
@@ -0,0 +1,9 @@
+---
+title: Usage
+---
+
+The detailed usage for Kaito supported models can be found in [**HERE**](./presets.md). In case users want to deploy their own containerized models, they can provide the pod template in the `inference` field of the workspace custom resource (please see [API definitions](../../../api/v1alpha1/workspace_types.go) for details). The controller will create a deployment workload using all provisioned GPU nodes. Note that currently the controller does **NOT** handle automatic model upgrade. It only creates inference workloads based on the preset configurations if the workloads do not exist.
+
+The number of the supported models in Kaito is growing! Please check [this](./preset-onboarding.md) document to see how to add a new supported model.
+
+Starting with version v0.3.0, Kaito supports model fine-tuning and using fine-tuned adapters in the inference service. Refer to the [tuning document](./tuning.md) and [inference document](./inference.md) for more information.
diff --git a/website/versioned_sidebars/version-v0.6.x-sidebars.json b/website/versioned_sidebars/version-v0.6.x-sidebars.json
new file mode 100644
index 0000000..2bc5619
--- /dev/null
+++ b/website/versioned_sidebars/version-v0.6.x-sidebars.json
@@ -0,0 +1,79 @@
+{
+  "sidebar": [
+    {
+      "type": "category",
+      "label": "Getting Started",
+      "collapsed": false,
+      "items": [
+        "intro",
+        "installation",
+        "quick-start",
+        "presets",
+        "usage",
+        "faq"
+      ]
+    },
+    {
+      "type": "category",
+      "label": "Cloud Providers",
+      "collapsed": false,
+      "items": [
+        "azure",
+        "aws"
+      ]
+    },
+    {
+      "type": "category",
+      "label": "Features",
+      "collapsed": false,
+      "items": [
+        "inference",
+        "multi-node-inference",
+        "tuning",
+        {
+          "type": "category",
+          "label": "Retrieval-Augmented Generation (RAG)",
+          "link": {
+            "type": "doc",
+            "id": "rag"
+          },
+          "items": [
+            "rag-api"
+          ]
+        },
+        "custom-model",
+        "tool-calling",
+        "model-as-oci-artifacts",
+        "headlamp-kaito"
+      ]
+    },
+    {
+      "type": "category",
+      "label": "Integrations",
+      "collapsed": false,
+      "items": [
+        "aikit"
+      ]
+    },
+    {
+      "type": "category",
+      "label": "Operations",
+      "collapsed": false,
+      "items": [
+        "monitoring",
+        "kaito-oom-prevention",
+        "kaito-on-byo-gpu-nodes"
+      ]
+    },
+    {
+      "type": "category",
+      "label": "Contributing",
+      "collapsed": false,
+      "items": [
+        "contributing",
+        "preset-onboarding",
+        "proposals"
+      ]
+    }
+  ]
+}
diff --git a/website/versions.json b/website/versions.json
index b02df32..22784ab 100644
--- a/website/versions.json
+++ b/website/versions.json
@@ -1,4 +1,5 @@
 [
+  "v0.6.x",
   "v0.5.1",
   "v0.5.0"
 ]
diff --git a/website/yarn.lock b/website/yarn.lock
index 2d170f8..122ae59 100644
--- a/website/yarn.lock
+++ b/website/yarn.lock
@@ -5028,6 +5028,11 @@ fs.realpath@^1.0.0:
   resolved "https://registry.npmjs.org/fs.realpath/-/fs.realpath-1.0.0.tgz"
   integrity sha512-OO0pH2lK6a0hZnAdau5ItzHPI6pUlvI7jMVnxUQRtw4owF2wk8lOSabtGDCTP4Ggrg2MbGnWO9X8K1t4+fGMDw==
 
+fsevents@~2.3.2:
+  version "2.3.3"
+  resolved "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz"
+  integrity sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==
+
 function-bind@^1.1.2:
   version "1.1.2"
   resolved "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz"



Test patch:
diff --git a/pkg/ragengine/controllers/preset-rag_test.go b/pkg/ragengine/controllers/preset_rag_test.go
similarity index 100%
rename from pkg/ragengine/controllers/preset-rag_test.go
rename to pkg/ragengine/controllers/preset_rag_test.go
diff --git a/pkg/workspace/controllers/workspace_controller_test.go b/pkg/workspace/controllers/workspace_controller_test.go
index 9e7c0c4..01292f4 100644
--- a/pkg/workspace/controllers/workspace_controller_test.go
+++ b/pkg/workspace/controllers/workspace_controller_test.go
@@ -24,19 +24,21 @@ import (
 	"time"
 
 	"github.com/awslabs/operatorpkg/status"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
-	"gotest.tools/assert"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/utils/ptr"
 	"knative.dev/pkg/apis"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	karpenterv1 "sigs.k8s.io/karpenter/pkg/apis/v1"
 
 	"github.com/kaito-project/kaito/api/v1beta1"
+	"github.com/kaito-project/kaito/pkg/featuregates"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
 	"github.com/kaito-project/kaito/pkg/utils/test"
@@ -378,7 +380,7 @@ func TestEnsureService(t *testing.T) {
 
 			err := reconciler.ensureService(ctx, tc.workspace)
 			if tc.expectedError == nil {
-				assert.Check(t, err == nil, "Not expected to return error")
+				assert.NoError(t, err)
 			} else {
 				assert.Equal(t, tc.expectedError.Error(), err.Error())
 			}
@@ -494,7 +496,7 @@ func TestApplyInferenceWithPreset(t *testing.T) {
 
 			err := reconciler.applyInference(ctx, &tc.workspace)
 			if tc.expectedError == nil {
-				assert.Check(t, err == nil, fmt.Sprintf("Not expected to return error: %v", err))
+				assert.NoError(t, err)
 			} else {
 				assert.Equal(t, tc.expectedError.Error(), err.Error())
 			}
@@ -551,7 +553,7 @@ func TestApplyInferenceWithTemplate(t *testing.T) {
 
 			err := reconciler.applyInference(ctx, &tc.workspace)
 			if tc.expectedError == nil {
-				assert.Check(t, err == nil, "Not expected to return error")
+				assert.NoError(t, err)
 			} else {
 				assert.Equal(t, tc.expectedError.Error(), err.Error())
 			}
@@ -560,6 +562,10 @@ func TestApplyInferenceWithTemplate(t *testing.T) {
 }
 
 func TestGetAllQualifiedNodes(t *testing.T) {
+	// Save and restore feature gate state
+	originalNAP := featuregates.FeatureGates[consts.FeatureFlagDisableNodeAutoProvisioning]
+	defer func() { featuregates.FeatureGates[consts.FeatureFlagDisableNodeAutoProvisioning] = originalNAP }()
+
 	deletedNode := corev1.Node{
 		ObjectMeta: v1.ObjectMeta{
 			Name: "node4",
@@ -571,12 +577,16 @@ func TestGetAllQualifiedNodes(t *testing.T) {
 	}
 	mockWorkspaceWithPreferredNodes := test.MockWorkspaceWithPreferredNodes.DeepCopy()
 	mockWorkspaceWithPreferredNodes.Resource.PreferredNodes = []string{"node-p1", "node-p2"}
+	mockWorkspaceWithPreferredNodes.Resource.Count = ptr.To(2)
+
+	mockWorkspaceWithSinglePreferredNode := test.MockWorkspaceWithPreferredNodes.DeepCopy()
 
 	testcases := map[string]struct {
 		callMocks     func(c *test.MockClient)
 		workspace     *v1beta1.Workspace
 		expectedError error
 		expectedNodes []string
+		disableNAP    bool
 	}{
 		"Fails to get qualified nodes because can't list nodes": {
 			callMocks: func(c *test.MockClient) {
@@ -585,6 +595,7 @@ func TestGetAllQualifiedNodes(t *testing.T) {
 			workspace:     test.MockWorkspaceDistributedModel,
 			expectedError: errors.New("Failed to list nodes"),
 			expectedNodes: nil,
+			disableNAP:    false,
 		},
 		"Gets all qualified nodes": {
 			callMocks: func(c *test.MockClient) {
@@ -606,6 +617,7 @@ func TestGetAllQualifiedNodes(t *testing.T) {
 			workspace:     test.MockWorkspaceDistributedModel,
 			expectedError: nil,
 			expectedNodes: []string{"node1"},
+			disableNAP:    false,
 		},
 		"Gets all qualified nodes with preferred": {
 			callMocks: func(c *test.MockClient) {
@@ -676,14 +688,187 @@ func TestGetAllQualifiedNodes(t *testing.T) {
 
 				c.On("List", mock.IsType(context.Background()), mock.IsType(&corev1.NodeList{}), mock.Anything).Return(nil)
 			},
-			workspace:     mockWorkspaceWithPreferredNodes,
+			workspace:     mockWorkspaceWithSinglePreferredNode,
 			expectedError: nil,
 			expectedNodes: []string{"node-p1"},
 		},
+		"NAP disabled: all preferred nodes present and ready, returns all": {
+			callMocks: func(c *test.MockClient) {
+				nodeList := &corev1.NodeList{Items: []corev1.Node{
+					{
+						ObjectMeta: v1.ObjectMeta{
+							Name: "node-p1",
+							Labels: map[string]string{
+								corev1.LabelInstanceTypeStable: "vendor1",
+								"apps":                         "test",
+							},
+						},
+						Status: corev1.NodeStatus{
+							Conditions: []corev1.NodeCondition{{
+								Type:   corev1.NodeReady,
+								Status: corev1.ConditionTrue,
+							}},
+						},
+					},
+					{
+						ObjectMeta: v1.ObjectMeta{
+							Name: "node-p2",
+							Labels: map[string]string{
+								corev1.LabelInstanceTypeStable: "vendor2",
+								"apps":                         "test",
+							},
+						},
+						Status: corev1.NodeStatus{
+							Conditions: []corev1.NodeCondition{{
+								Type:   corev1.NodeReady,
+								Status: corev1.ConditionTrue,
+							}},
+						},
+					},
+				}}
+
+				c.CreateOrUpdateObjectInMap(&nodeList.Items[0])
+				c.CreateOrUpdateObjectInMap(&nodeList.Items[1])
+
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+
+			},
+			workspace:     mockWorkspaceWithPreferredNodes.DeepCopy(),
+			expectedError: nil,
+			expectedNodes: []string{"node-p1", "node-p2"},
+			disableNAP:    true,
+		},
+		"NAP disabled: one preferred node not ready, returns error": {
+			callMocks: func(c *test.MockClient) {
+				nodeList := &corev1.NodeList{Items: []corev1.Node{
+					{
+						ObjectMeta: v1.ObjectMeta{
+							Name: "node-p1",
+							Labels: map[string]string{
+								corev1.LabelInstanceTypeStable: "vendor1",
+								"apps":                         "test",
+							},
+						},
+						Status: corev1.NodeStatus{
+							Conditions: []corev1.NodeCondition{{
+								Type:   corev1.NodeReady,
+								Status: corev1.ConditionTrue,
+							}},
+						},
+					},
+					{
+						ObjectMeta: v1.ObjectMeta{
+							Name: "node-p2",
+							Labels: map[string]string{
+								corev1.LabelInstanceTypeStable: "vendor2",
+								"apps":                         "test",
+							},
+						},
+						Status: corev1.NodeStatus{
+							Conditions: []corev1.NodeCondition{{
+								Type:   corev1.NodeReady,
+								Status: corev1.ConditionFalse,
+							}},
+						},
+					},
+				}}
+
+				c.CreateOrUpdateObjectInMap(&nodeList.Items[0])
+				c.CreateOrUpdateObjectInMap(&nodeList.Items[1])
+
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+			},
+			workspace:     mockWorkspaceWithPreferredNodes.DeepCopy(),
+			expectedError: errors.New("when node auto-provisioning is disabled, at least 2 preferred nodes must match the label selector and be ready and not deleting, only have 1"),
+			expectedNodes: nil,
+			disableNAP:    true,
+		},
+		"NAP disabled: one preferred node deleting, returns error": {
+			callMocks: func(c *test.MockClient) {
+				deletingTime := v1.Time{Time: time.Now()}
+				nodeList := &corev1.NodeList{Items: []corev1.Node{
+					{
+						ObjectMeta: v1.ObjectMeta{
+							Name: "node-p1",
+							Labels: map[string]string{
+								corev1.LabelInstanceTypeStable: "vendor1",
+								"apps":                         "test",
+							},
+						},
+						Status: corev1.NodeStatus{
+							Conditions: []corev1.NodeCondition{{
+								Type:   corev1.NodeReady,
+								Status: corev1.ConditionTrue,
+							}},
+						},
+					},
+					{
+						ObjectMeta: v1.ObjectMeta{
+							Name: "node-p2",
+							Labels: map[string]string{
+								corev1.LabelInstanceTypeStable: "vendor2",
+								"apps":                         "test",
+							},
+							DeletionTimestamp: &deletingTime,
+						},
+						Status: corev1.NodeStatus{
+							Conditions: []corev1.NodeCondition{{
+								Type:   corev1.NodeReady,
+								Status: corev1.ConditionTrue,
+							}},
+						},
+					},
+				}}
+
+				c.CreateOrUpdateObjectInMap(&nodeList.Items[0])
+				c.CreateOrUpdateObjectInMap(&nodeList.Items[1])
+
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+			},
+			workspace:     mockWorkspaceWithPreferredNodes.DeepCopy(),
+			expectedError: errors.New("when node auto-provisioning is disabled, at least 2 preferred nodes must match the label selector and be ready and not deleting, only have 1"),
+			expectedNodes: nil,
+			disableNAP:    true,
+		},
+		"NAP disabled: missing preferred node, returns error": {
+			callMocks: func(c *test.MockClient) {
+				nodeList := &corev1.NodeList{Items: []corev1.Node{
+					{
+						ObjectMeta: v1.ObjectMeta{
+							Name: "node-p1",
+							Labels: map[string]string{
+								corev1.LabelInstanceTypeStable: "vendor1",
+								"apps":                         "test",
+							},
+						},
+						Status: corev1.NodeStatus{
+							Conditions: []corev1.NodeCondition{{
+								Type:   corev1.NodeReady,
+								Status: corev1.ConditionTrue,
+							}},
+						},
+					},
+				}}
+
+				c.CreateOrUpdateObjectInMap(&nodeList.Items[0])
+
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(nil)
+				c.On("Get", mock.IsType(context.Background()), mock.Anything, mock.IsType(&corev1.Node{}), mock.Anything).Return(errors.New("nodes \"node-p2\" not found"))
+			},
+			workspace:     mockWorkspaceWithPreferredNodes.DeepCopy(),
+			expectedError: errors.New("when node auto-provisioning is disabled, at least 2 preferred nodes must match the label selector and be ready and not deleting, only have 1"),
+			expectedNodes: nil,
+			disableNAP:    true,
+		},
 	}
 
 	for k, tc := range testcases {
 		t.Run(k, func(t *testing.T) {
+			featuregates.FeatureGates[consts.FeatureFlagDisableNodeAutoProvisioning] = tc.disableNAP
+
 			mockClient := test.NewClient()
 			reconciler := &WorkspaceReconciler{
 				Client: mockClient,
@@ -696,14 +881,15 @@ func TestGetAllQualifiedNodes(t *testing.T) {
 			nodes, err := reconciler.getAllQualifiedNodes(ctx, tc.workspace)
 
 			if tc.expectedError != nil {
-				assert.Equal(t, tc.expectedError.Error(), err.Error())
-				assert.Check(t, nodes == nil, "Response node array should be nil")
+				assert.NotNil(t, err)
+				assert.Equal(t, tc.expectedError, err)
+				assert.Nil(t, nodes)
 				return
 			}
 
-			assert.Check(t, err == nil, "Not expected to return error")
-			assert.Check(t, nodes != nil, "Response node array should not be nil")
-			assert.Check(t, len(nodes) == len(tc.expectedNodes), "Unexpected qualified nodes")
+			assert.NoError(t, err)
+			assert.NotNil(t, nodes)
+			assert.Equal(t, len(tc.expectedNodes), len(nodes))
 		})
 	}
 }
@@ -855,7 +1041,7 @@ func TestApplyWorkspaceResource(t *testing.T) {
 
 			err := reconciler.applyWorkspaceResource(ctx, &tc.workspace)
 			if tc.expectedError == nil {
-				assert.Check(t, err == nil, "Not expected to return error")
+				assert.NoError(t, err)
 			} else {
 				assert.Equal(t, tc.expectedError.Error(), err.Error())
 			}
@@ -1062,7 +1248,7 @@ func TestSyncControllerRevision(t *testing.T) {
 
 			err := reconciler.syncControllerRevision(ctx, &tc.workspace)
 			if tc.expectedError == nil {
-				assert.Check(t, err == nil, "Not expected to return error")
+				assert.NoError(t, err)
 			} else {
 				assert.Equal(t, tc.expectedError.Error(), err.Error())
 			}
diff --git a/presets/ragengine/tests/api/conftest.py b/presets/ragengine/tests/api/conftest.py
index 5f20223..62c6d3c 100644
--- a/presets/ragengine/tests/api/conftest.py
+++ b/presets/ragengine/tests/api/conftest.py
@@ -12,32 +12,40 @@
 # limitations under the License.
 
 
-import sys
 import os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # Force CPU-only execution for testing
-os.environ["OMP_NUM_THREADS"] = "1" # Force single-threaded for testing to prevent segfault while loading embedding model
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force CPU-only execution for testing
+# Force single-threaded for testing to prevent segfault while loading embedding model
+os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["MKL_NUM_THREADS"] = "1"  # Force MKL to use a single thread
 
-import nest_asyncio
 import asyncio
-import pytest_asyncio
+
 import httpx
+import nest_asyncio
+import pytest_asyncio
+
 from ragengine.main import app, vector_store_handler
+
 nest_asyncio.apply()
 
+
 @pytest_asyncio.fixture
 async def async_client():
     """Use an async HTTP client to interact with FastAPI app."""
     async with httpx.AsyncClient(app=app, base_url="http://localhost") as client:
         yield client
 
+
 @pytest_asyncio.fixture(scope="session")
 def event_loop():
     loop = asyncio.new_event_loop()
     yield loop
     loop.close()
 
+
 @pytest_asyncio.fixture(autouse=True)
 def clear_index():
-    vector_store_handler.index_map.clear()
\ No newline at end of file
+    vector_store_handler.index_map.clear()
diff --git a/presets/ragengine/tests/api/test_chat_completions.py b/presets/ragengine/tests/api/test_chat_completions.py
index 2e53806..25c4aa6 100644
--- a/presets/ragengine/tests/api/test_chat_completions.py
+++ b/presets/ragengine/tests/api/test_chat_completions.py
@@ -11,20 +11,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import re
+import time
 from unittest.mock import patch
 
-import pytest
 import httpx
+import pytest
 import respx
-import time
-import re
+
 
 @pytest.fixture(autouse=True)
 def overwrite_inference_url(monkeypatch):
-    import ragengine.inference.inference
     import ragengine.config
-    monkeypatch.setattr(ragengine.config, "LLM_INFERENCE_URL", "http://localhost:5000/v1/chat/completions")
-    monkeypatch.setattr(ragengine.inference.inference, "LLM_INFERENCE_URL", "http://localhost:5000/v1/chat/completions")
+    import ragengine.inference.inference
+
+    monkeypatch.setattr(
+        ragengine.config,
+        "LLM_INFERENCE_URL",
+        "http://localhost:5000/v1/chat/completions",
+    )
+    monkeypatch.setattr(
+        ragengine.inference.inference,
+        "LLM_INFERENCE_URL",
+        "http://localhost:5000/v1/chat/completions",
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
@@ -48,26 +59,24 @@ async def test_chat_completions_basic_success(mock_get, async_client):
                 "index": 0,
                 "message": {
                     "role": "assistant",
-                    "content": "This is a helpful response about the test document."
+                    "content": "This is a helpful response about the test document.",
                 },
-                "finish_reason": "stop"
+                "finish_reason": "stop",
             }
         ],
-        "usage": {
-            "prompt_tokens": 25,
-            "completion_tokens": 12,
-            "total_tokens": 37
-        }
+        "usage": {"prompt_tokens": 25, "completion_tokens": 12, "total_tokens": 37},
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index some test documents
     index_request = {
         "index_name": "test_index",
         "documents": [
             {"text": "This is a test document about AI and machine learning."},
-            {"text": "Another document discussing natural language processing."}
-        ]
+            {"text": "Another document discussing natural language processing."},
+        ],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -77,17 +86,15 @@ async def test_chat_completions_basic_success(mock_get, async_client):
     chat_request = {
         "index_name": "test_index",
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "What can you tell me about AI?"}
-        ],
+        "messages": [{"role": "user", "content": "What can you tell me about AI?"}],
         "temperature": 0.7,
         "max_tokens": 100,
-        "top_k": 2
+        "top_k": 2,
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
     assert "id" in response_data
     assert response_data["object"] == "chat.completion"
@@ -95,16 +102,35 @@ async def test_chat_completions_basic_success(mock_get, async_client):
     assert response_data["model"] == "mock-model"
     assert len(response_data["choices"]) == 1
     assert response_data["choices"][0]["message"]["role"] == "assistant"
-    assert response_data["choices"][0]["message"]["content"] == "This is a helpful response about the test document."
+    assert (
+        response_data["choices"][0]["message"]["content"]
+        == "This is a helpful response about the test document."
+    )
     assert response_data["choices"][0]["finish_reason"] == "stop"
     assert response_data["choices"][0]["index"] == 0
     assert "source_nodes" in response_data
     assert len(response_data["source_nodes"]) > 0
 
-    response = await async_client.get(f"/metrics")
+    response = await async_client.get("/metrics")
     assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_chat_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_chat_requests_total{status="success"} ([1-9]\d*).0', response.text
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
@@ -123,33 +149,42 @@ async def test_chat_completions_without_index_name(mock_get, async_client):
         "object": "chat.completion",
         "created": int(time.time()),
         "model": "mock-model",
-        "choices": [{
-            "index": 0,
-            "message": {"role": "assistant", "content": "This is a direct LLM response"},
-            "finish_reason": "stop"
-        }]
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "This is a direct LLM response",
+                },
+                "finish_reason": "stop",
+            }
+        ],
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Test request without index_name (should trigger passthrough)
     chat_request = {
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "Hello, how are you?"}
-        ],
+        "messages": [{"role": "user", "content": "Hello, how are you?"}],
         "temperature": 0.7,
-        "max_tokens": 100
+        "max_tokens": 100,
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
     assert response_data["id"] == "chatcmpl-test123"
-    assert response_data["choices"][0]["message"]["content"] == "This is a direct LLM response"
+    assert (
+        response_data["choices"][0]["message"]["content"]
+        == "This is a direct LLM response"
+    )
     # Should have source_nodes field but it should be None for passthrough requests
     assert response_data["source_nodes"] is None
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -167,24 +202,35 @@ async def test_chat_completions_with_tools(mock_get, async_client):
         "object": "chat.completion",
         "created": int(time.time()),
         "model": "mock-model",
-        "choices": [{
-            "index": 0,
-            "message": {
-                "role": "assistant", 
-                "content": None,
-                "tool_calls": [{"id": "call123", "type": "function", "function": {"name": "test_tool", "arguments": '{"param1": "value1"}'}}]
-            },
-            "finish_reason": "tool_calls"
-        }]
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {
+                            "id": "call123",
+                            "type": "function",
+                            "function": {
+                                "name": "test_tool",
+                                "arguments": '{"param1": "value1"}',
+                            },
+                        }
+                    ],
+                },
+                "finish_reason": "tool_calls",
+            }
+        ],
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Test request with tools (should trigger passthrough)
     chat_request = {
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "Use a tool to help me"}
-        ],
+        "messages": [{"role": "user", "content": "Use a tool to help me"}],
         "tools": [
             {
                 "type": "function",
@@ -194,36 +240,39 @@ async def test_chat_completions_with_tools(mock_get, async_client):
                     "parameters": {
                         "type": "object",
                         "properties": {
-                            "param1": {"type": "string", "description": "A test parameter"}
+                            "param1": {
+                                "type": "string",
+                                "description": "A test parameter",
+                            }
                         },
-                    }
-                }
+                    },
+                },
             }
-        ]
+        ],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
     assert response_data["choices"][0]["finish_reason"] == "tool_calls"
     assert "tool_calls" in response_data["choices"][0]["message"]
 
+
 @pytest.mark.asyncio
 async def test_chat_completions_nonexistent_index(async_client):
     """Test chat completion with non-existent index."""
     chat_request = {
         "index_name": "nonexistent_index",
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "Test question"}
-        ]
+        "messages": [{"role": "user", "content": "Test question"}],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 404
     assert "No such index: 'nonexistent_index' exists" in response.json()["detail"]
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -236,17 +285,18 @@ async def test_chat_completions_invalid_request_format(mock_get, async_client):
     }
 
     # Mock HTTPX response for passthrough LLM call (in case it gets that far)
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(400, json={"error": "Invalid request"}))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(400, json={"error": "Invalid request"})
+    )
 
     # Test missing messages
-    chat_request = {
-        "model": "mock-model"
-    }
+    chat_request = {"model": "mock-model"}
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 400
     assert "Invalid request" in response.json()["detail"]
 
+
 @pytest.mark.asyncio
 async def test_chat_completions_missing_role_in_message(async_client):
     """Test chat completion with message missing role."""
@@ -255,8 +305,8 @@ async def test_chat_completions_missing_role_in_message(async_client):
         "index_name": "test_index",
         "documents": [
             {"text": "This is a test document about AI and machine learning."},
-            {"text": "Another document discussing natural language processing."}
-        ]
+            {"text": "Another document discussing natural language processing."},
+        ],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -265,18 +315,17 @@ async def test_chat_completions_missing_role_in_message(async_client):
     chat_request = {
         "index_name": "test_index",
         "model": "mock-model",
-        "messages": [
-            {"content": "What can you tell me about AI?"}
-        ],
+        "messages": [{"content": "What can you tell me about AI?"}],
         "temperature": 0.7,
         "max_tokens": 100,
-        "top_k": 2
+        "top_k": 2,
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 400
     assert "messages must contain 'role'" in response.json()["detail"]
 
+
 @pytest.mark.asyncio
 async def test_chat_completions_missing_content_for_user_role(async_client):
     """Test chat completion with user message missing content."""
@@ -285,8 +334,8 @@ async def test_chat_completions_missing_content_for_user_role(async_client):
         "index_name": "test_index",
         "documents": [
             {"text": "This is a test document about AI and machine learning."},
-            {"text": "Another document discussing natural language processing."}
-        ]
+            {"text": "Another document discussing natural language processing."},
+        ],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -295,17 +344,18 @@ async def test_chat_completions_missing_content_for_user_role(async_client):
     chat_request = {
         "index_name": "test_index",
         "model": "mock-model",
-        "messages": [
-            {"role": "user"}
-        ],
+        "messages": [{"role": "user"}],
         "temperature": 0.7,
         "max_tokens": 100,
-        "top_k": 2
+        "top_k": 2,
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 400
-    assert "messages must contain 'content' for role 'user'" in response.json()["detail"]
+    assert (
+        "messages must contain 'content' for role 'user'" in response.json()["detail"]
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
@@ -329,25 +379,21 @@ async def test_chat_completions_system_message(mock_get, async_client):
                 "index": 0,
                 "message": {
                     "role": "assistant",
-                    "content": "This is a helpful response about the test document."
+                    "content": "This is a helpful response about the test document.",
                 },
-                "finish_reason": "stop"
+                "finish_reason": "stop",
             }
         ],
-        "usage": {
-            "prompt_tokens": 25,
-            "completion_tokens": 12,
-            "total_tokens": 37
-        }
+        "usage": {"prompt_tokens": 25, "completion_tokens": 12, "total_tokens": 37},
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index some test documents
     index_request = {
         "index_name": "test_index",
-        "documents": [
-            {"text": "Document about machine learning algorithms."}
-        ]
+        "documents": [{"text": "Document about machine learning algorithms."}],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -358,19 +404,26 @@ async def test_chat_completions_system_message(mock_get, async_client):
         "index_name": "test_index",
         "model": "mock-model",
         "messages": [
-            {"role": "system", "content": "You are a helpful AI assistant specializing in machine learning."},
-            {"role": "user", "content": "Tell me about algorithms."}
+            {
+                "role": "system",
+                "content": "You are a helpful AI assistant specializing in machine learning.",
+            },
+            {"role": "user", "content": "Tell me about algorithms."},
         ],
-        "temperature": 0.5
+        "temperature": 0.5,
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
-    assert response_data["choices"][0]["message"]["content"] == "This is a helpful response about the test document."
+    assert (
+        response_data["choices"][0]["message"]["content"]
+        == "This is a helpful response about the test document."
+    )
     assert len(response_data["source_nodes"]) > 0
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -384,22 +437,29 @@ async def test_chat_completions_unsupported_message_role(mock_get, async_client)
 
     # Mock HTTPX response for passthrough LLM call
     mock_response = {"detail": "bad request format"}
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(400, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(400, json=mock_response)
+    )
 
     # Test request with unsupported role (should trigger passthrough)
     chat_request = {
         "model": "mock-model",
         "messages": [
-            {"role": "function", "content": "Function response", "name": "test_function"}
-        ]
+            {
+                "role": "function",
+                "content": "Function response",
+                "name": "test_function",
+            }
+        ],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 400
-    
+
     response_data = response.json()
     assert "bad request format" in response_data["detail"]
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -417,36 +477,46 @@ async def test_chat_completions_complex_user_content(mock_get, async_client):
         "object": "chat.completion",
         "created": int(time.time()),
         "model": "mock-model",
-        "choices": [{
-            "index": 0,
-            "message": {"role": "assistant", "content": "Complex content response"},
-            "finish_reason": "stop"
-        }]
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Complex content response"},
+                "finish_reason": "stop",
+            }
+        ],
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Test request with complex content (should trigger passthrough)
     chat_request = {
         "model": "mock-model",
         "messages": [
             {
-                "role": "user", 
+                "role": "user",
                 "content": [
                     {"type": "text", "text": "What's in this image?"},
-                    {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
-                ]
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "data:image/jpeg;base64,..."},
+                    },
+                ],
             }
-        ]
+        ],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
-    assert response_data["choices"][0]["message"]["content"] == "Complex content response"
+    assert (
+        response_data["choices"][0]["message"]["content"] == "Complex content response"
+    )
     # Should have source_nodes field but it should be None for passthrough requests
     assert response_data["source_nodes"] is None
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -468,20 +538,20 @@ async def test_chat_completions_developer_role(mock_get, async_client):
                 "index": 0,
                 "message": {
                     "role": "assistant",
-                    "content": "This is a helpful response about the test document."
+                    "content": "This is a helpful response about the test document.",
                 },
-                "finish_reason": "stop"
+                "finish_reason": "stop",
             }
-        ]
+        ],
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index some test documents
     index_request = {
         "index_name": "test_index",
-        "documents": [
-            {"text": "Technical documentation about APIs."}
-        ]
+        "documents": [{"text": "Technical documentation about APIs."}],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -493,15 +563,19 @@ async def test_chat_completions_developer_role(mock_get, async_client):
         "model": "mock-model",
         "messages": [
             {"role": "developer", "content": "Debug information: API call failed"},
-            {"role": "user", "content": "Help me understand the API."}
-        ]
+            {"role": "user", "content": "Help me understand the API."},
+        ],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
-    assert response_data["choices"][0]["message"]["content"] == "This is a helpful response about the test document."
+    assert (
+        response_data["choices"][0]["message"]["content"]
+        == "This is a helpful response about the test document."
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
@@ -516,7 +590,9 @@ async def test_chat_completions_with_top_k_parameter(mock_get, async_client):
 
     # Mock HTTPX response for Custom Inference API
     mock_response = {"result": "Response based on top documents."}
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index multiple test documents
     index_request = {
@@ -526,8 +602,8 @@ async def test_chat_completions_with_top_k_parameter(mock_get, async_client):
             {"text": "Second document about deep learning."},
             {"text": "Third document about neural networks."},
             {"text": "Fourth document about AI applications."},
-            {"text": "Fifth document about data science."}
-        ]
+            {"text": "Fifth document about data science."},
+        ],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -537,19 +613,18 @@ async def test_chat_completions_with_top_k_parameter(mock_get, async_client):
     chat_request = {
         "index_name": "test_index",
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "Tell me about machine learning."}
-        ],
-        "top_k": 3
+        "messages": [{"role": "user", "content": "Tell me about machine learning."}],
+        "top_k": 3,
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
     # Should return up to 3 source nodes based on top_k parameter
     assert len(response_data["source_nodes"]) <= 3
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -562,14 +637,14 @@ async def test_chat_completions_error_handling(mock_get, async_client):
     }
 
     # Mock HTTPX response with error
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(500, json={"error": "Internal server error"}))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(500, json={"error": "Internal server error"})
+    )
 
     # Index some test documents
     index_request = {
         "index_name": "test_index",
-        "documents": [
-            {"text": "Test document."}
-        ]
+        "documents": [{"text": "Test document."}],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -579,17 +654,16 @@ async def test_chat_completions_error_handling(mock_get, async_client):
     chat_request = {
         "index_name": "test_index",
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "Test question."}
-        ]
+        "messages": [{"role": "user", "content": "Test question."}],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 500
     assert "An unexpected error occurred" in response.json()["detail"]
 
+
 @pytest.mark.asyncio
-@respx.mock 
+@respx.mock
 @patch("requests.get")
 async def test_chat_completions_assistant_message_with_content(mock_get, async_client):
     """Test chat completion with assistant message that has content."""
@@ -601,28 +675,28 @@ async def test_chat_completions_assistant_message_with_content(mock_get, async_c
 
     # Mock HTTPX response for Custom Inference API
     mock_response = {
-            "id": "chatcmpl-test123",
-            "object": "chat.completion",
-            "model": "mock-model",
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": "This is a helpful response about the test document."
-                    },
-                    "finish_reason": "stop"
-                }
-            ]
-        }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+        "id": "chatcmpl-test123",
+        "object": "chat.completion",
+        "model": "mock-model",
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "This is a helpful response about the test document.",
+                },
+                "finish_reason": "stop",
+            }
+        ],
+    }
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index some test documents
     index_request = {
         "index_name": "test_index",
-        "documents": [
-            {"text": "Conversation about AI."}
-        ]
+        "documents": [{"text": "Conversation about AI."}],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -634,16 +708,23 @@ async def test_chat_completions_assistant_message_with_content(mock_get, async_c
         "model": "mock-model",
         "messages": [
             {"role": "user", "content": "Hello"},
-            {"role": "assistant", "content": "Hello! How can I help you?"},  # Assistant message with content
-            {"role": "user", "content": "Can you help me?"}
-        ]
+            {
+                "role": "assistant",
+                "content": "Hello! How can I help you?",
+            },  # Assistant message with content
+            {"role": "user", "content": "Can you help me?"},
+        ],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
-    assert response_data["choices"][0]["message"]["content"] == "This is a helpful response about the test document."
+    assert (
+        response_data["choices"][0]["message"]["content"]
+        == "This is a helpful response about the test document."
+    )
+
 
 @pytest.mark.asyncio
 async def test_chat_completions_metrics_tracking(async_client):
@@ -651,25 +732,24 @@ async def test_chat_completions_metrics_tracking(async_client):
     # Test successful request
     chat_request = {
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "Hello"}
-        ]
+        "messages": [{"role": "user", "content": "Hello"}],
     }
 
     # This will fail but should still track metrics
     await async_client.post("/v1/chat/completions", json=chat_request)
     # Should fail due to validation error or missing setup
-    
+
     # Check metrics endpoint
     metrics_response = await async_client.get("/metrics")
     assert metrics_response.status_code == 200
-    
+
     # Should have at least one chat request recorded (success or failure)
     metrics_text = metrics_response.text
     assert "rag_chat_requests_total" in metrics_text
     assert "rag_chat_latency" in metrics_text
 
-@pytest.mark.asyncio 
+
+@pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
 async def test_chat_completions_mixed_message_types(mock_get, async_client):
@@ -690,20 +770,20 @@ async def test_chat_completions_mixed_message_types(mock_get, async_client):
                 "index": 0,
                 "message": {
                     "role": "assistant",
-                    "content": "This is a helpful response about the test document."
+                    "content": "This is a helpful response about the test document.",
                 },
-                "finish_reason": "stop"
+                "finish_reason": "stop",
             }
-        ]
+        ],
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index some test documents
     index_request = {
         "index_name": "test_index",
-        "documents": [
-            {"text": "Technical documentation about software development."}
-        ]
+        "documents": [{"text": "Technical documentation about software development."}],
     }
 
     response = await async_client.post("/index", json=index_request)
@@ -718,17 +798,24 @@ async def test_chat_completions_mixed_message_types(mock_get, async_client):
             {"role": "user", "content": "I need help with coding."},
             {"role": "assistant", "content": "I'd be happy to help with coding."},
             {"role": "user", "content": "What about software development?"},
-            {"role": "developer", "content": "Debug: User asking about software development"}
-        ]
+            {
+                "role": "developer",
+                "content": "Debug: User asking about software development",
+            },
+        ],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
-    assert response_data["choices"][0]["message"]["content"] == "This is a helpful response about the test document."
+    assert (
+        response_data["choices"][0]["message"]["content"]
+        == "This is a helpful response about the test document."
+    )
     assert len(response_data["source_nodes"]) > 0
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -741,17 +828,17 @@ async def test_chat_completions_empty_messages_list(mock_get, async_client):
     }
 
     # Mock HTTPX response for passthrough LLM call (in case it gets that far)
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(400, json={"error": "Invalid request"}))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(400, json={"error": "Invalid request"})
+    )
 
-    chat_request = {
-        "model": "mock-model",
-        "messages": []
-    }
+    chat_request = {"model": "mock-model", "messages": []}
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 400
     assert "Invalid request" in response.json()["detail"]
 
+
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")
@@ -769,32 +856,34 @@ async def test_chat_completions_with_functions(mock_get, async_client):
         "object": "chat.completion",
         "created": int(time.time()),
         "model": "mock-model",
-        "choices": [{
-            "index": 0,
-            "message": {"role": "assistant", "content": "Function-enabled response"},
-            "finish_reason": "stop"
-        }]
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "Function-enabled response",
+                },
+                "finish_reason": "stop",
+            }
+        ],
     }
-    respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/chat/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Test request with functions (should trigger passthrough)
     chat_request = {
         "model": "mock-model",
-        "messages": [
-            {"role": "user", "content": "Use a function to help me"}
-        ],
-        "functions": [
-            {
-                "name": "test_function",
-                "description": "A test function"
-            }
-        ]
+        "messages": [{"role": "user", "content": "Use a function to help me"}],
+        "functions": [{"name": "test_function", "description": "A test function"}],
     }
 
     response = await async_client.post("/v1/chat/completions", json=chat_request)
     assert response.status_code == 200
-    
+
     response_data = response.json()
-    assert response_data["choices"][0]["message"]["content"] == "Function-enabled response"
+    assert (
+        response_data["choices"][0]["message"]["content"] == "Function-enabled response"
+    )
     # Should have source_nodes field but it should be None for passthrough requests
-    assert response_data["source_nodes"] is None
\ No newline at end of file
+    assert response_data["source_nodes"] is None
diff --git a/presets/ragengine/tests/api/test_main.py b/presets/ragengine/tests/api/test_main.py
index 381e685..5e6f550 100644
--- a/presets/ragengine/tests/api/test_main.py
+++ b/presets/ragengine/tests/api/test_main.py
@@ -12,47 +12,53 @@
 # limitations under the License.
 
 
-from unittest.mock import patch, ANY
-
-from llama_index.core.storage.index_store import SimpleIndexStore
-
-from ragengine.main import app, vector_store_handler, rag_ops
-from ragengine.config import DEFAULT_VECTOR_DB_PERSIST_DIR
-
+import json
 import os
+import re
+from unittest.mock import patch
 
-import pytest
 import httpx
+import pytest
 import respx
-import json
-import re
+
+from ragengine.config import DEFAULT_VECTOR_DB_PERSIST_DIR
 
 AUTO_GEN_DOC_ID_LEN = 64
 
+
 @pytest.mark.asyncio
 async def test_index_documents_success(async_client):
     request_data = {
         "index_name": "test_index",
         "documents": [
             {"text": "This is a test document"},
-            {"text": "Another test document"}
-        ]
+            {"text": "Another test document"},
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
     assert response.status_code == 200
     doc1, doc2 = response.json()
-    assert (doc1["text"] == "This is a test document")
+    assert doc1["text"] == "This is a test document"
     assert len(doc1["doc_id"]) == AUTO_GEN_DOC_ID_LEN
     assert not doc1["metadata"]
 
-    assert (doc2["text"] == "Another test document")
+    assert doc2["text"] == "Another test document"
     assert len(doc2["doc_id"]) == AUTO_GEN_DOC_ID_LEN
     assert not doc2["metadata"]
 
-    response = await async_client.get(f"/metrics")
+    response = await async_client.get("/metrics")
     assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
@@ -66,15 +72,17 @@ async def test_query_index_success(mock_get, async_client):
 
     # Mock HTTPX response for Custom Inference API
     mock_response = {"result": "This is the completion from the API"}
-    respx.post("http://localhost:5000/v1/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index Request
     request_data = {
         "index_name": "test_index",
         "documents": [
             {"text": "This is a test document"},
-            {"text": "Another test document"}
-        ]
+            {"text": "Another test document"},
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
@@ -85,27 +93,49 @@ async def test_query_index_success(mock_get, async_client):
         "index_name": "test_index",
         "query": "test query",
         "top_k": 1,
-        "llm_params": {"temperature": 0.7}
+        "llm_params": {"temperature": 0.7},
     }
 
     response = await async_client.post("/query", json=request_data)
     assert response.status_code == 200
-    assert response.json()["response"] == "{'result': 'This is the completion from the API'}"
+    assert (
+        response.json()["response"]
+        == "{'result': 'This is the completion from the API'}"
+    )
     assert len(response.json()["source_nodes"]) == 1
     assert response.json()["source_nodes"][0]["text"] == "This is a test document"
-    assert response.json()["source_nodes"][0]["score"] == pytest.approx(0.5354418754577637, rel=1e-6)
+    assert response.json()["source_nodes"][0]["score"] == pytest.approx(
+        0.5354418754577637, rel=1e-6
+    )
     assert response.json()["source_nodes"][0]["metadata"] == {}
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_query_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_query_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
 @patch("requests.get")  # Mock the requests.get call for fetching model metadata
 async def test_document_update_success(mock_get, async_client):
-    #Mock the response for the default model fetch
+    # Mock the response for the default model fetch
     mock_get.return_value.status_code = 200
     mock_get.return_value.json.return_value = {
         "data": [{"id": "mock-model", "max_model_len": 2048}]
@@ -113,15 +143,17 @@ async def test_document_update_success(mock_get, async_client):
 
     # Mock HTTPX response for Custom Inference API
     mock_response = {"result": "This is the completion from the API"}
-    respx.post("http://localhost:5000/v1/completions").mock(return_value=httpx.Response(200, json=mock_response))
+    respx.post("http://localhost:5000/v1/completions").mock(
+        return_value=httpx.Response(200, json=mock_response)
+    )
 
     # Index Request
     request_data = {
         "index_name": "test_update_index",
         "documents": [
             {"text": "This is a test document"},
-            {"text": "Another test document"}
-        ]
+            {"text": "Another test document"},
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
@@ -130,18 +162,25 @@ async def test_document_update_success(mock_get, async_client):
     assert doc2["doc_id"] != ""
 
     doc2["text"] = "This is an updated test document"
-    not_existing_doc = {"doc_id": "nonexistingdoc", "text": "This is a new test document"}
+    not_existing_doc = {
+        "doc_id": "nonexistingdoc",
+        "text": "This is a new test document",
+    }
     update_request_data = {
-        "documents": [
-            doc2,
-            not_existing_doc,
-            doc1
-        ],
+        "documents": [doc2, not_existing_doc, doc1],
     }
-    response = await async_client.post(f"/indexes/test_update_index/documents", json=update_request_data)
-    assert response.status_code == 200
-    assert response.json()["updated_documents"][0]["text"] == "This is an updated test document"
-    assert response.json()["not_found_documents"][0]["doc_id"] == not_existing_doc["doc_id"]
+    response = await async_client.post(
+        "/indexes/test_update_index/documents", json=update_request_data
+    )
+    assert response.status_code == 200
+    assert (
+        response.json()["updated_documents"][0]["text"]
+        == "This is an updated test document"
+    )
+    assert (
+        response.json()["not_found_documents"][0]["doc_id"]
+        == not_existing_doc["doc_id"]
+    )
     assert response.json()["unchanged_documents"][0]["text"] == doc1["text"]
 
     # Query Request
@@ -149,24 +188,56 @@ async def test_document_update_success(mock_get, async_client):
         "index_name": "test_update_index",
         "query": "updates test query",
         "top_k": 1,
-        "llm_params": {"temperature": 0.7}
+        "llm_params": {"temperature": 0.7},
     }
 
     response = await async_client.post("/query", json=request_data)
     assert response.status_code == 200
-    assert response.json()["response"] == "{'result': 'This is the completion from the API'}"
+    assert (
+        response.json()["response"]
+        == "{'result': 'This is the completion from the API'}"
+    )
     assert len(response.json()["source_nodes"]) == 1
-    assert response.json()["source_nodes"][0]["text"] == "This is an updated test document"
-    assert response.json()["source_nodes"][0]["score"] == pytest.approx(0.48061275482177734, rel=1e-6)
+    assert (
+        response.json()["source_nodes"][0]["text"] == "This is an updated test document"
+    )
+    assert response.json()["source_nodes"][0]["score"] == pytest.approx(
+        0.48061275482177734, rel=1e-6
+    )
     assert response.json()["source_nodes"][0]["metadata"] == {}
 
     assert respx.calls.call_count == 1
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_query_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_indexes_update_document_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_query_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_indexes_update_document_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 async def test_document_delete_success(async_client):
@@ -175,8 +246,8 @@ async def test_document_delete_success(async_client):
         "index_name": "test_delete_index",
         "documents": [
             {"text": "This is a test document"},
-            {"text": "Another test document"}
-        ]
+            {"text": "Another test document"},
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
@@ -184,26 +255,52 @@ async def test_document_delete_success(async_client):
     doc1, doc2 = response.json()
     assert doc2["doc_id"] != ""
 
-
     delete_request_data = {
         "doc_ids": [doc2["doc_id"], "nonexistingdoc"],
     }
-    response = await async_client.post(f"/indexes/test_delete_index/documents/delete", json=delete_request_data)
+    response = await async_client.post(
+        "/indexes/test_delete_index/documents/delete", json=delete_request_data
+    )
     assert response.status_code == 200
     assert response.json()["deleted_doc_ids"] == [doc2["doc_id"]]
     assert response.json()["not_found_doc_ids"] == ["nonexistingdoc"]
 
-    response = await async_client.get(f"/indexes/test_delete_index/documents")
+    response = await async_client.get("/indexes/test_delete_index/documents")
     assert response.status_code == 200
     assert response.json()["count"] == 1
     assert len(response.json()["documents"]) == 1
     assert response.json()["documents"][0]["text"] == "This is a test document"
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_indexes_delete_document_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_indexes_document_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_indexes_delete_document_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_indexes_document_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
@@ -241,15 +338,21 @@ async def test_reranker_and_query_with_index(mock_get, async_client):
     }
 
     # Mock HTTPX response for reranker and query API calls
-    reranker_mock_response = {"choices": [{"text": "Doc: 4, Relevance: 10\nDoc: 5, Relevance: 10"}]}
+    reranker_mock_response = {
+        "choices": [{"text": "Doc: 4, Relevance: 10\nDoc: 5, Relevance: 10"}]
+    }
     query_mock_response = {"choices": [{"text": "his is the completion from the API"}]}
 
     # Mock reranker response
-    respx.post("http://localhost:5000/v1/completions", content__contains="A list of documents is shown below") \
-        .mock(return_value=httpx.Response(200, json=reranker_mock_response))
+    respx.post(
+        "http://localhost:5000/v1/completions",
+        content__contains="A list of documents is shown below",
+    ).mock(return_value=httpx.Response(200, json=reranker_mock_response))
 
     # Mock query response
-    respx.post("http://localhost:5000/v1/completions").mock(return_value=httpx.Response(200, json=query_mock_response))
+    respx.post("http://localhost:5000/v1/completions").mock(
+        return_value=httpx.Response(200, json=query_mock_response)
+    )
 
     # Define input documents for indexing
     documents = [
@@ -258,13 +361,13 @@ async def test_reranker_and_query_with_index(mock_get, async_client):
         "The capital of France is beautiful.",
         "Have you ever visited Paris? It is a beautiful city where you can eat delicious food and see the Eiffel Tower. I really enjoyed all the cities in France, but its capital with the Eiffel Tower is my favorite city.",
         "I really enjoyed my trip to Paris, France. The city is beautiful and the food is delicious. I would love to visit again. "
-        "Such a great capital city."
+        "Such a great capital city.",
     ]
 
     # Indexing request payload
     index_request_payload = {
         "index_name": "test_index",
-        "documents": [{"text": doc.strip()} for doc in documents]
+        "documents": [{"text": doc.strip()} for doc in documents],
     }
 
     # Perform indexing
@@ -279,7 +382,7 @@ async def test_reranker_and_query_with_index(mock_get, async_client):
         "query": "what is the capital of france?",
         "top_k": 5,
         "llm_params": {"temperature": 0, "max_tokens": 2000},
-        "rerank_params": {"top_n": top_n}
+        "rerank_params": {"top_n": top_n},
     }
 
     # Perform query
@@ -292,11 +395,19 @@ async def test_reranker_and_query_with_index(mock_get, async_client):
 
     # Validate each source node in the query response
     expected_source_nodes = [
-        {"text": "Have you ever visited Paris? It is a beautiful city where you can eat delicious food and see the Eiffel Tower. I really enjoyed all the cities in France, but its capital with the Eiffel Tower is my favorite city.",
-         "score": 10.0, "metadata": {}, "doc_id": index_response[3]["doc_id"]},
-        {"text": "I really enjoyed my trip to Paris, France. The city is beautiful and the "
-                 "food is delicious. I would love to visit again. Such a great capital city.",
-         "score": 10.0, "metadata": {}, "doc_id": index_response[4]["doc_id"]},
+        {
+            "text": "Have you ever visited Paris? It is a beautiful city where you can eat delicious food and see the Eiffel Tower. I really enjoyed all the cities in France, but its capital with the Eiffel Tower is my favorite city.",
+            "score": 10.0,
+            "metadata": {},
+            "doc_id": index_response[3]["doc_id"],
+        },
+        {
+            "text": "I really enjoyed my trip to Paris, France. The city is beautiful and the "
+            "food is delicious. I would love to visit again. Such a great capital city.",
+            "score": 10.0,
+            "metadata": {},
+            "doc_id": index_response[4]["doc_id"],
+        },
     ]
     for i, expected_node in enumerate(expected_source_nodes):
         actual_node = query_response["source_nodes"][i]
@@ -308,10 +419,27 @@ async def test_reranker_and_query_with_index(mock_get, async_client):
     # Ensure HTTPX requests were made
     assert respx.calls.call_count == 2  # One for rerank, one for query completion
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_query_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_query_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 @respx.mock
@@ -330,8 +458,10 @@ async def test_reranker_failed_and_query_with_index(mock_get, async_client):
     reranker_mock_response = {"choices": [{"text": "Empty Response"}]}
 
     # Mock reranker response
-    respx.post("http://localhost:5000/v1/completions", content__contains="A list of documents is shown below") \
-        .mock(return_value=httpx.Response(200, json=reranker_mock_response))
+    respx.post(
+        "http://localhost:5000/v1/completions",
+        content__contains="A list of documents is shown below",
+    ).mock(return_value=httpx.Response(200, json=reranker_mock_response))
 
     # Define input documents for indexing
     documents = [
@@ -340,13 +470,13 @@ async def test_reranker_failed_and_query_with_index(mock_get, async_client):
         "The capital of France is beautiful.",
         "Have you ever visited Paris? It is a beautiful city where you can eat delicious food and see the Eiffel Tower. I really enjoyed all the cities in France, but its capital with the Eiffel Tower is my favorite city.",
         "I really enjoyed my trip to Paris, France. The city is beautiful and the food is delicious. I would love to visit again. "
-        "Such a great capital city."
+        "Such a great capital city.",
     ]
 
     # Indexing request payload
     index_request_payload = {
         "index_name": "test_index",
-        "documents": [{"text": doc.strip()} for doc in documents]
+        "documents": [{"text": doc.strip()} for doc in documents],
     }
 
     # Perform indexing
@@ -360,21 +490,41 @@ async def test_reranker_failed_and_query_with_index(mock_get, async_client):
         "query": "what is the capital of france?",
         "top_k": 5,
         "llm_params": {"temperature": 0, "max_tokens": 2000},
-        "rerank_params": {"top_n": top_n}
+        "rerank_params": {"top_n": top_n},
     }
 
     # Perform query
     response = await async_client.post("/query", json=query_request_payload)
     assert response.status_code == 422
-    assert response.content == b'{"detail":"Rerank operation failed: Invalid response from LLM. This feature is experimental."}'
+    assert (
+        response.content
+        == b'{"detail":"Rerank operation failed: Invalid response from LLM. This feature is experimental."}'
+    )
 
     # Ensure HTTPX requests were made
     assert respx.calls.call_count == 1  # One for rerank
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_query_requests_total{status="failure"} ([1-9]\d*).0', response.text)) == 1
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_query_requests_total{status="failure"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 async def test_query_index_failure(async_client):
@@ -383,16 +533,25 @@ async def test_query_index_failure(async_client):
         "index_name": "non_existent_index",  # Use an index name that doesn't exist
         "query": "test query",
         "top_k": 1,
-        "llm_params": {"temperature": 0.7}
+        "llm_params": {"temperature": 0.7},
     }
 
     response = await async_client.post("/query", json=request_data)
     assert response.status_code == 404
     assert response.json()["detail"] == "No such index: 'non_existent_index' exists."
 
-    response = await async_client.get(f"/metrics")
+    response = await async_client.get("/metrics")
     assert response.status_code == 200
-    assert len(re.findall(r'rag_query_requests_total{status="failure"} ([1-9]\d*).0', response.text)) == 1
+    assert (
+        len(
+            re.findall(
+                r'rag_query_requests_total{status="failure"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 async def test_list_documents_in_index_success(async_client):
@@ -402,14 +561,14 @@ async def test_list_documents_in_index_success(async_client):
     response = await async_client.get(f"/indexes/{index_name}/documents")
 
     assert response.status_code == 404
-    assert response.json() == {'detail': "No such index: 'test_index' exists."}
+    assert response.json() == {"detail": "No such index: 'test_index' exists."}
 
     request_data = {
         "index_name": index_name,
         "documents": [
             {"text": "This is a test document"},
-            {"text": "Another test document"}
-        ]
+            {"text": "Another test document"},
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
@@ -424,11 +583,18 @@ async def test_list_documents_in_index_success(async_client):
     # Ensure documents exist correctly in the specific index
     assert response_json["count"] == 2
     assert len(response_json["documents"]) == 2
-    assert all(((item['doc_id'] == doc1['doc_id'] and item['text'] == doc1['text']) or
-               (item['doc_id'] == doc2['doc_id'] and item['text'] == doc2['text'])) for item in response_json["documents"])
-    
-    assert ({item["text"] for item in response_json["documents"]}
-            == {item["text"] for item in request_data["documents"]})
+    assert all(
+        (
+            (item["doc_id"] == doc1["doc_id"] and item["text"] == doc1["text"])
+            or (item["doc_id"] == doc2["doc_id"] and item["text"] == doc2["text"])
+        )
+        for item in response_json["documents"]
+    )
+
+    assert {item["text"] for item in response_json["documents"]} == {
+        item["text"] for item in request_data["documents"]
+    }
+
 
 @pytest.mark.asyncio
 async def test_list_documents_with_metadata_filter_success(async_client):
@@ -438,14 +604,20 @@ async def test_list_documents_with_metadata_filter_success(async_client):
     response = await async_client.get(f"/indexes/{index_name}/documents")
 
     assert response.status_code == 404
-    assert response.json() == {'detail': "No such index: 'test_index' exists."}
+    assert response.json() == {"detail": "No such index: 'test_index' exists."}
 
     request_data = {
         "index_name": index_name,
         "documents": [
-            {"text": "This is a test document", "metadata": {"filename": "test.txt", "branch": "main"}},
-            {"text": "Another test document", "metadata": {"filename": "main.py", "branch": "main"}}
-        ]
+            {
+                "text": "This is a test document",
+                "metadata": {"filename": "test.txt", "branch": "main"},
+            },
+            {
+                "text": "Another test document",
+                "metadata": {"filename": "main.py", "branch": "main"},
+            },
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
@@ -455,7 +627,9 @@ async def test_list_documents_with_metadata_filter_success(async_client):
     filters = {
         "filename": "test.txt",
     }
-    response = await async_client.get(f"/indexes/{index_name}/documents?metadata_filter={json.dumps(filters)}")
+    response = await async_client.get(
+        f"/indexes/{index_name}/documents?metadata_filter={json.dumps(filters)}"
+    )
     assert response.status_code == 200
     response_json = response.json()
 
@@ -464,6 +638,7 @@ async def test_list_documents_with_metadata_filter_success(async_client):
     assert len(response_json["documents"]) == 1
     assert response_json["documents"][0]["text"] == "This is a test document"
 
+
 @pytest.mark.asyncio
 async def test_list_documents_with_metadata_filter_failure(async_client):
     index_name = "test_index"
@@ -472,22 +647,31 @@ async def test_list_documents_with_metadata_filter_failure(async_client):
     response = await async_client.get(f"/indexes/{index_name}/documents")
 
     assert response.status_code == 404
-    assert response.json() == {'detail': "No such index: 'test_index' exists."}
+    assert response.json() == {"detail": "No such index: 'test_index' exists."}
 
     request_data = {
         "index_name": index_name,
         "documents": [
-            {"text": "This is a test document", "metadata": {"filename": "test.txt", "branch": "main"}},
-            {"text": "Another test document", "metadata": {"filename": "main.py", "branch": "main"}}
-        ]
+            {
+                "text": "This is a test document",
+                "metadata": {"filename": "test.txt", "branch": "main"},
+            },
+            {
+                "text": "Another test document",
+                "metadata": {"filename": "main.py", "branch": "main"},
+            },
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
     assert response.status_code == 200
 
-    response = await async_client.get(f"/indexes/{index_name}/documents?metadata_filter=invalidjsonstring")
+    response = await async_client.get(
+        f"/indexes/{index_name}/documents?metadata_filter=invalidjsonstring"
+    )
     assert response.status_code == 400
 
+
 @pytest.mark.asyncio
 async def test_persist_documents(async_client):
     index_name = "test_index"
@@ -496,14 +680,14 @@ async def test_persist_documents(async_client):
     response = await async_client.get(f"/indexes/{index_name}/documents")
 
     assert response.status_code == 404
-    assert response.json() == {'detail': "No such index: 'test_index' exists."}
+    assert response.json() == {"detail": "No such index: 'test_index' exists."}
 
     request_data = {
         "index_name": index_name,
         "documents": [
             {"text": "This is a test document"},
-            {"text": "Another test document"}
-        ]
+            {"text": "Another test document"},
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
@@ -513,7 +697,9 @@ async def test_persist_documents(async_client):
     response = await async_client.post(f"/persist/{index_name}")
     assert response.status_code == 200
     response_json = response.json()
-    assert response_json == {"message": f"Successfully persisted index {index_name} to {DEFAULT_VECTOR_DB_PERSIST_DIR}/{index_name}."}
+    assert response_json == {
+        "message": f"Successfully persisted index {index_name} to {DEFAULT_VECTOR_DB_PERSIST_DIR}/{index_name}."
+    }
     assert os.path.exists(os.path.join(DEFAULT_VECTOR_DB_PERSIST_DIR, index_name))
 
     # Persist documents for the specific index at a custom path
@@ -521,28 +707,59 @@ async def test_persist_documents(async_client):
     response = await async_client.post(f"/persist/{index_name}?path={custom_path}")
     assert response.status_code == 200
     response_json = response.json()
-    assert response_json == {"message": f"Successfully persisted index {index_name} to {custom_path}."}
+    assert response_json == {
+        "message": f"Successfully persisted index {index_name} to {custom_path}."
+    }
     assert os.path.exists(custom_path)
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_indexes_document_requests_total{status="failure"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_persist_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_indexes_document_requests_total{status="failure"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_persist_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 async def test_load_documents(async_client):
     index_name = "test_index"
-    response = await async_client.post(f"/load/{index_name}?path={DEFAULT_VECTOR_DB_PERSIST_DIR}/{index_name}")
+    response = await async_client.post(
+        f"/load/{index_name}?path={DEFAULT_VECTOR_DB_PERSIST_DIR}/{index_name}"
+    )
 
     assert response.status_code == 200
-    assert response.json() == {'message': 'Successfully loaded index test_index from storage/test_index.'}
+    assert response.json() == {
+        "message": "Successfully loaded index test_index from storage/test_index."
+    }
 
-    response = await async_client.get(f"/indexes")
+    response = await async_client.get("/indexes")
     assert response.status_code == 200
     assert response.json() == [index_name]
 
-    response = await async_client.get(f"/indexes/test_index/documents")
+    response = await async_client.get("/indexes/test_index/documents")
     assert response.status_code == 200
     response_data = response.json()
 
@@ -551,11 +768,35 @@ async def test_load_documents(async_client):
     assert response_data["documents"][0]["text"] == "This is a test document"
     assert response_data["documents"][1]["text"] == "Another test document"
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_load_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_indexes_document_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_indexes_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_load_requests_total{status="success"} ([1-9]\d*).0', response.text
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_indexes_document_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_indexes_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+
 
 @pytest.mark.asyncio
 async def test_delete_index(async_client):
@@ -565,14 +806,14 @@ async def test_delete_index(async_client):
     response = await async_client.get(f"/indexes/{index_name}/documents")
 
     assert response.status_code == 404
-    assert response.json() == {'detail': "No such index: 'test_index' exists."}
+    assert response.json() == {"detail": "No such index: 'test_index' exists."}
 
     request_data = {
         "index_name": index_name,
         "documents": [
             {"text": "This is a test document"},
-            {"text": "Another test document"}
-        ]
+            {"text": "Another test document"},
+        ],
     }
 
     response = await async_client.post("/index", json=request_data)
@@ -588,13 +829,38 @@ async def test_delete_index(async_client):
     response = await async_client.get(f"/indexes/{index_name}/documents")
 
     assert response.status_code == 404
-    assert response.json() == {'detail': "No such index: 'test_index' exists."}
+    assert response.json() == {"detail": "No such index: 'test_index' exists."}
+
+    response = await async_client.get("/metrics")
+    assert response.status_code == 200
+    assert (
+        len(
+            re.findall(
+                r'rag_indexes_document_requests_total{status="failure"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_delete_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
+    assert (
+        len(
+            re.findall(
+                r'rag_index_requests_total{status="success"} ([1-9]\d*).0',
+                response.text,
+            )
+        )
+        == 1
+    )
 
-    response = await async_client.get(f"/metrics")
-    assert response.status_code == 200
-    assert len(re.findall(r'rag_indexes_document_requests_total{status="failure"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_delete_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
-    assert len(re.findall(r'rag_index_requests_total{status="success"} ([1-9]\d*).0', response.text)) == 1
 
 """
 Example of a live query test. This test is currently commented out as it requires a valid 
@@ -626,4 +892,4 @@ def test_live_query_test():
 
     response = client.post("/query", json=request_data)
     assert response.status_code == 200
-"""
\ No newline at end of file
+"""
diff --git a/presets/ragengine/tests/vector_store/conftest.py b/presets/ragengine/tests/vector_store/conftest.py
index 4a18fed..8461658 100644
--- a/presets/ragengine/tests/vector_store/conftest.py
+++ b/presets/ragengine/tests/vector_store/conftest.py
@@ -12,10 +12,12 @@
 # limitations under the License.
 
 
-import sys
 import os
+import sys
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # Force CPU-only execution for testing
-os.environ["OMP_NUM_THREADS"] = "1" # Force single-threaded for testing to prevent segfault while loading embedding model
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force CPU-only execution for testing
+os.environ["OMP_NUM_THREADS"] = (
+    "1"  # Force single-threaded for testing to prevent segfault while loading embedding model
+)
 os.environ["MKL_NUM_THREADS"] = "1"  # Force MKL to use a single thread
diff --git a/presets/ragengine/tests/vector_store/test_base_store.py b/presets/ragengine/tests/vector_store/test_base_store.py
index 1e8274d..e7d5261 100644
--- a/presets/ragengine/tests/vector_store/test_base_store.py
+++ b/presets/ragengine/tests/vector_store/test_base_store.py
@@ -13,22 +13,27 @@
 
 
 import os
-from unittest.mock import patch
-import pytest
 from abc import ABC, abstractmethod
+from unittest.mock import patch
 
-from ragengine.vector_store.base import BaseVectorStore
-from ragengine.models import Document
-from ragengine.embedding.huggingface_local_embedding import LocalHuggingFaceEmbedding
-from ragengine.config import (LOCAL_EMBEDDING_MODEL_ID, LLM_INFERENCE_URL,
-                              LLM_ACCESS_SECRET, DEFAULT_VECTOR_DB_PERSIST_DIR)
 import httpx
+import pytest
 import respx
 
+from ragengine.config import (
+    DEFAULT_VECTOR_DB_PERSIST_DIR,
+    LLM_INFERENCE_URL,
+    LOCAL_EMBEDDING_MODEL_ID,
+)
+from ragengine.embedding.huggingface_local_embedding import LocalHuggingFaceEmbedding
+from ragengine.models import Document
+from ragengine.vector_store.base import BaseVectorStore
+
+
 class BaseVectorStoreTest(ABC):
     """Base class for vector store tests that defines the test structure."""
-    
-    @pytest.fixture(scope='session')
+
+    @pytest.fixture(scope="session")
     def init_embed_manager(self):
         return LocalHuggingFaceEmbedding(LOCAL_EMBEDDING_MODEL_ID)
 
@@ -37,7 +42,7 @@ class BaseVectorStoreTest(ABC):
     def vector_store_manager(self, init_embed_manager):
         """Each implementation must provide its own vector store manager."""
         pass
-    
+
     @property
     @abstractmethod
     def expected_query_score(self):
@@ -49,14 +54,16 @@ class BaseVectorStoreTest(ABC):
         first_doc_text, second_doc_text = "First document", "Second document"
         documents = [
             Document(text=first_doc_text, metadata={"type": "text"}),
-            Document(text=second_doc_text, metadata={"type": "text"})
+            Document(text=second_doc_text, metadata={"type": "text"}),
         ]
-        
+
         doc_ids = await vector_store_manager.index_documents("test_index", documents)
-        
+
         assert len(doc_ids) == 2
-        assert set(doc_ids) == {BaseVectorStore.generate_doc_id(first_doc_text),
-                                BaseVectorStore.generate_doc_id(second_doc_text)}
+        assert set(doc_ids) == {
+            BaseVectorStore.generate_doc_id(first_doc_text),
+            BaseVectorStore.generate_doc_id(second_doc_text),
+        }
 
     @pytest.mark.asyncio
     async def test_index_documents_isolation(self, vector_store_manager):
@@ -90,33 +97,54 @@ class BaseVectorStoreTest(ABC):
         }
 
         mock_response = {"result": "This is the completion from the API"}
-        respx.post(LLM_INFERENCE_URL).mock(return_value=httpx.Response(200, json=mock_response))
+        respx.post(LLM_INFERENCE_URL).mock(
+            return_value=httpx.Response(200, json=mock_response)
+        )
 
         documents = [
             Document(text="First document", metadata={"type": "text"}),
-            Document(text="Second document", metadata={"type": "text"})
+            Document(text="Second document", metadata={"type": "text"}),
         ]
-        index_doc_resp = await vector_store_manager.index_documents("test_index", documents)
+        index_doc_resp = await vector_store_manager.index_documents(
+            "test_index", documents
+        )
 
         params = {"temperature": 0.7}
-        query_result = await vector_store_manager.query("test_index", "First", top_k=1,
-                                                  llm_params=params, rerank_params={})
+        query_result = await vector_store_manager.query(
+            "test_index", "First", top_k=1, llm_params=params, rerank_params={}
+        )
 
         assert query_result is not None
-        assert query_result["response"] == "{'result': 'This is the completion from the API'}"
+        assert (
+            query_result["response"]
+            == "{'result': 'This is the completion from the API'}"
+        )
         assert query_result["source_nodes"][0]["text"] == "First document"
-        assert query_result["source_nodes"][0]["score"] == pytest.approx(self.expected_query_score, rel=1e-6)
+        assert query_result["source_nodes"][0]["score"] == pytest.approx(
+            self.expected_query_score, rel=1e-6
+        )
         assert query_result["source_nodes"][0]["doc_id"] == index_doc_resp[0]
-        assert respx.calls.call_count == 1  # Ensure only one LLM inference request was made
-    
+        assert (
+            respx.calls.call_count == 1
+        )  # Ensure only one LLM inference request was made
+
     @pytest.mark.asyncio
     @respx.mock
     @patch("requests.get")
     async def test_chat_completions(self, mock_get, vector_store_manager, monkeypatch):
-        import ragengine.inference.inference
         import ragengine.config
-        monkeypatch.setattr(ragengine.config, "LLM_INFERENCE_URL", "http://localhost:5000/v1/chat/completions")
-        monkeypatch.setattr(ragengine.inference.inference, "LLM_INFERENCE_URL", "http://localhost:5000/v1/chat/completions")
+        import ragengine.inference.inference
+
+        monkeypatch.setattr(
+            ragengine.config,
+            "LLM_INFERENCE_URL",
+            "http://localhost:5000/v1/chat/completions",
+        )
+        monkeypatch.setattr(
+            ragengine.inference.inference,
+            "LLM_INFERENCE_URL",
+            "http://localhost:5000/v1/chat/completions",
+        )
         mock_get.return_value.status_code = 200
         mock_get.return_value.json.return_value = {
             "data": [{"id": "mock-model", "max_model_len": 2048}]
@@ -131,19 +159,23 @@ class BaseVectorStoreTest(ABC):
                     "index": 0,
                     "message": {
                         "role": "assistant",
-                        "content": "This is a helpful response about the test document."
+                        "content": "This is a helpful response about the test document.",
                     },
-                    "finish_reason": "stop"
+                    "finish_reason": "stop",
                 }
-            ]
+            ],
         }
-        respx.post("http://localhost:5000/v1/chat/completions").mock(return_value=httpx.Response(200, json=mock_response))
+        respx.post("http://localhost:5000/v1/chat/completions").mock(
+            return_value=httpx.Response(200, json=mock_response)
+        )
 
         documents = [
             Document(text="First document", metadata={"type": "text"}),
-            Document(text="Second document", metadata={"type": "text"})
+            Document(text="Second document", metadata={"type": "text"}),
         ]
-        index_doc_resp = await vector_store_manager.index_documents("test_index", documents)
+        index_doc_resp = await vector_store_manager.index_documents(
+            "test_index", documents
+        )
 
         chat_results = await vector_store_manager.chat_completion(
             {
@@ -172,9 +204,14 @@ class BaseVectorStoreTest(ABC):
         assert chat_results.choices[0].finish_reason == "stop"
         assert chat_results.choices[0].index == 0
         assert chat_results.choices[0].message.role == "assistant"
-        assert chat_results.choices[0].message.content == "This is a helpful response about the test document."
+        assert (
+            chat_results.choices[0].message.content
+            == "This is a helpful response about the test document."
+        )
 
-        assert respx.calls.call_count == 1  # Ensure only one LLM inference request was made
+        assert (
+            respx.calls.call_count == 1
+        )  # Ensure only one LLM inference request was made
 
     @pytest.mark.asyncio
     async def test_add_document(self, vector_store_manager):
@@ -184,37 +221,58 @@ class BaseVectorStoreTest(ABC):
         new_document = [Document(text="Fourth document", metadata={"type": "text"})]
         await vector_store_manager.index_documents("test_index", new_document)
 
-        assert await vector_store_manager.document_exists("test_index", new_document[0],
-                                                    BaseVectorStore.generate_doc_id("Fourth document"))
-    
+        assert await vector_store_manager.document_exists(
+            "test_index",
+            new_document[0],
+            BaseVectorStore.generate_doc_id("Fourth document"),
+        )
+
     @pytest.mark.asyncio
     async def test_add_code_documents_with_code_splitting(self, vector_store_manager):
         sample_go_code = """package main
 func main() {}"""
-        documents = [Document(text=sample_go_code, metadata={"split_type": "code", "language": "go"})]
-        result = await vector_store_manager.index_documents("test_code_index", documents)
+        documents = [
+            Document(
+                text=sample_go_code, metadata={"split_type": "code", "language": "go"}
+            )
+        ]
+        result = await vector_store_manager.index_documents(
+            "test_code_index", documents
+        )
+
+        assert await vector_store_manager.document_exists(
+            "test_code_index", documents[0], result[0]
+        )
 
-        assert await vector_store_manager.document_exists("test_code_index", documents[0],
-                                                    result[0])
-        
-        all_docs = await vector_store_manager.list_documents_in_index("test_code_index", limit=10, offset=0)
+        all_docs = await vector_store_manager.list_documents_in_index(
+            "test_code_index", limit=10, offset=0
+        )
         assert len(all_docs) == 1
-        assert all_docs[0]['text'] == sample_go_code
+        assert all_docs[0]["text"] == sample_go_code
 
         try:
             # Attempt to index a document with no language
-            unsupported_doc = Document(text="Unsupported language code", metadata={"split_type": "code"})
-            await vector_store_manager.index_documents("test_code_index", [unsupported_doc])
+            unsupported_doc = Document(
+                text="Unsupported language code", metadata={"split_type": "code"}
+            )
+            await vector_store_manager.index_documents(
+                "test_code_index", [unsupported_doc]
+            )
             assert False, "Expected ValueError for unsupported language"
         except ValueError as e:
             assert str(e) == "Language not specified in node metadata."
 
         try:
             # Attempt to index a document with an invalid language
-            unsupported_doc = Document(text="Unsupported language code", metadata={"split_type": "code", "language": "invalid"})
-            await vector_store_manager.index_documents("test_code_index", [unsupported_doc])
+            unsupported_doc = Document(
+                text="Unsupported language code",
+                metadata={"split_type": "code", "language": "invalid"},
+            )
+            await vector_store_manager.index_documents(
+                "test_code_index", [unsupported_doc]
+            )
             assert False, "Expected ValueError for unsupported language"
-        except LookupError as e:
+        except LookupError:
             assert True
 
     @pytest.mark.asyncio
@@ -227,36 +285,67 @@ func main() {}"""
         }
 
         mock_response = {"result": "This is the completion from the API"}
-        respx.post(LLM_INFERENCE_URL).mock(return_value=httpx.Response(200, json=mock_response))
+        respx.post(LLM_INFERENCE_URL).mock(
+            return_value=httpx.Response(200, json=mock_response)
+        )
 
         documents = [Document(text="Fifth document", metadata={"type": "text"})]
         ids = await vector_store_manager.index_documents("test_index", documents)
 
         result = await vector_store_manager.update_documents(
-            "test_index", [Document(doc_id=ids[0], text="Updated Fifth document", metadata={"type": "text"})]
+            "test_index",
+            [
+                Document(
+                    doc_id=ids[0],
+                    text="Updated Fifth document",
+                    metadata={"type": "text"},
+                )
+            ],
         )
         assert result["updated_documents"][0].doc_id == ids[0]
 
         # Check preexisting unchanged document case
         result = await vector_store_manager.update_documents(
-            "test_index", [Document(doc_id=ids[0], text="Updated Fifth document", metadata={"type": "text"})]
+            "test_index",
+            [
+                Document(
+                    doc_id=ids[0],
+                    text="Updated Fifth document",
+                    metadata={"type": "text"},
+                )
+            ],
         )
         assert result["unchanged_documents"][0].doc_id == ids[0]
 
-        assert await vector_store_manager.document_exists("test_index", Document(text="Updated Fifth document", metadata={"type": "text"}),
-                                                    ids[0])
+        assert await vector_store_manager.document_exists(
+            "test_index",
+            Document(text="Updated Fifth document", metadata={"type": "text"}),
+            ids[0],
+        )
 
         # Check if the document was updated
-        result = await vector_store_manager.query("test_index", "Updated Fifth document", top_k=1,
-                                              llm_params={}, rerank_params={})
+        result = await vector_store_manager.query(
+            "test_index",
+            "Updated Fifth document",
+            top_k=1,
+            llm_params={},
+            rerank_params={},
+        )
         assert result["source_nodes"][0]["text"] == "Updated Fifth document"
 
         # Check documents not found case
         result = await vector_store_manager.update_documents(
-            "test_index", [Document(doc_id="baddocid", text="Updated Fifth document", metadata={"type": "text"})]
+            "test_index",
+            [
+                Document(
+                    doc_id="baddocid",
+                    text="Updated Fifth document",
+                    metadata={"type": "text"},
+                )
+            ],
         )
         assert result["not_found_documents"][0].doc_id == "baddocid"
-    
+
     @pytest.mark.asyncio
     @respx.mock
     @patch("requests.get")
@@ -267,11 +356,12 @@ func main() {}"""
         }
 
         mock_response = {"result": "This is the completion from the API"}
-        respx.post(LLM_INFERENCE_URL).mock(return_value=httpx.Response(200, json=mock_response))
+        respx.post(LLM_INFERENCE_URL).mock(
+            return_value=httpx.Response(200, json=mock_response)
+        )
 
         documents = [
-            Document(text=f"Document {i}", metadata={"type": "text"})
-            for i in range(10)
+            Document(text=f"Document {i}", metadata={"type": "text"}) for i in range(10)
         ]
         ids = await vector_store_manager.index_documents("test_index", documents)
 
@@ -285,21 +375,24 @@ func main() {}"""
     @pytest.mark.asyncio
     async def test_add_document_on_existing_index(self, vector_store_manager):
         # Create index with single doc
-        await vector_store_manager.index_documents("test_add_index", [Document(text=f"Initial Doc", metadata={"type": "text"})])
+        await vector_store_manager.index_documents(
+            "test_add_index", [Document(text="Initial Doc", metadata={"type": "text"})]
+        )
 
         documents = [
-            Document(text=f"Document {i}", metadata={"type": "text"})
-            for i in range(10)
+            Document(text=f"Document {i}", metadata={"type": "text"}) for i in range(10)
         ]
         # load documents into the same index
         ids = await vector_store_manager.index_documents("test_add_index", documents)
 
         # Check if the documents were indexed correctly
-        all_docs = await vector_store_manager.list_documents_in_index("test_add_index", limit=10, offset=1)
+        all_docs = await vector_store_manager.list_documents_in_index(
+            "test_add_index", limit=10, offset=1
+        )
         print(all_docs)
 
         # Validate id's from index_documents match the expected ids
-        assert all(doc['doc_id'] == ids[idx] for idx, doc in enumerate(all_docs))
+        assert all(doc["doc_id"] == ids[idx] for idx, doc in enumerate(all_docs))
 
     @pytest.mark.asyncio
     async def test_persist_index(self, vector_store_manager):
@@ -307,19 +400,21 @@ func main() {}"""
         await vector_store_manager.index_documents("test_index", documents)
         await vector_store_manager.persist("test_index", DEFAULT_VECTOR_DB_PERSIST_DIR)
         assert os.path.exists(DEFAULT_VECTOR_DB_PERSIST_DIR)
-    
+
     @pytest.mark.asyncio
     async def test_delete_index(self, vector_store_manager):
-        one_mb = b'\x00' * 1024
-        documents = [Document(text=one_mb.decode(), metadata={"type": "text"}) for _ in range(10)]
+        one_mb = b"\x00" * 1024
+        documents = [
+            Document(text=one_mb.decode(), metadata={"type": "text"}) for _ in range(10)
+        ]
         await vector_store_manager.index_documents("test_index", documents)
-        
+
         # Ensure index exists before deletion
         indexes = vector_store_manager.list_indexes()
         assert "test_index" in indexes
 
         await vector_store_manager.delete_index("test_index")
-        
+
         # Ensure index is deleted
         indexes = vector_store_manager.list_indexes()
         assert "test_index" not in indexes
@@ -330,79 +425,114 @@ func main() {}"""
         index_name = "test_index"
         # Create multiple documents
         documents = [
-            Document(text=f"Document {i}", metadata={"type": "text"})
-            for i in range(10)
+            Document(text=f"Document {i}", metadata={"type": "text"}) for i in range(10)
         ]
-        
+
         await vector_store_manager.index_documents(index_name, documents)
 
         # 1. Offset 0, Limit 5 (Basic Case)
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=0)
+        result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=5, offset=0
+        )
         assert len(result) == 5
 
         # 2. Offset 5, Limit 5 (Next Batch)
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=5)
+        result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=5, offset=5
+        )
         assert len(result) == 5
 
         # 3. Offset at max (Empty Case)
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=10)
+        result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=5, offset=10
+        )
         assert index_name not in result or len(result) == 0
 
         # 4. Limit larger than available docs
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=15, offset=0)
+        result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=15, offset=0
+        )
         assert len(result) == 10  # Should return only available docs
 
         # 5. Limit exactly matches available docs
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=10, offset=0)
+        result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=10, offset=0
+        )
         assert len(result) == 10
 
         # 6. Limit of 1 (Single-Doc Retrieval)
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=1, offset=0)
+        result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=1, offset=0
+        )
         assert len(result) == 1
 
         # 7. max_text_length truncation check
-        truncated_result = await vector_store_manager.list_documents_in_index(index_name, limit=1, offset=0, max_text_length=5)
-        assert len(next(iter(truncated_result))['text']) == 5  # Ensure truncation
+        truncated_result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=1, offset=0, max_text_length=5
+        )
+        assert len(next(iter(truncated_result))["text"]) == 5  # Ensure truncation
 
         # 8. max_text_length is None (Full text should return)
-        full_text_result = await vector_store_manager.list_documents_in_index(index_name, limit=1, offset=0, max_text_length=None)
-        assert "Document" in next(iter(full_text_result))['text']  # Ensure no truncation
-    
+        full_text_result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=1, offset=0, max_text_length=None
+        )
+        assert (
+            "Document" in next(iter(full_text_result))["text"]
+        )  # Ensure no truncation
+
     @pytest.mark.asyncio
     async def test_list_documents_with_filter_index(self, vector_store_manager):
         """Test various list document scenarios with different limit and offset values."""
         index_name = "test_index"
         # Create multiple documents
         documents = [
-            Document(text=f"Document {i}", metadata={"type": "text", "filename": f"file_{i}", "branch": "main"})
+            Document(
+                text=f"Document {i}",
+                metadata={"type": "text", "filename": f"file_{i}", "branch": "main"},
+            )
             for i in range(10)
         ]
-        
+
         await vector_store_manager.index_documents(index_name, documents)
-        
+
         # single filter
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=0, metadata_filter={"filename": "file_1"})
+        result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=5, offset=0, metadata_filter={"filename": "file_1"}
+        )
         assert len(result) == 1
-        assert result[0]['metadata']['filename'] == "file_1"
+        assert result[0]["metadata"]["filename"] == "file_1"
 
-        
-        first_five_results = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=0, metadata_filter={"branch": "main"})
+        first_five_results = await vector_store_manager.list_documents_in_index(
+            index_name, limit=5, offset=0, metadata_filter={"branch": "main"}
+        )
         assert len(first_five_results) == 5
-        assert all(doc['metadata']['branch'] == "main" for doc in first_five_results)
+        assert all(doc["metadata"]["branch"] == "main" for doc in first_five_results)
 
-        second_five_results = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=5, metadata_filter={"branch": "main"})
+        second_five_results = await vector_store_manager.list_documents_in_index(
+            index_name, limit=5, offset=5, metadata_filter={"branch": "main"}
+        )
         assert len(second_five_results) == 5
-        assert all(doc['metadata']['branch'] == "main" for doc in second_five_results)
+        assert all(doc["metadata"]["branch"] == "main" for doc in second_five_results)
         assert first_five_results != second_five_results
 
         # multiple filters
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=0, metadata_filter={"filename": "file_5", "branch": "main"})
+        result = await vector_store_manager.list_documents_in_index(
+            index_name,
+            limit=5,
+            offset=0,
+            metadata_filter={"filename": "file_5", "branch": "main"},
+        )
         assert len(result) == 1
-        assert result[0]['metadata']['filename'] == "file_5"
-        assert result[0]['metadata']['branch'] == "main"
+        assert result[0]["metadata"]["filename"] == "file_5"
+        assert result[0]["metadata"]["branch"] == "main"
 
         # no results
-        result = await vector_store_manager.list_documents_in_index(index_name, limit=5, offset=0, metadata_filter={"filename": "file_15", "branch": "main"})
+        result = await vector_store_manager.list_documents_in_index(
+            index_name,
+            limit=5,
+            offset=0,
+            metadata_filter={"filename": "file_15", "branch": "main"},
+        )
         assert len(result) == 0
 
     @pytest.mark.asyncio
@@ -410,31 +540,58 @@ func main() {}"""
         index_name, second_index_name = "test_index", "second_test_index"
         # Create multiple documents
         documents = [
-            Document(text=f"Document {i}", metadata={"type": "text", "filename": f"file_{i}", "branch": "main"})
+            Document(
+                text=f"Document {i}",
+                metadata={"type": "text", "filename": f"file_{i}", "branch": "main"},
+            )
             for i in range(10)
         ]
-        
+
         await vector_store_manager.index_documents(index_name, documents)
 
         await vector_store_manager.persist(index_name, DEFAULT_VECTOR_DB_PERSIST_DIR)
-        await vector_store_manager.load(second_index_name, DEFAULT_VECTOR_DB_PERSIST_DIR, overwrite=True)
+        await vector_store_manager.load(
+            second_index_name, DEFAULT_VECTOR_DB_PERSIST_DIR, overwrite=True
+        )
 
-        result = await vector_store_manager.list_documents_in_index(second_index_name, limit=5, offset=0)
+        result = await vector_store_manager.list_documents_in_index(
+            second_index_name, limit=5, offset=0
+        )
         assert len(result) == 5
 
         # validate loaded index doesnt change the original index
-        await vector_store_manager.delete_documents(second_index_name, [result[0]['doc_id']])
+        await vector_store_manager.delete_documents(
+            second_index_name, [result[0]["doc_id"]]
+        )
 
-        first_index_result = await vector_store_manager.list_documents_in_index(index_name, limit=10, offset=0)
-        second_index_result = await vector_store_manager.list_documents_in_index(second_index_name, limit=10, offset=0)
+        first_index_result = await vector_store_manager.list_documents_in_index(
+            index_name, limit=10, offset=0
+        )
+        second_index_result = await vector_store_manager.list_documents_in_index(
+            second_index_name, limit=10, offset=0
+        )
         assert len(first_index_result) == 10
         assert len(second_index_result) == 9
 
-        second_index_result[0]['text'] = "Modified text"
-        second_update_result = await vector_store_manager.update_documents(second_index_name, [Document(doc_id=second_index_result[0]['doc_id'], text="Modified text", metadata=second_index_result[0]['metadata'])])
+        second_index_result[0]["text"] = "Modified text"
+        second_update_result = await vector_store_manager.update_documents(
+            second_index_name,
+            [
+                Document(
+                    doc_id=second_index_result[0]["doc_id"],
+                    text="Modified text",
+                    metadata=second_index_result[0]["metadata"],
+                )
+            ],
+        )
         assert len(second_update_result["updated_documents"]) == 1
         assert second_update_result["updated_documents"][0].text == "Modified text"
 
-        second_delete_result = await vector_store_manager.delete_documents(second_index_name, [second_index_result[0]['doc_id']])
+        second_delete_result = await vector_store_manager.delete_documents(
+            second_index_name, [second_index_result[0]["doc_id"]]
+        )
         assert len(second_delete_result["deleted_doc_ids"]) == 1
-        assert second_delete_result["deleted_doc_ids"][0] == second_index_result[0]['doc_id']
\ No newline at end of file
+        assert (
+            second_delete_result["deleted_doc_ids"][0]
+            == second_index_result[0]["doc_id"]
+        )
diff --git a/presets/ragengine/tests/vector_store/test_faiss_store.py b/presets/ragengine/tests/vector_store/test_faiss_store.py
index a68087f..b276d1b 100644
--- a/presets/ragengine/tests/vector_store/test_faiss_store.py
+++ b/presets/ragengine/tests/vector_store/test_faiss_store.py
@@ -12,53 +12,54 @@
 # limitations under the License.
 
 
-import pytest
 import os
-
 from tempfile import TemporaryDirectory
+
+import pytest
+
 from ragengine.tests.vector_store.test_base_store import BaseVectorStoreTest
 from ragengine.vector_store.faiss_store import FaissVectorStoreHandler
 
+
 class TestFaissVectorStore(BaseVectorStoreTest):
     """Test implementation for FAISS vector store."""
-    
+
     @pytest.fixture
     def vector_store_manager(self, init_embed_manager):
         with TemporaryDirectory() as temp_dir:
             print(f"Saving temporary test storage at: {temp_dir}")
-            os.environ['PERSIST_DIR'] = temp_dir
+            os.environ["PERSIST_DIR"] = temp_dir
             yield FaissVectorStoreHandler(init_embed_manager)
 
     @pytest.mark.asyncio
     async def check_indexed_documents(self, vector_store_manager):
         expected_output_1 = [
             {
-                'hash_value': '1e64a170be48c45efeaa8667ab35919106da0489ec99a11d0029f2842db133aa',
-                'text': 'First document in index1',
-                'is_truncated': False,
-                'metadata': {
-                    'type': 'text',
+                "hash_value": "1e64a170be48c45efeaa8667ab35919106da0489ec99a11d0029f2842db133aa",
+                "text": "First document in index1",
+                "is_truncated": False,
+                "metadata": {
+                    "type": "text",
                 },
             }
         ]
 
         expected_output_2 = [
             {
-                'hash_value': 'a222f875b83ce8b6eb72b3cae278b620de9bcc7c6b73222424d3ce979d1a463b',
-                'text': 'First document in index2',
-                'is_truncated': False,
-                'metadata': {
-                    'type': 'text',
+                "hash_value": "a222f875b83ce8b6eb72b3cae278b620de9bcc7c6b73222424d3ce979d1a463b",
+                "text": "First document in index2",
+                "is_truncated": False,
+                "metadata": {
+                    "type": "text",
                 },
             }
         ]
 
-        for index, expected_output in zip(["index1", "index2"], [expected_output_1, expected_output_2]):
+        for index, expected_output in zip(
+            ["index1", "index2"], [expected_output_1, expected_output_2], strict=False
+        ):
             response = await vector_store_manager.list_documents_in_index(
-                index,
-                limit=10,
-                offset=0,
-                max_text_length=1000
+                index, limit=10, offset=0, max_text_length=1000
             )
 
             # Remove "doc_id" from each document in the specified index
@@ -70,4 +71,4 @@ class TestFaissVectorStore(BaseVectorStoreTest):
     @property
     def expected_query_score(self):
         """Override this in implementation-specific test classes."""
-        return 0.5795239210128784
\ No newline at end of file
+        return 0.5795239210128784
diff --git a/presets/workspace/inference/text-generation/tests/test_inference_api.py b/presets/workspace/inference/text-generation/tests/test_inference_api.py
index 2a24528..b4df8a0 100644
--- a/presets/workspace/inference/text-generation/tests/test_inference_api.py
+++ b/presets/workspace/inference/text-generation/tests/test_inference_api.py
@@ -25,29 +25,48 @@ parent_dir = str(Path(__file__).resolve().parent.parent)
 # Add the parent directory to sys.path
 sys.path.append(parent_dir)
 
-CHAT_TEMPLATE = ("{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') %}"
+CHAT_TEMPLATE = (
+    "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') %}"
     "{{'
' + '\n' + message['content'] + '
' + '\n' + '
' + '\n'}}"
-    "{% elif (message['role'] == 'assistant') %}{{message['content'] + '
' + '\n'}}{% endif %}{% endfor %}")
-
-@pytest.fixture(params=[
-    {"pipeline": "text-generation", "model_path": "stanford-crfm/alias-gpt2-small-x21", "device": "cpu"},
-    {"pipeline": "conversational", "model_path": "stanford-crfm/alias-gpt2-small-x21", "device": "cpu"},
-])
+    "{% elif (message['role'] == 'assistant') %}{{message['content'] + '
' + '\n'}}{% endif %}{% endfor %}"
+)
+
+
+@pytest.fixture(
+    params=[
+        {
+            "pipeline": "text-generation",
+            "model_path": "stanford-crfm/alias-gpt2-small-x21",
+            "device": "cpu",
+        },
+        {
+            "pipeline": "conversational",
+            "model_path": "stanford-crfm/alias-gpt2-small-x21",
+            "device": "cpu",
+        },
+    ]
+)
 def configured_app(request):
     original_argv = sys.argv.copy()
     # Use request.param to set correct test arguments for each configuration
     test_args = [
-        'program_name',
-        '--pipeline', request.param['pipeline'],
-        '--pretrained_model_name_or_path', request.param['model_path'],
-        '--device_map', request.param['device'],
-        '--allow_remote_files', 'True',
-        '--chat_template', CHAT_TEMPLATE
+        "program_name",
+        "--pipeline",
+        request.param["pipeline"],
+        "--pretrained_model_name_or_path",
+        request.param["model_path"],
+        "--device_map",
+        request.param["device"],
+        "--allow_remote_files",
+        "True",
+        "--chat_template",
+        CHAT_TEMPLATE,
     ]
     sys.argv = test_args
 
     import inference_api
-    importlib.reload(inference_api) # Reload to prevent module caching
+
+    importlib.reload(inference_api)  # Reload to prevent module caching
     from inference_api import app
 
     # Attach the request params to the app instance for access in tests
@@ -56,18 +75,22 @@ def configured_app(request):
 
     sys.argv = original_argv
 
+
 def test_conversational(configured_app):
-    if configured_app.test_config['pipeline'] != 'conversational':
+    if configured_app.test_config["pipeline"] != "conversational":
         pytest.skip("Skipping non-conversational tests")
     client = TestClient(configured_app)
     messages = [
         {"role": "user", "content": "What is your favourite condiment?"},
-        {"role": "assistant", "content": "Well, im quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever im cooking up in the kitchen!"},
-        {"role": "user", "content": "Do you have mayonnaise recipes?"}
+        {
+            "role": "assistant",
+            "content": "Well, im quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever im cooking up in the kitchen!",
+        },
+        {"role": "user", "content": "Do you have mayonnaise recipes?"},
     ]
     request_data = {
         "messages": messages,
-        "generate_kwargs": {"max_new_tokens": 20, "do_sample": True}
+        "generate_kwargs": {"max_new_tokens": 20, "do_sample": True},
     }
     response = client.post("/chat", json=request_data)
 
@@ -76,26 +99,35 @@ def test_conversational(configured_app):
     assert "Result" in data
     assert len(data["Result"]) > 0  # Check if the conversation result is not empty
 
+
 def test_missing_messages_for_conversation(configured_app):
-    if configured_app.test_config['pipeline'] != 'conversational':
+    if configured_app.test_config["pipeline"] != "conversational":
         pytest.skip("Skipping non-conversational tests")
     client = TestClient(configured_app)
     request_data = {
         # "messages" is missing for conversational pipeline
     }
     response = client.post("/chat", json=request_data)
-    assert response.status_code == 400  # Expecting a Bad Request response due to missing messages
-    assert "Conversational parameter messages required" in response.json().get("detail", "")
+    assert (
+        response.status_code == 400
+    )  # Expecting a Bad Request response due to missing messages
+    assert "Conversational parameter messages required" in response.json().get(
+        "detail", ""
+    )
+
 
 def test_text_generation(configured_app):
-    if configured_app.test_config['pipeline'] != 'text-generation':
+    if configured_app.test_config["pipeline"] != "text-generation":
         pytest.skip("Skipping non-text-generation tests")
     client = TestClient(configured_app)
     request_data = {
         "prompt": "Hello, world!",
         "return_full_text": True,
         "clean_up_tokenization_spaces": False,
-        "generate_kwargs": {"max_length": 50, "min_length": 10}  # Example generate_kwargs
+        "generate_kwargs": {
+            "max_length": 50,
+            "min_length": 10,
+        },  # Example generate_kwargs
     }
     response = client.post("/chat", json=request_data)
     assert response.status_code == 200
@@ -103,19 +135,25 @@ def test_text_generation(configured_app):
     assert "Result" in data
     assert len(data["Result"]) > 0  # Check if the result text is not empty
 
+
 def test_missing_prompt(configured_app):
-    if configured_app.test_config['pipeline'] != 'text-generation':
+    if configured_app.test_config["pipeline"] != "text-generation":
         pytest.skip("Skipping non-text-generation tests")
     client = TestClient(configured_app)
     request_data = {
         # "prompt" is missing
         "return_full_text": True,
         "clean_up_tokenization_spaces": False,
-        "generate_kwargs": {"max_length": 50}
+        "generate_kwargs": {"max_length": 50},
     }
     response = client.post("/chat", json=request_data)
-    assert response.status_code == 400  # Expecting a Bad Request response due to missing prompt
-    assert "Text generation parameter prompt required" in response.json().get("detail", "")
+    assert (
+        response.status_code == 400
+    )  # Expecting a Bad Request response due to missing prompt
+    assert "Text generation parameter prompt required" in response.json().get(
+        "detail", ""
+    )
+
 
 def test_read_main(configured_app):
     client = TestClient(configured_app)
@@ -123,20 +161,24 @@ def test_read_main(configured_app):
     assert response.status_code == 200
     assert response.json() == {"message": "Server is running"}
 
+
 def test_health_check(configured_app):
     client = TestClient(configured_app)
     response = client.get("/health")
     assert response.status_code == 200
     assert response.json() == {"status": "Healthy"}
 
+
 def test_get_metrics(configured_app):
     client = TestClient(configured_app)
     response = client.get("/metrics")
     assert response.status_code == 200
     assert "gpu_info" in response.json()
 
+
 def test_get_metrics_with_gpus(configured_app):
     client = TestClient(configured_app)
+
     # Define a simple mock GPU object with the necessary attributes
     class MockGPU:
         def __init__(self, id, name, load, temperature, memoryUsed, memoryTotal):
@@ -153,14 +195,16 @@ def test_get_metrics_with_gpus(configured_app):
         name="GeForce GTX 950",
         load=0.25,  # 25%
         temperature=55,  # 55 C
-        memoryUsed=1 * (1024 ** 3),  # 1 GB
-        memoryTotal=2 * (1024 ** 3)  # 2 GB
+        memoryUsed=1 * (1024**3),  # 1 GB
+        memoryTotal=2 * (1024**3),  # 2 GB
     )
 
     # Mock torch.cuda.is_available to simulate an environment with GPUs
     # Mock GPUtil.getGPUs to return a list containing the mock GPU object
-    with patch('torch.cuda.is_available', return_value=True), \
-            patch('GPUtil.getGPUs', return_value=[mock_gpu]):
+    with (
+        patch("torch.cuda.is_available", return_value=True),
+        patch("GPUtil.getGPUs", return_value=[mock_gpu]),
+    ):
         response = client.get("/metrics")
         assert response.status_code == 200
         data = response.json()
@@ -176,17 +220,22 @@ def test_get_metrics_with_gpus(configured_app):
         assert gpu_data["temperature"] == "55 C"
         assert gpu_data["memory"]["used"] == "1.00 GB"
         assert gpu_data["memory"]["total"] == "2.00 GB"
-        assert data["cpu_info"] is None  # Assuming CPU info is not present when GPUs are available
+        assert (
+            data["cpu_info"] is None
+        )  # Assuming CPU info is not present when GPUs are available
+
 
 def test_get_metrics_no_gpus(configured_app):
     client = TestClient(configured_app)
     # Mock GPUtil.getGPUs to simulate an environment without GPUs
-    with patch('torch.cuda.is_available', return_value=False), \
-            patch('psutil.cpu_percent', return_value=20.0), \
-            patch('psutil.cpu_count', side_effect=[4, 8]), \
-            patch('psutil.virtual_memory') as mock_virtual_memory:
-        mock_virtual_memory.return_value.used = 4 * (1024 ** 3)  # 4 GB
-        mock_virtual_memory.return_value.total = 16 * (1024 ** 3)  # 16 GB
+    with (
+        patch("torch.cuda.is_available", return_value=False),
+        patch("psutil.cpu_percent", return_value=20.0),
+        patch("psutil.cpu_count", side_effect=[4, 8]),
+        patch("psutil.virtual_memory") as mock_virtual_memory,
+    ):
+        mock_virtual_memory.return_value.used = 4 * (1024**3)  # 4 GB
+        mock_virtual_memory.return_value.total = 16 * (1024**3)  # 16 GB
         response = client.get("/metrics")
         assert response.status_code == 200
         data = response.json()
@@ -198,8 +247,9 @@ def test_get_metrics_no_gpus(configured_app):
         assert data["cpu_info"]["memory"]["used"] == "4.00 GB"
         assert data["cpu_info"]["memory"]["total"] == "16.00 GB"
 
+
 def test_default_generation_params(configured_app):
-    if configured_app.test_config['pipeline'] != 'text-generation':
+    if configured_app.test_config["pipeline"] != "text-generation":
         pytest.skip("Skipping non-text-generation tests")
 
     client = TestClient(configured_app)
@@ -207,35 +257,40 @@ def test_default_generation_params(configured_app):
     request_data = {
         "prompt": "Test default params",
         "return_full_text": True,
-        "clean_up_tokenization_spaces": False
+        "clean_up_tokenization_spaces": False,
         # Note: generate_kwargs is not provided, so defaults should be used
     }
 
-    with patch('inference_api.pipeline') as mock_pipeline:
-        mock_pipeline.return_value = [{"generated_text": "Mocked response"}]  # Mock the response of the pipeline function
+    with patch("inference_api.pipeline") as mock_pipeline:
+        mock_pipeline.return_value = [
+            {"generated_text": "Mocked response"}
+        ]  # Mock the response of the pipeline function
 
         response = client.post("/chat", json=request_data)
 
         assert response.status_code == 200
         data = response.json()
         assert "Result" in data
-        assert data["Result"] == "Mocked response", "The response content doesn't match the expected mock response"
+        assert data["Result"] == "Mocked response", (
+            "The response content doesn't match the expected mock response"
+        )
 
         # Check the default args
         _, kwargs = mock_pipeline.call_args
-        assert kwargs['max_length'] == 200
-        assert kwargs['min_length'] == 0
-        assert kwargs['do_sample'] is True
-        assert kwargs['temperature'] == 1.0
-        assert kwargs['top_k'] == 10
-        assert kwargs['top_p'] == 1
-        assert kwargs['typical_p'] == 1
-        assert kwargs['repetition_penalty'] == 1
-        assert kwargs['num_beams'] == 1
-        assert kwargs['early_stopping'] is False
+        assert kwargs["max_length"] == 200
+        assert kwargs["min_length"] == 0
+        assert kwargs["do_sample"] is True
+        assert kwargs["temperature"] == 1.0
+        assert kwargs["top_k"] == 10
+        assert kwargs["top_p"] == 1
+        assert kwargs["typical_p"] == 1
+        assert kwargs["repetition_penalty"] == 1
+        assert kwargs["num_beams"] == 1
+        assert kwargs["early_stopping"] is False
+
 
 def test_generation_with_max_length(configured_app):
-    if configured_app.test_config['pipeline'] != 'text-generation':
+    if configured_app.test_config["pipeline"] != "text-generation":
         pytest.skip("Skipping non-text-generation tests")
 
     client = TestClient(configured_app)
@@ -247,7 +302,7 @@ def test_generation_with_max_length(configured_app):
         "prompt": prompt,
         "return_full_text": True,
         "clean_up_tokenization_spaces": False,
-        "generate_kwargs": {"max_length": max_length}
+        "generate_kwargs": {"max_length": max_length},
     }
 
     response = client.post("/chat", json=request_data)
@@ -257,20 +312,29 @@ def test_generation_with_max_length(configured_app):
     print("Response: ", data["Result"])
     assert "Result" in data, "The response should contain a 'Result' key"
 
-    tokenizer = AutoTokenizer.from_pretrained(configured_app.test_config['model_path'])
+    tokenizer = AutoTokenizer.from_pretrained(configured_app.test_config["model_path"])
     prompt_tokens = tokenizer.tokenize(prompt)
-    total_tokens = tokenizer.tokenize(data["Result"])  # data["Result"] includes the input prompt
+    total_tokens = tokenizer.tokenize(
+        data["Result"]
+    )  # data["Result"] includes the input prompt
 
     prompt_tokens_len = len(prompt_tokens)
     max_new_tokens = max_length - prompt_tokens_len
     new_tokens = len(total_tokens) - prompt_tokens_len
 
-    assert avg_res_len <= new_tokens, f"Ideally response should generate at least 15 tokens"
-    assert new_tokens <= max_new_tokens, "Response must not generate more than max new tokens"
-    assert len(total_tokens) <= max_length, "Total # of tokens has to be less than or equal to max_length"
+    assert avg_res_len <= new_tokens, (
+        "Ideally response should generate at least 15 tokens"
+    )
+    assert new_tokens <= max_new_tokens, (
+        "Response must not generate more than max new tokens"
+    )
+    assert len(total_tokens) <= max_length, (
+        "Total # of tokens has to be less than or equal to max_length"
+    )
+
 
 def test_generation_with_min_length(configured_app):
-    if configured_app.test_config['pipeline'] != 'text-generation':
+    if configured_app.test_config["pipeline"] != "text-generation":
         pytest.skip("Skipping non-text-generation tests")
 
     client = TestClient(configured_app)
@@ -282,7 +346,7 @@ def test_generation_with_min_length(configured_app):
         "prompt": prompt,
         "return_full_text": True,
         "clean_up_tokenization_spaces": False,
-        "generate_kwargs": {"min_length": min_length, "max_length": max_length}
+        "generate_kwargs": {"min_length": min_length, "max_length": max_length},
     }
 
     response = client.post("/chat", json=request_data)
@@ -291,9 +355,11 @@ def test_generation_with_min_length(configured_app):
     data = response.json()
     assert "Result" in data, "The response should contain a 'Result' key"
 
-    tokenizer = AutoTokenizer.from_pretrained(configured_app.test_config['model_path'])
+    tokenizer = AutoTokenizer.from_pretrained(configured_app.test_config["model_path"])
     prompt_tokens = tokenizer.tokenize(prompt)
-    total_tokens = tokenizer.tokenize(data["Result"])  # data["Result"] includes the input prompt
+    total_tokens = tokenizer.tokenize(
+        data["Result"]
+    )  # data["Result"] includes the input prompt
 
     prompt_tokens_len = len(prompt_tokens)
 
@@ -301,5 +367,9 @@ def test_generation_with_min_length(configured_app):
     max_new_tokens = max_length - prompt_tokens_len
     new_tokens = len(total_tokens) - prompt_tokens_len
 
-    assert min_new_tokens <= new_tokens <= max_new_tokens, "Response should generate at least min_new_tokens and at most max_new_tokens new tokens"
-    assert len(total_tokens) <= max_length, "Total # of tokens has to be less than or equal to max_length"
+    assert min_new_tokens <= new_tokens <= max_new_tokens, (
+        "Response should generate at least min_new_tokens and at most max_new_tokens new tokens"
+    )
+    assert len(total_tokens) <= max_length, (
+        "Total # of tokens has to be less than or equal to max_length"
+    )
diff --git a/presets/workspace/inference/text-generation/tests/test_model_config.py b/presets/workspace/inference/text-generation/tests/test_model_config.py
index b7a787e..2b276ec 100644
--- a/presets/workspace/inference/text-generation/tests/test_model_config.py
+++ b/presets/workspace/inference/text-generation/tests/test_model_config.py
@@ -22,28 +22,41 @@ parent_dir = str(Path(__file__).resolve().parent.parent)
 # Add the parent directory to sys.path
 sys.path.append(parent_dir)
 
-@pytest.fixture(params=[
-    {"pipeline": "text-generation", "model_path": "stanford-crfm/alias-gpt2-small-x21"},
-    {"pipeline": "conversational", "model_path": "stanford-crfm/alias-gpt2-small-x21"},
-])
+
+@pytest.fixture(
+    params=[
+        {
+            "pipeline": "text-generation",
+            "model_path": "stanford-crfm/alias-gpt2-small-x21",
+        },
+        {
+            "pipeline": "conversational",
+            "model_path": "stanford-crfm/alias-gpt2-small-x21",
+        },
+    ]
+)
 def configured_model_config(request):
     original_argv = sys.argv.copy()
 
     sys.argv = [
-        'program_name',
-        '--pipeline', request.param['pipeline'],
-        '--pretrained_model_name_or_path', request.param['model_path'],
-        '--allow_remote_files', 'True'
+        "program_name",
+        "--pipeline",
+        request.param["pipeline"],
+        "--pretrained_model_name_or_path",
+        request.param["model_path"],
+        "--allow_remote_files",
+        "True",
     ]
 
     import inference_api
+
     importlib.reload(inference_api)
     from inference_api import ModelConfig
 
     # Create and configure the ModelConfig instance
     model_config = ModelConfig(
-        pipeline=request.param['pipeline'],
-        pretrained_model_name_or_path=request.param['model_path'],
+        pipeline=request.param["pipeline"],
+        pretrained_model_name_or_path=request.param["model_path"],
     )
 
     yield model_config
@@ -51,15 +64,18 @@ def configured_model_config(request):
     # Restore the original sys.argv after the test is done
     sys.argv = original_argv
 
+
 def test_process_additional_args(configured_model_config):
     config = configured_model_config
 
     # Simulate additional command-line arguments
     additional_args = [
-        "--new_arg1", "value1",
+        "--new_arg1",
+        "value1",
         "--new_arg2",
-        "--new_arg3", "value3",
-        "--flag_arg"
+        "--new_arg3",
+        "value3",
+        "--flag_arg",
     ]
 
     # Process the additional arguments
@@ -71,13 +87,11 @@ def test_process_additional_args(configured_model_config):
     assert getattr(config, "new_arg3", None) == "value3"
     assert getattr(config, "flag_arg", None) is True
 
+
 # Test case for ignoring arguments prefixed with '--' when expecting a value
 def test_ignore_double_dash_arguments(configured_model_config):
     config = configured_model_config
-    additional_args = [
-        "--new_arg1", "--new_arg2",
-        "--new_arg3", "correct_value"
-    ]
+    additional_args = ["--new_arg1", "--new_arg2", "--new_arg3", "correct_value"]
 
     config.process_additional_args(additional_args)
 
@@ -86,16 +100,20 @@ def test_ignore_double_dash_arguments(configured_model_config):
     assert getattr(config, "new_arg2", None) is True
     assert getattr(config, "new_arg3", None) == "correct_value"
 
+
 # Test case to verify handling unsupported pipeline values
 def test_unsupported_pipeline_raises_value_error(configured_model_config):
     with pytest.raises(ValueError) as excinfo:
         from inference_api import ModelConfig
+
         ModelConfig(pipeline="unsupported_pipeline")
     assert "Unsupported pipeline" in str(excinfo.value)
 
+
 # Test case for validating torch_dtype
 def test_invalid_torch_dtype_raises_value_error(configured_model_config):
     with pytest.raises(ValueError) as excinfo:
         from inference_api import ModelConfig
+
         ModelConfig(pipeline="text-generation", torch_dtype="unsupported_dtype")
-    assert "Invalid torch dtype" in str(excinfo.value)
\ No newline at end of file
+    assert "Invalid torch dtype" in str(excinfo.value)
diff --git a/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py b/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py
index d9d1b5c..90bc774 100644
--- a/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py
+++ b/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py
@@ -11,33 +11,39 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
 import os
+import shutil
+import socket
 import subprocess
+import sys
 import time
-import socket
 from pathlib import Path
 
 import pytest
 import requests
+from huggingface_hub import snapshot_download
 
 # Get the parent directory of the current file
 parent_dir = str(Path(__file__).resolve().parent.parent)
 # Add the parent directory to sys.path
 sys.path.append(parent_dir)
 
-from inference_api import binary_search_with_limited_steps, KaitoConfig
-from huggingface_hub import snapshot_download
-import shutil
+try:
+    from inference_api import KaitoConfig, binary_search_with_limited_steps
+except ImportError as e:
+    raise ImportError(f"Error importing required module(s): {e}") from e
 
 TEST_MODEL = "facebook/opt-125m"
 TEST_ADAPTER_NAME1 = "mylora1"
 TEST_ADAPTER_NAME2 = "mylora2"
 TEST_MODEL_NAME = "mymodel"
 TEST_MODEL_LEN = 1024
-CHAT_TEMPLATE = ("{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') %}"
+CHAT_TEMPLATE = (
+    "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') %}"
     "{{'
' + '\n' + message['content'] + '
' + '\n' + '
' + '\n'}}"
-    "{% elif (message['role'] == 'assistant') %}{{message['content'] + '
' + '\n'}}{% endif %}{% endfor %}")
+    "{% elif (message['role'] == 'assistant') %}{{message['content'] + '
' + '\n'}}{% endif %}{% endfor %}"
+)
+
 
 @pytest.fixture(scope="session", autouse=True)
 def setup_server(request, tmp_path_factory, autouse=True):
@@ -51,17 +57,21 @@ def setup_server(request, tmp_path_factory, autouse=True):
     # prepare testing adapter images
     tmp_file_dir = tmp_path_factory.mktemp("adapter")
     print(f"Downloading adapter image to {tmp_file_dir}")
-    snapshot_download(repo_id="slall/facebook-opt-125M-imdb-lora", local_dir=str(tmp_file_dir / TEST_ADAPTER_NAME1))
-    snapshot_download(repo_id="slall/facebook-opt-125M-imdb-lora", local_dir=str(tmp_file_dir / TEST_ADAPTER_NAME2))
+    snapshot_download(
+        repo_id="slall/facebook-opt-125M-imdb-lora",
+        local_dir=str(tmp_file_dir / TEST_ADAPTER_NAME1),
+    )
+    snapshot_download(
+        repo_id="slall/facebook-opt-125M-imdb-lora",
+        local_dir=str(tmp_file_dir / TEST_ADAPTER_NAME2),
+    )
 
     # prepare testing config file
     config_file = tmp_file_dir / "config.yaml"
     kaito_config = KaitoConfig(
-        vllm={
-            "max-model-len": TEST_MODEL_LEN,
-            "served-model-name": TEST_MODEL_NAME
-        },
+        vllm={"max-model-len": TEST_MODEL_LEN, "served-model-name": TEST_MODEL_NAME},
         max_probe_steps=0,
+        kv_cache_cpu_memory_utilization=0.6,
     )
     with open(config_file, "w") as f:
         f.write(kaito_config.to_yaml())
@@ -69,17 +79,25 @@ def setup_server(request, tmp_path_factory, autouse=True):
     args = [
         "python3",
         os.path.join(parent_dir, "inference_api.py"),
-        "--model", TEST_MODEL,
-        "--chat-template", CHAT_TEMPLATE,
-        "--max-model-len", "2048", # expected to be overridden by config file
-        "--port", str(TEST_PORT),
-        "--kaito-adapters-dir", tmp_file_dir,
-        "--kaito-config-file", config_file,
+        "--model",
+        TEST_MODEL,
+        "--chat-template",
+        CHAT_TEMPLATE,
+        "--max-model-len",
+        "2048",  # expected to be overridden by config file
+        "--port",
+        str(TEST_PORT),
+        "--kaito-adapters-dir",
+        tmp_file_dir,
+        "--kaito-config-file",
+        config_file,
     ]
     print(f">>> Starting server on port {TEST_PORT}")
     env = os.environ.copy()
     env["MAX_PROBE_STEPS"] = "0"
-    process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
+    process = subprocess.Popen(
+        args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
+    )
 
     def fin():
         process.terminate()
@@ -88,7 +106,7 @@ def setup_server(request, tmp_path_factory, autouse=True):
         print(f">>> Server stderr: {stderr}")
         stdout = process.stdout.read().decode()
         print(f">>> Server stdout: {stdout}")
-        print ("\n>>> Doing teardown")
+        print("\n>>> Doing teardown")
         print(f"Removing adapter image from {tmp_file_dir}")
         shutil.rmtree(tmp_file_dir)
 
@@ -98,6 +116,7 @@ def setup_server(request, tmp_path_factory, autouse=True):
 
     request.addfinalizer(fin)
 
+
 def is_port_open(host, port, timeout=60):
     start_time = time.time()
     while time.time() - start_time < timeout:
@@ -111,12 +130,14 @@ def is_port_open(host, port, timeout=60):
             time.sleep(1)  # Wait a bit before retrying
     return False
 
+
 def find_available_port(start_port=5000, end_port=8000):
     for port in range(start_port, end_port + 1):
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            if s.connect_ex(('localhost', port)) != 0:
+            if s.connect_ex(("localhost", port)) != 0:
                 return port
-    raise RuntimeError('No available ports found')
+    raise RuntimeError("No available ports found")
+
 
 def test_completions_api(setup_server):
     request_data = {
@@ -124,10 +145,12 @@ def test_completions_api(setup_server):
         "prompt": "Say this is a test",
         "max_tokens": 7,
         "temperature": 0.5,
-        "n": 2
+        "n": 2,
     }
 
-    response = requests.post(f"http://127.0.0.1:{TEST_PORT}/v1/completions", json=request_data)
+    response = requests.post(
+        f"http://127.0.0.1:{TEST_PORT}/v1/completions", json=request_data
+    )
     data = response.json()
     assert "choices" in data, "The response should contain a 'choices' key"
     assert len(data["choices"]) == 2, "The response should contain two completion"
@@ -136,19 +159,22 @@ def test_completions_api(setup_server):
         assert "text" in choice, "Each choice should contain a 'text' key"
         assert len(choice["text"]) > 0, "The completion text should not be empty"
 
+
 def test_chat_completions_api(setup_server):
     request_data = {
         "model": TEST_MODEL_NAME,
         "messages": [
             {"role": "user", "content": "Hello!"},
-            {"role": "assistant", "content": "Hi there! How can I help you today?"}
+            {"role": "assistant", "content": "Hi there! How can I help you today?"},
         ],
         "max_tokens": 7,
         "temperature": 0.5,
-        "n": 2
+        "n": 2,
     }
 
-    response = requests.post(f"http://127.0.0.1:{TEST_PORT}/v1/chat/completions", json=request_data)
+    response = requests.post(
+        f"http://127.0.0.1:{TEST_PORT}/v1/chat/completions", json=request_data
+    )
     data = response.json()
 
     assert "choices" in data, "The response should contain a 'choices' key"
@@ -156,24 +182,43 @@ def test_chat_completions_api(setup_server):
 
     for choice in data["choices"]:
         assert "message" in choice, "Each choice should contain a 'message' key"
-        assert "content" in choice["message"], "Each message should contain a 'content' key"
-        assert len(choice["message"]["content"]) > 0, "The completion text should not be empty"
+        assert "content" in choice["message"], (
+            "Each message should contain a 'content' key"
+        )
+        assert len(choice["message"]["content"]) > 0, (
+            "The completion text should not be empty"
+        )
+
 
 def test_model_list(setup_server):
     response = requests.get(f"http://127.0.0.1:{TEST_PORT}/v1/models")
     data = response.json()
 
     assert "data" in data, f"The response should contain a 'data' key, but got {data}"
-    assert len(data["data"]) == 3, f"The response should contain three models, but got {data['data']}"
-    assert data["data"][0]["id"] == TEST_MODEL_NAME, f"The first model should be the test model, but got {data['data'][0]['id']}"
-    assert data["data"][0]["max_model_len"] == TEST_MODEL_LEN, f"The first model should have the test model length, but got {data['data'][0]['max_model_len']}"
-    assert data["data"][1]["id"] == TEST_ADAPTER_NAME2, f"The second model should be the test adapter, but got {data['data'][1]['id']}"
-    assert data["data"][1]["parent"] == TEST_MODEL_NAME, f"The second model should have the test model as parent, but got {data['data'][1]['parent']}"
-    assert data["data"][2]["id"] == TEST_ADAPTER_NAME1, f"The third model should be the test adapter, but got {data['data'][2]['id']}"
-    assert data["data"][2]["parent"] == TEST_MODEL_NAME, f"The third model should have the test model as parent, but got {data['data'][2]['parent']}"
+    assert len(data["data"]) == 3, (
+        f"The response should contain three models, but got {data['data']}"
+    )
+    assert data["data"][0]["id"] == TEST_MODEL_NAME, (
+        f"The first model should be the test model, but got {data['data'][0]['id']}"
+    )
+    assert data["data"][0]["max_model_len"] == TEST_MODEL_LEN, (
+        f"The first model should have the test model length, but got {data['data'][0]['max_model_len']}"
+    )
+    assert data["data"][1]["id"] == TEST_ADAPTER_NAME2, (
+        f"The second model should be the test adapter, but got {data['data'][1]['id']}"
+    )
+    assert data["data"][1]["parent"] == TEST_MODEL_NAME, (
+        f"The second model should have the test model as parent, but got {data['data'][1]['parent']}"
+    )
+    assert data["data"][2]["id"] == TEST_ADAPTER_NAME1, (
+        f"The third model should be the test adapter, but got {data['data'][2]['id']}"
+    )
+    assert data["data"][2]["parent"] == TEST_MODEL_NAME, (
+        f"The third model should have the test model as parent, but got {data['data'][2]['parent']}"
+    )
 
-def test_binary_search_with_limited_steps():
 
+def test_binary_search_with_limited_steps():
     def is_safe_fn(x):
         return x <= 10
 
diff --git a/presets/workspace/test/falcon-benchmark/inference-benchmark.py b/presets/workspace/test/falcon-benchmark/inference-benchmark.py
index 148c2ff..01e6b32 100644
--- a/presets/workspace/test/falcon-benchmark/inference-benchmark.py
+++ b/presets/workspace/test/falcon-benchmark/inference-benchmark.py
@@ -11,15 +11,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import csv
 import argparse
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import transformers
-import torch
-from datetime import datetime
+import csv
 import time
 import uuid
+from datetime import datetime
+
+import torch
+import transformers
 from accelerate import Accelerator
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
 
 def get_args():
     parser = argparse.ArgumentParser()
@@ -32,9 +34,14 @@ def get_args():
     parser.add_argument("--data_parallelism", type=str, required=True)
     parser.add_argument("--quantization", type=str, required=True)
     parser.add_argument("--machine", type=str, required=True)
-    parser.add_argument("--use_accelerator", action='store_true', help="Use the Accelerator for parallel processing.")
+    parser.add_argument(
+        "--use_accelerator",
+        action="store_true",
+        help="Use the Accelerator for parallel processing.",
+    )
     return parser.parse_args()
 
+
 def inference(requests):
     for request in requests:
         start_time = time.time()
@@ -50,8 +57,8 @@ def inference(requests):
 
         end_time = time.time()
         inference_time = end_time - start_time
-        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-    
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
         result = {
             "model": args.model,
             "num_nodes": args.num_nodes,
@@ -64,14 +71,15 @@ def inference(requests):
             "quantization": args.quantization,
             "machine": args.machine,
             "inference_time": inference_time,
-            "request_id": str(uuid.uuid4()), # Generate a unique UUID
-            "timestamp": timestamp
+            "request_id": str(uuid.uuid4()),  # Generate a unique UUID
+            "timestamp": timestamp,
         }
         writer.writerow(result)
 
         for seq in sequences:
             print(f"Result: {seq['generated_text']}")
 
+
 args = get_args()
 
 model_id = "tiiuae/falcon-7b-instruct"
@@ -93,13 +101,26 @@ pipeline = transformers.pipeline(
 )
 
 
-
-with open("../common-gpt-questions.csv", "r") as f:
+with open("../common-gpt-questions.csv") as f:
     requests = [line.strip() for line in f.readlines()]
 
-fieldnames = ["model", "num_nodes", "num_processes", "num_gpus", "num_prompts", "prompt_len", "model_parallelism", "data_parallelism", "quantization", "machine", "inference_time", "request_id", "timestamp"]
-
-with open("results.csv", "a", newline='') as f:
+fieldnames = [
+    "model",
+    "num_nodes",
+    "num_processes",
+    "num_gpus",
+    "num_prompts",
+    "prompt_len",
+    "model_parallelism",
+    "data_parallelism",
+    "quantization",
+    "machine",
+    "inference_time",
+    "request_id",
+    "timestamp",
+]
+
+with open("results.csv", "a", newline="") as f:
     writer = csv.DictWriter(f, fieldnames=fieldnames)
     writer.writeheader()
 
@@ -110,5 +131,3 @@ with open("results.csv", "a", newline='') as f:
             inference(split_requests)
     else:
         inference(requests)
-
-        
diff --git a/presets/workspace/test/falcon-benchmark/plot.ipynb b/presets/workspace/test/falcon-benchmark/plot.ipynb
index 2b49852..459913c 100644
--- a/presets/workspace/test/falcon-benchmark/plot.ipynb
+++ b/presets/workspace/test/falcon-benchmark/plot.ipynb
@@ -1,15 +1,23 @@
-# Copyright (c) KAITO authors.
-# Licensed under the MIT license.
 {
  "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Copyright (c) KAITO authors.\n",
+    "# Licensed under the MIT license."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pandas as pd\n",
     "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
     "import seaborn as sns"
    ]
   },
@@ -737,8 +745,8 @@
     }
    ],
    "source": [
-    "import pandas as pd\n",
     "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
     "\n",
     "# Assuming the CSV is already read into combined_df\n",
     "# Convert the timestamp column to datetime type\n",
@@ -757,7 +765,7 @@
     "# Plotting experiment duration\n",
     "plt.figure(figsize=(10, 6))\n",
     "\n",
-    "for parallelism, duration in zip(df_grouped_duration['model_parallelism'], df_grouped_duration['experiment_duration']):\n",
+    "for parallelism, duration in zip(df_grouped_duration['model_parallelism'], df_grouped_duration['experiment_duration'], strict=False):\n",
     "    plt.bar(parallelism, duration, label=f\"falcon-40b ({parallelism})\")\n",
     "    plt.annotate(f\"{duration:.2f}s\", \n",
     "                 (parallelism, duration),\n",
diff --git a/presets/workspace/test/llama-test/inference-benchmark.py b/presets/workspace/test/llama-test/inference-benchmark.py
index 07257f2..cab35f4 100644
--- a/presets/workspace/test/llama-test/inference-benchmark.py
+++ b/presets/workspace/test/llama-test/inference-benchmark.py
@@ -11,48 +11,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import requests
-import time
 import csv
 import random
+import time
 from datetime import datetime
 
+import requests
+
+
 def is_file_empty(filename):
     try:
         return not bool(len(open(filename).readline()))
-    except:
+    except FileNotFoundError:
         return True
 
+
 def load_search_terms_from_csv(filename):
-    with open(filename, 'r', newline='') as file:
+    with open(filename, newline="") as file:
         reader = csv.DictReader(file)
-        return [row['search_terms'] for row in reader]
+        return [row["search_terms"] for row in reader]
+
 
 # Load search terms and shuffle them for randomness
 search_terms = load_search_terms_from_csv("../common-gpt-questions.csv")
 random.shuffle(search_terms)
 
 # Constants
-URL = "http://20.241.194.198/chat" # Replace with service URL
+URL = "http://20.241.194.198/chat"  # Replace with service URL
 input_payload = {
     "input_data": {
-        "input_string": [
-            [
-                {
-                    "role": "user",
-                    "content": ""
-                }
-            ]
-        ], 
+        "input_string": [[{"role": "user", "content": ""}]],
     },
-    "parameters": {
-        "temperature": 0.6,
-        "top_p": 0.9,
-        "max_gen_len": 64
-    }
+    "parameters": {"temperature": 0.6, "top_p": 0.9, "max_gen_len": 64},
 }
 
-NUM_REQUESTS = min(1000, len(search_terms))  # Cannot exceed the number of questions available
+NUM_REQUESTS = min(
+    1000, len(search_terms)
+)  # Cannot exceed the number of questions available
 
 # Generate a unique run_id based on the current timestamp
 run_id = int(time.time())
@@ -61,34 +56,36 @@ run_id = int(time.time())
 latencies = []
 
 # Open the CSV for writing
-with open('gpt-requests.csv', 'a', newline='') as file:
+with open("gpt-requests.csv", "a", newline="") as file:
     writer = csv.writer(file)
-    
+
     # If the file is empty, write the header
-    if is_file_empty('gpt-requests.csv'):
+    if is_file_empty("gpt-requests.csv"):
         writer.writerow(["run_id", "request_id", "request_num", "latency", "timestamp"])
 
     for i in range(NUM_REQUESTS):
         question = search_terms.pop()  # Get a random question without replacement
         print("Question Asked: ", question)
-        input_payload['input_data']['input_string'][0][0]['content'] = question
+        input_payload["input_data"]["input_string"][0][0]["content"] = question
 
         start_time = time.time()
-        
+
         response = requests.post(URL, json=input_payload)
         print("Response Status:", response.status_code)
         print("Response Body:", response.text)
         # Get the date/time for this request
-        request_date = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+        request_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
         if response.status_code == 200:
             elapsed_time = (time.time() - start_time) * 1000
-            print(f"Request #{i+1} elapsed time", elapsed_time)
+            print(f"Request #{i + 1} elapsed time", elapsed_time)
             latencies.append(elapsed_time)
-            request_id = f"{run_id}-{i+1}"
-            writer.writerow([run_id, request_id, i+1, elapsed_time, request_date])
+            request_id = f"{run_id}-{i + 1}"
+            writer.writerow([run_id, request_id, i + 1, elapsed_time, request_date])
         else:
-            print(f"Request #{i+1} failed with status code {response.status_code}. Error: {response.text}")
+            print(
+                f"Request #{i + 1} failed with status code {response.status_code}. Error: {response.text}"
+            )
 
 # Calculate statistics
 average_latency = sum(latencies) / len(latencies)
diff --git a/presets/workspace/test/manifests/inference-tmpl/manifest.yaml b/presets/workspace/test/manifests/inference-tmpl/manifest.yaml
index 1181dc2..ee94393 100644
--- a/presets/workspace/test/manifests/inference-tmpl/manifest.yaml
+++ b/presets/workspace/test/manifests/inference-tmpl/manifest.yaml
@@ -28,7 +28,7 @@ deployment:
             mountPath: /workspace/weights
         containers:
         - name: WORKLOAD_NAME-container
-          image: REPO/base:0.0.4
+          image: REPO/base:0.0.5
           command:
             - /bin/sh
             - -c
@@ -94,6 +94,7 @@ config:
         swap-space: 4
         served-model-name: test
         dtype: float16
+        max-model-len: 2048
 
         # max-seq-len-to-capture: 8192
         # num-scheduler-steps: 1
diff --git a/presets/workspace/test/scripts/generate_manifests.py b/presets/workspace/test/scripts/generate_manifests.py
index b676592..d3bd808 100644
--- a/presets/workspace/test/scripts/generate_manifests.py
+++ b/presets/workspace/test/scripts/generate_manifests.py
@@ -12,114 +12,147 @@
 # limitations under the License.
 
 #!/usr/bin/env python3
+import argparse
 import json
-import yaml
-import sys
 import os
-import argparse
+import sys
 from pathlib import Path
 
+import yaml
+
+
 def get_project_root():
     # Get the directory containing the current script
     current_dir = Path(__file__).resolve().parent
     # Go up 4 levels to reach project root (from presets/workspace/test/scripts)
     return current_dir.parents[3]
 
+
 def load_json_config():
     project_root = get_project_root()
-    config_path = project_root / '.github' / 'e2e-preset-configs.json'
-    with open(config_path, 'r') as f:
+    config_path = project_root / ".github" / "e2e-preset-configs.json"
+    with open(config_path) as f:
         data = json.load(f)
-        return data['matrix']['image']  # Return the array of model configs
+        return data["matrix"]["image"]  # Return the array of model configs
+
 
 def load_template():
     project_root = get_project_root()
-    template_path = project_root / 'presets' / 'workspace' / 'test' / 'manifests' / 'inference-tmpl' / 'manifest.yaml'
-    with open(template_path, 'r') as f:
+    template_path = os.path.join(
+        project_root,
+        "presets",
+        "workspace",
+        "test",
+        "manifests",
+        "inference-tmpl",
+        "manifest.yaml",
+    )
+    with open(template_path) as f:
         return yaml.safe_load(f)
 
+
 def check_predefined_manifest(model_name):
     project_root = get_project_root()
-    manifest_path = project_root / 'presets' / 'workspace' / 'test' / 'manifests' / f'{model_name}' / f'{model_name}.yaml'
+    manifest_path = os.path.join(
+        str(project_root),
+        "presets",
+        "workspace",
+        "test",
+        "manifests",
+        f"{model_name}",
+        f"{model_name}.yaml",
+    )
     if not os.path.exists(manifest_path):
         return (None, False)
-    with open(manifest_path, 'r') as f:
+    with open(manifest_path) as f:
         return (f.read(), True)
 
+
 def process_model(model_name, runtime, repo=None, tag=None):
     configs = load_json_config()
-    model_config = next((m for m in configs if m['name'] == model_name), None)
+    model_config = next((m for m in configs if m["name"] == model_name), None)
 
     if not model_config:
         print(f"Model {model_name} not found in configs", file=sys.stderr)
         sys.exit(1)
 
-    predefined_manifest, predefined_manifest_exists = check_predefined_manifest(model_name)
+    predefined_manifest, predefined_manifest_exists = check_predefined_manifest(
+        model_name
+    )
     if predefined_manifest_exists:
-        return process_predefined_manifest(model_name, runtime, predefined_manifest, repo, tag)
-
-    if runtime not in model_config.get('runtimes', {}):
-        print(f"Runtime {runtime} not configured for model {model_name}", file=sys.stderr)
+        return process_predefined_manifest(
+            model_name, runtime, predefined_manifest, repo, tag
+        )
+
+    if runtime not in model_config.get("runtimes", {}):
+        print(
+            f"Runtime {runtime} not configured for model {model_name}", file=sys.stderr
+        )
         sys.exit(1)
 
-    runtime_config = model_config['runtimes'][runtime]
-    workload_name = model_config.get('workload', model_name)
+    runtime_config = model_config["runtimes"][runtime]
+    workload_name = model_config.get("workload", model_name)
 
     templates = load_template()
-    manifest_str = yaml.dump(templates['deployment'])
+    manifest_str = yaml.dump(templates["deployment"])
 
     # Replace placeholders in template
     manifest_str = (
-        manifest_str
-        .replace('WORKLOAD_NAME', workload_name)
-        .replace('MODEL_NAME', model_name)
-        .replace('RUNTIME_COMMAND', runtime_config['command'])
-        .replace('GPU_COUNT', str(runtime_config['gpu_count']))
-        .replace('NODE_POOL', model_config['node_pool'])
-        .replace('NODE_COUNT', str(model_config['node-count']))
+        manifest_str.replace("WORKLOAD_NAME", workload_name)
+        .replace("MODEL_NAME", model_name)
+        .replace("RUNTIME_COMMAND", runtime_config["command"])
+        .replace("GPU_COUNT", str(runtime_config["gpu_count"]))
+        .replace("NODE_POOL", model_config["node_pool"])
+        .replace("NODE_COUNT", str(model_config["node-count"]))
     )
 
     # Replace repo and tag if provided
     if repo:
-        manifest_str = manifest_str.replace('REPO', repo)
+        manifest_str = manifest_str.replace("REPO", repo)
     if tag:
-        manifest_str = manifest_str.replace('TAG', tag)
+        manifest_str = manifest_str.replace("TAG", tag)
 
     # Parse the template string back into YAML
     manifest = yaml.safe_load(manifest_str)
 
     # Generate service manifest
-    service_template = templates['service']
+    service_template = templates["service"]
     service_str = yaml.dump(service_template)
-    service_str = service_str.replace('WORKLOAD_NAME', workload_name)
+    service_str = service_str.replace("WORKLOAD_NAME", workload_name)
     service_manifest = yaml.safe_load(service_str)
 
     # Generate config manifest
-    config_template = templates['config']
+    config_template = templates["config"]
     config_str = yaml.dump(config_template)
     config_manifest = yaml.safe_load(config_str)
 
     # Print manifests to stdout
     yaml.dump(manifest, sys.stdout, default_flow_style=False)
-    print('---')  # Document separator
+    print("---")  # Document separator
     yaml.dump(service_manifest, sys.stdout, default_flow_style=False)
-    print('---')  # Document separator
+    print("---")  # Document separator
     yaml.dump(config_manifest, sys.stdout, default_flow_style=False)
 
-def process_predefined_manifest(model_name, runtime, predefined_manifest, repo=None, tag=None):
+
+def process_predefined_manifest(
+    model_name, runtime, predefined_manifest, repo=None, tag=None
+):
     if repo:
-        predefined_manifest = predefined_manifest.replace('REPO', repo)
+        predefined_manifest = predefined_manifest.replace("REPO", repo)
     if tag:
-        predefined_manifest = predefined_manifest.replace('TAG', tag)
+        predefined_manifest = predefined_manifest.replace("TAG", tag)
 
     print(predefined_manifest)
 
+
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description='Process model template with optional repo and tag.')
-    parser.add_argument('model_name', help='Name of the model')
-    parser.add_argument('runtime', help='Runtime to use')
-    parser.add_argument('--repo', help='Repository name to use instead of REPO')
-    parser.add_argument('--tag', help='Tag to use instead of TAG')
+    parser = argparse.ArgumentParser(
+        description="Process model template with optional repo and tag."
+    )
+    parser.add_argument("model_name", help="Name of the model")
+    parser.add_argument("runtime", help="Runtime to use")
+    parser.add_argument("--repo", help="Repository name to use instead of REPO")
+    parser.add_argument("--tag", help="Tag to use instead of TAG")
 
     args = parser.parse_args()
     process_model(args.model_name, args.runtime, args.repo, args.tag)