fix: types for the auto param in MIPRO (#2018)

Co-authored-by: Jamesb <experimentallearning0@gmail.com>
stanfordnlp · Jan 8, 2025 · 238e312 · 238e312
1 parent 9515397
commit 238e312
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dspy/teleprompt/mipro_optimizer_v2.py b/dspy/teleprompt/mipro_optimizer_v2.py
@@ -2,7 +2,7 @@
 import random
 import textwrap
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
 
 import numpy as np
 import optuna
@@ -52,7 +52,7 @@ def __init__(
         teacher_settings: Dict = {},
         max_bootstrapped_demos: int = 4,
         max_labeled_demos: int = 16,
-        auto: Optional[str] = None,
+        auto: Optional[Literal["light", "medium", "heavy"]] = None,
         num_candidates: int = 10,
         num_threads: int = 6,
         max_errors: int = 10,