Undo API key

src/fireworks/inference.py

@@ -7,25 +7,22 @@ from src.config import EMBEDDING_MODEL
 load_dotenv()
 
 
-def create_client() -> OpenAI:
+def create_client(api_key: str = None) -> OpenAI:
     """
     Create client for FW inference
     """
-    assert (
-        os.getenv("FIREWORKS_API_KEY") is not None
-    ), "FIREWORKS_API_KEY not found in environment variables"
+    api_key = os.getenv("FIREWORKS_API_KEY", api_key)
+    assert api_key is not None, "FIREWORKS_API_KEY not found in environment variables"
     return OpenAI(
-        api_key=os.getenv("FIREWORKS_API_KEY"),
+        api_key=api_key,
         base_url="https://api.fireworks.ai/inference/v1",
     )
 
 
-FW_CLIENT = create_client()
-
-
 def get_embedding(text: str) -> List[float]:
     """
     Get embedding for a given text
     """
-    response = FW_CLIENT.embeddings.create(model=EMBEDDING_MODEL, input=text)
+    client = create_client()
+    response = client.embeddings.create(model=EMBEDDING_MODEL, input=text)
     return response.data[0].embedding