From acf754f7a8b247cad9d7f9ca27631a58a8bcba82 Mon Sep 17 00:00:00 2001
From: Nathan Nowack <thrast36@gmail.com>
Date: Tue, 5 Dec 2023 22:49:09 -0600
Subject: [PATCH] mark flaky test / fix more imports

---
 src/marvin/components/speech.py        | 3 +--
 src/marvin/settings.py                 | 3 ++-
 tests/components/test_ai_classifier.py | 2 ++
 tests/components/test_ai_model.py      | 3 ++-
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/marvin/components/speech.py b/src/marvin/components/speech.py
index fbc45595d..0bbcfe00f 100644
--- a/src/marvin/components/speech.py
+++ b/src/marvin/components/speech.py
@@ -4,12 +4,11 @@
     Any,
     Callable,
     Coroutine,
-    Literal,
     Optional,
     TypeVar,
 )
 
-from typing_extensions import ParamSpec
+from typing_extensions import Literal, ParamSpec
 
 if TYPE_CHECKING:
     from openai._base_client import HttpxBinaryResponseContent
diff --git a/src/marvin/settings.py b/src/marvin/settings.py
index 413d7ab12..45a74479d 100644
--- a/src/marvin/settings.py
+++ b/src/marvin/settings.py
@@ -12,10 +12,11 @@
 import os
 from contextlib import contextmanager
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from pydantic import Field, SecretStr
 from pydantic_settings import BaseSettings, SettingsConfigDict
+from typing_extensions import Literal
 
 if TYPE_CHECKING:
     from openai import AsyncClient, Client
diff --git a/tests/components/test_ai_classifier.py b/tests/components/test_ai_classifier.py
index 74755625f..ec64f79a0 100644
--- a/tests/components/test_ai_classifier.py
+++ b/tests/components/test_ai_classifier.py
@@ -1,5 +1,6 @@
 from enum import Enum
 
+import pytest
 from marvin import ai_classifier
 from typing_extensions import Literal
 
@@ -27,6 +28,7 @@ def sentiment(text: str) -> Sentiment:
 
             assert result == "Positive"
 
+        @pytest.mark.flaky(reruns=3)
         def test_ai_classifier_literal_return_type_with_docstring(self):
             @ai_classifier
             def sentiment(text: str) -> Sentiment:
diff --git a/tests/components/test_ai_model.py b/tests/components/test_ai_model.py
index 0fb9efa38..00885bf9f 100644
--- a/tests/components/test_ai_model.py
+++ b/tests/components/test_ai_model.py
@@ -1,8 +1,9 @@
-from typing import List, Literal, Optional
+from typing import List, Optional
 
 import pytest
 from marvin import ai_model
 from pydantic import BaseModel, Field
+from typing_extensions import Literal
 
 from tests.utils import pytest_mark_class