change to ues api keys

Author: DJSaunders1997

backend/fastapi_generate_quiz.py

@@ -2,12 +2,18 @@
 # AI-powered quiz generation and image creation service
 # https://platform.openai.com/docs/api-reference/streaming
 import logging
-from generate_quiz import QuizGenerator
-from generate_image import ImageGenerator
+from typing import Optional
+
+from dotenv import load_dotenv
 from fastapi import FastAPI, Query
-from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from typing import Optional
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from generate_image import ImageGenerator
+from generate_quiz import QuizGenerator
+
+# Load environment variables from .env file
+load_dotenv()
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)

backend/generate_quiz.py

@@ -1,9 +1,15 @@
-from typing import Generator, Optional
-import logging
 import json
+import logging
 import os
-from response_stream_parser import ResponseStreamParser
+from typing import Generator, Optional
+
 import litellm
+from dotenv import load_dotenv
+
+from response_stream_parser import ResponseStreamParser
+
+# Load environment variables
+load_dotenv()
 
 # Set up logging
 logger = logging.getLogger(__name__)
@@ -66,25 +72,23 @@ class QuizGenerator:
 
     @classmethod
     def check_api_key_from_env(cls) -> None:
-        """Retrieves the API keys from environment variables.
+        """Check if at least one API key is available.
 
         Raises:
-            ValueError: If the environment variable is not set or empty.
+            ValueError: If no API keys are found.
         """
-
-        for key in [
-            "OPENAI_API_KEY",
-            "GEMINI_API_KEY",
-            "DEEPSEEK_API_KEY",
-            "AZURE_AI_API_KEY",
-            "AZURE_AI_API_BASE",
-        ]:
-            api_key = os.getenv(key)
-            if not api_key:
-                raise ValueError(
-                    f"Environment variable {key} is not set."
-                    "Please ensure it's set and try again."
-                )
+        api_keys = [
+            os.getenv("OPENAI_API_KEY"),
+            os.getenv("GEMINI_API_KEY"),
+            os.getenv("DEEPSEEK_API_KEY"),
+            os.getenv("AZURE_AI_API_KEY"),
+        ]
+        
+        if not any(api_keys):
+            raise ValueError(
+                "No API keys found. Please set at least one of: "
+                "OPENAI_API_KEY, GEMINI_API_KEY, DEEPSEEK_API_KEY, or AZURE_AI_API_KEY"
+            )
 
     @staticmethod
     def check_model_is_supported(model: str) -> str:
@@ -218,7 +222,7 @@ def print_quiz(generator: Generator[str, None, None]):
         "azure_ai/DeepSeek-R1",
     ]
 
-    quiz_generator = QuizGenerator(model="gemini/gemini-2.0-flash")
+    quiz_generator = QuizGenerator(model="o1-mini")
 
     topic = "Crested Gecko"
     difficulty = "Medium"

backend/pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
     "fastapi",
     "uvicorn",
     "litellm",
+    "python-dotenv",
 ]
 
 [project.optional-dependencies]

backend/tests/test_generate_quiz.py

@@ -43,14 +43,16 @@ def test_check_model_is_supported(self):
 
     def test_environment_variable_not_set(self, monkeypatch):
         """
-        Test that initializing QuizGenerator without an API key (i.e. if the environment variable
-        is not set) raises a ValueError.
-
-        We remove the environment variable and expect the constructor to raise an error.
+        Test that initializing QuizGenerator without any API keys raises a ValueError.
         """
+        # Remove all API keys
         monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+        monkeypatch.delenv("GEMINI_API_KEY", raising=False)
+        monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False)
+        monkeypatch.delenv("AZURE_AI_API_KEY", raising=False)
+        
         with pytest.raises(
-            ValueError, match="Environment variable OPENAI_API_KEY is not set"
+            ValueError, match="No API keys found"
         ):
             QuizGenerator()
 
@@ -70,7 +72,7 @@ def test_create_role(self, quiz_generator):
         assert str(n_questions) in role
         assert quiz_generator.EXAMPLE_RESPONSE in role
 
-    @patch("backend.generate_quiz.completion")
+    @patch("backend.generate_quiz.litellm.completion")
     def test_generate_quiz(self, mock_completion, quiz_generator):
         """Test generate_quiz to ensure it streams responses properly."""
         mock_stream = iter(['{"question": "What is 2+2?", "answer": "4"}\n'])