fix small bug

generate.py

@@ -1,6 +1,6 @@
 import os
 import re
-from typing import Optional
+from typing import Optional, Union
 from llms import build_prompt, get_completion, PROMPT_TEMPLATE, SYSTEM_PROMPT
 
 
@@ -26,15 +26,17 @@ def extract_code(completion: str) -> str:
     if code_block:
         return code_block.group(1)
 
-def get_solution_file_path(model: str, day: Optional[str | int] = None):
+def get_solution_file_path(model: str, day: Optional[Union[str, int]] = None):
     """Returns the path to the solution file for the given day and model. If day is None, returns the path to the solution file only."""
 
+    if day is None:
+        return f"solution_{model}.py"
+
+    # We want it formatted properly, so we convert to int and back if its already a string for convenience
     if isinstance(day, str):
         day = int(day)
-    day_str = f"{day:02d}"
 
-    if day is None:
-        return f"solution_{model}.py"
+    day_str = f"{day:02d}"
 
     return f"day{day_str}/solution_{model}.py"