Repository: kaito-project/kaito

Issue:
add linter/formatter for python code
**Is your feature request related to a problem? Please describe.**

**Describe the solution you'd like**

- [ ] add python linter make command
- [ ] add python linter pipeline
- [ ] add pre-commit config for linter https://pre-commit.com/

**Describe alternatives you've considered**

**Additional context**

Discussion:
@farooqameen: Are you open to using [ruff](https://docs.astral.sh/ruff/) as a linter and formatter? If yes, I'd like to work on this.@zhuangqh: > Are you open to using [ruff](https://docs.astral.sh/ruff/) as a linter and formatter? If yes, I'd like to work on this.

ruff is great. let's do that.

Solution patch:
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/workflows/lint-go.yaml b/.github/workflows/lint-go.yaml
index 0705f2d..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:
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/unit-tests.yaml b/.github/workflows/unit-tests.yaml
index 6ab1541..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]
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/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/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/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..074e7d2 100644
--- a/presets/workspace/inference/vllm/inference_api.py
+++ b/presets/workspace/inference/vllm/inference_api.py
@@ -11,30 +11,33 @@
 # 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 torch
-from vllm.utils import FlexibleArgumentParser
+import uvloop
 import vllm.entrypoints.openai.api_server as api_server
+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 +51,26 @@ 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.",
+        )
 
     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,14 +86,14 @@ 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:
@@ -98,6 +114,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.
@@ -108,19 +125,20 @@ class KaitoConfig:
     max_probe_steps: int
 
     @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),
         )
 
     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 +151,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 +162,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 +175,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 +208,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 +217,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 +260,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 +273,9 @@ 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
 
+
 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
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/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/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/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.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.6.x/example_rag_client.py b/website/versioned_docs/version-v0.6.x/example_rag_client.py
index 60d5e83..e2a5760 100644
--- a/website/versioned_docs/version-v0.6.x/example_rag_client.py
+++ b/website/versioned_docs/version-v0.6.x/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()



Test patch:
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..f0e00f3 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,16 +57,19 @@ 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,
     )
     with open(config_file, "w") as f:
@@ -69,17 +78,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 +105,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 +115,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 +129,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 +144,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 +158,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 +181,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/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)