diff --git a/llm-service/app/ai/vector_stores/qdrant.py b/llm-service/app/ai/vector_stores/qdrant.py
index 9d4a022..6afbdff 100644
--- a/llm-service/app/ai/vector_stores/qdrant.py
+++ b/llm-service/app/ai/vector_stores/qdrant.py
@@ -37,7 +37,7 @@
 #
 
 import os
-from typing import Optional, List, Tuple, Any
+from typing import Optional
 
 import qdrant_client
 from llama_index.core.indices import VectorStoreIndex
@@ -45,7 +45,7 @@
 from llama_index.vector_stores.qdrant import (
     QdrantVectorStore as LlamaIndexQdrantVectorStore,
 )
-from qdrant_client.http.models import CountResult, ScoredPoint, Record
+from qdrant_client.http.models import CountResult, Record
 
 from ...services import models
 from .vector_store import VectorStore
diff --git a/ui/src/pages/DataSources/VisualizationTab/VectorGraph.tsx b/ui/src/pages/DataSources/VisualizationTab/VectorGraph.tsx
index 12acaf4..9357d79 100644
--- a/ui/src/pages/DataSources/VisualizationTab/VectorGraph.tsx
+++ b/ui/src/pages/DataSources/VisualizationTab/VectorGraph.tsx
@@ -71,19 +71,20 @@ const VectorGraph = ({ rawData }: { rawData: Point2d[] }) => {
       hash = (hash << 5) - hash + str.charCodeAt(i);
       hash |= 0; // Convert to 32bit integer
     }
-    return Math.abs(hash % 10);
+    return Math.abs(hash % colors.length);
   };
 
+  const pickColor = (label: string) => colors[hashStringToIndex(label)];
+
   const vizDatasets: DataSets = Object.entries(points).map(
     ([label, points]) => {
       const userQuery = label === "USER_QUERY";
+      const color = pickColor(label);
       return {
         label: userQuery ? "User Query" : label,
         data: points,
-        backgroundColor: userQuery
-          ? "lightgray"
-          : colors[hashStringToIndex(label)],
-        borderColor: userQuery ? "black" : colors[hashStringToIndex(label)],
+        backgroundColor: userQuery ? "lightgray" : color,
+        borderColor: userQuery ? "black" : color,
         borderWidth: 1,
         pointStyle: userQuery ? "circle" : "circle",
         pointRadius: userQuery ? 15 : 3,