fix: implement backwards-compatible StrEnum

awslabs · Apr 8, 2024 · 7d8072b · 7d8072b
1 parent 289f57d
commit 7d8072b
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/src/agenteval/evaluators/aws/bedrock/claude/claude_evaluator.py b/src/agenteval/evaluators/aws/bedrock/claude/claude_evaluator.py
@@ -9,16 +9,6 @@
 from agenteval.evaluators.aws.bedrock.claude import model_configs
 from agenteval.test_result import TestResult
 
-# enable backwards-compatible StrEnum
-try:
-    from enum import StrEnum
-except ImportError:
-    from enum import Enum
-
-    class StrEnum(str, Enum):
-        pass
-
-
 logger = logging.getLogger(__name__)
 
 _PROMPT_TEMPLATE_ROOT = "evaluators/claude"
@@ -30,6 +20,15 @@ class StrEnum(str, Enum):
     "generate_evaluation",
 ]
 
+# enable backwards-compatible StrEnum
+try:
+    from enum import StrEnum
+except ImportError:
+    from enum import Enum
+
+    class StrEnum(str, Enum):
+        pass
+
 
 class TestStatusCategories(StrEnum):
     ALL_STEPS_ATTEMPTED = "A"