diff --git a/benchmark/over_time.py b/benchmark/over_time.py
index 565038a8e..f72bac31e 100644
--- a/benchmark/over_time.py
+++ b/benchmark/over_time.py
@@ -6,6 +6,17 @@ from matplotlib import rc
 from aider.dump import dump  # noqa: 401
 
 
+def get_model_color(model):
+    if "-4o" in model and "gpt-4o-mini" not in model:
+        return "purple"
+    elif "gpt-4" in model:
+        return "red"
+    elif "gpt-3.5" in model:
+        return "green"
+    else:
+        return "lightblue"
+
+
 def plot_over_time(yaml_file):
     with open(yaml_file, "r") as file:
         data = yaml.safe_load(file)
@@ -49,14 +60,7 @@ def plot_over_time(yaml_file):
         spine.set_edgecolor("#DDDDDD")
         spine.set_linewidth(0.5)
 
-    colors = [
-        (
-            "purple"
-            if "-4o" in model and "gpt-4o-mini" not in model
-            else "red" if "gpt-4" in model else "green" if "gpt-3.5" in model else "lightblue"
-        )
-        for model in models
-    ]
+    colors = [get_model_color(model) for model in models]
 
     # Separate data points by color
     purple_points = [(d, r) for d, r, c in zip(dates, pass_rates, colors) if c == "purple"]