diff --git a/servers/fastapi/api/lifespan.py b/servers/fastapi/api/lifespan.py index a74785dc..184a55eb 100644 --- a/servers/fastapi/api/lifespan.py +++ b/servers/fastapi/api/lifespan.py @@ -5,6 +5,7 @@ from fastapi import FastAPI from sqlmodel import SQLModel from services import SQL_ENGINE +from utils.get_env import get_app_data_directory_env from utils.model_availability import check_llm_and_image_provider_api_or_model_availability @@ -15,7 +16,7 @@ async def app_lifespan(_: FastAPI): Initializes the application data directory and checks LLM model availability. """ - os.makedirs(os.getenv("APP_DATA_DIRECTORY"), exist_ok=True) + os.makedirs(get_app_data_directory_env(), exist_ok=True) SQLModel.metadata.create_all(SQL_ENGINE) await check_llm_and_image_provider_api_or_model_availability() yield diff --git a/servers/fastapi/services/image_generation_service.py b/servers/fastapi/services/image_generation_service.py index 4a1d979c..2bad97dd 100644 --- a/servers/fastapi/services/image_generation_service.py +++ b/servers/fastapi/services/image_generation_service.py @@ -122,7 +122,7 @@ class ImageGenerationService: async def get_image_from_pixabay(self, prompt: str) -> str: async with aiohttp.ClientSession() as session: response = await session.get( - f"https://pixabay.com/api/?key={os.getenv('PIXABAY_API_KEY')}&q={prompt}&image_type=photo&per_page=3" + f"https://pixabay.com/api/?key={get_pixabay_api_key_env()}&q={prompt}&image_type=photo&per_page=3" ) data = await response.json() image_url = data["hits"][0]["largeImageURL"] diff --git a/servers/fastapi/services/redis_service.py b/servers/fastapi/services/redis_service.py index 5d9dff22..f2e3d8c9 100644 --- a/servers/fastapi/services/redis_service.py +++ b/servers/fastapi/services/redis_service.py @@ -1,15 +1,21 @@ -import os from typing import Any, Optional import redis from redis.exceptions import RedisError +from utils.get_env import ( + get_redis_db_env, + get_redis_host_env, + get_redis_password_env, + get_redis_port_env, +) + class RedisService: def __init__(self): - self.redis_host = os.getenv("REDIS_HOST", "localhost") - self.redis_port = int(os.getenv("REDIS_PORT", "6379")) - self.redis_db = int(os.getenv("REDIS_DB", "0")) - self.redis_password = os.getenv("REDIS_PASSWORD") + self.redis_host = get_redis_host_env() or "localhost" + self.redis_port = int(get_redis_port_env() or "6379") + self.redis_db = int(get_redis_db_env() or "0") + self.redis_password = get_redis_password_env() or None self.client = self._create_client() def _create_client(self) -> redis.Redis: diff --git a/servers/fastapi/services/temp_file_service.py b/servers/fastapi/services/temp_file_service.py index 31a39035..f4c59cf5 100644 --- a/servers/fastapi/services/temp_file_service.py +++ b/servers/fastapi/services/temp_file_service.py @@ -2,11 +2,13 @@ import os import uuid from typing import Optional, Union +from utils.get_env import get_temp_directory_env + class TempFileService: def __init__(self): - self.base_dir = os.getenv("TEMP_DIRECTORY") + self.base_dir = get_temp_directory_env() # TODO: Uncomment this when we want to cleanup the base dir on startup # self.cleanup_base_dir() os.makedirs(self.base_dir, exist_ok=True) diff --git a/servers/fastapi/utils/get_env.py b/servers/fastapi/utils/get_env.py index f8ca88c7..265e3698 100644 --- a/servers/fastapi/utils/get_env.py +++ b/servers/fastapi/utils/get_env.py @@ -56,8 +56,26 @@ def get_custom_model_env(): def get_pexels_api_key_env(): return os.getenv("PEXELS_API_KEY") + def get_image_provider_env(): return os.getenv("IMAGE_PROVIDER") + def get_pixabay_api_key_env(): - return os.getenv("PIXABAY_API_KEY") \ No newline at end of file + return os.getenv("PIXABAY_API_KEY") + + +def get_redis_host_env(): + return os.getenv("REDIS_HOST") + + +def get_redis_port_env(): + return os.getenv("REDIS_PORT") + + +def get_redis_db_env(): + return os.getenv("REDIS_DB") + + +def get_redis_password_env(): + return os.getenv("REDIS_PASSWORD") diff --git a/servers/fastapi/utils/image_provider.py b/servers/fastapi/utils/image_provider.py index 1dd8f217..53167d1e 100644 --- a/servers/fastapi/utils/image_provider.py +++ b/servers/fastapi/utils/image_provider.py @@ -1,5 +1,12 @@ import os from enums.image_provider import ImageProvider +from utils.get_env import ( + get_google_api_key_env, + get_image_provider_env, + get_openai_api_key_env, + get_pexels_api_key_env, + get_pixabay_api_key_env, +) def is_pixels_selected() -> bool: @@ -24,18 +31,18 @@ def get_selected_image_provider() -> ImageProvider: Returns: ImageProvider: The selected image provider. """ - return ImageProvider(os.getenv("IMAGE_PROVIDER")) + return ImageProvider(get_image_provider_env()) def get_image_provider_api_key() -> str: selected_image_provider = get_selected_image_provider() if selected_image_provider == ImageProvider.PEXELS: - return os.getenv("PEXELS_API_KEY") + return get_pexels_api_key_env() elif selected_image_provider == ImageProvider.PIXABAY: - return os.getenv("PIXABAY_API_KEY") + return get_pixabay_api_key_env() elif selected_image_provider == ImageProvider.GEMINI_FLASH: - return os.getenv("GOOGLE_API_KEY") + return get_google_api_key_env() elif selected_image_provider == ImageProvider.DALLE3: - return os.getenv("OPENAI_API_KEY") + return get_openai_api_key_env() else: raise ValueError(f"Invalid image provider: {selected_image_provider}") diff --git a/servers/fastapi/utils/llm_provider.py b/servers/fastapi/utils/llm_provider.py index bb069773..7999ad4a 100644 --- a/servers/fastapi/utils/llm_provider.py +++ b/servers/fastapi/utils/llm_provider.py @@ -7,8 +7,10 @@ from enums.llm_provider import LLMProvider from utils.get_env import ( get_custom_llm_api_key_env, get_custom_llm_url_env, + get_custom_model_env, get_google_api_key_env, get_llm_provider_env, + get_ollama_model_env, get_ollama_url_env, get_openai_api_key_env, ) @@ -93,9 +95,9 @@ def get_large_model(): elif selected_llm == LLMProvider.GOOGLE: return "gemini-2.0-flash" elif selected_llm == LLMProvider.OLLAMA: - return os.getenv("OLLAMA_MODEL") + return get_ollama_model_env() elif selected_llm == LLMProvider.CUSTOM: - return os.getenv("CUSTOM_MODEL") + return get_custom_model_env() else: raise ValueError(f"Invalid LLM model") @@ -107,9 +109,9 @@ def get_small_model(): elif selected_llm == LLMProvider.GOOGLE: return "gemini-2.0-flash" elif selected_llm == LLMProvider.OLLAMA: - return os.getenv("OLLAMA_MODEL") + return get_ollama_model_env() elif selected_llm == LLMProvider.CUSTOM: - return os.getenv("CUSTOM_MODEL") + return get_custom_model_env() else: raise ValueError(f"Invalid LLM model") @@ -121,8 +123,8 @@ def get_nano_model(): elif selected_llm == LLMProvider.GOOGLE: return "gemini-2.0-flash" elif selected_llm == LLMProvider.OLLAMA: - return os.getenv("OLLAMA_MODEL") + return get_ollama_model_env() elif selected_llm == LLMProvider.CUSTOM: - return os.getenv("CUSTOM_MODEL") + return get_custom_model_env() else: raise ValueError(f"Invalid LLM model") diff --git a/servers/fastapi/utils/model_availability.py b/servers/fastapi/utils/model_availability.py index 4b3f1c10..c1dce981 100644 --- a/servers/fastapi/utils/model_availability.py +++ b/servers/fastapi/utils/model_availability.py @@ -2,7 +2,17 @@ import os from constants.supported_ollama_models import SUPPORTED_OLLAMA_MODELS from enums.llm_provider import LLMProvider from utils.custom_llm_provider import list_available_custom_models -from utils.get_env import get_can_change_keys_env +from utils.get_env import ( + get_can_change_keys_env, + get_openai_api_key_env, + get_pixabay_api_key_env, + get_pexels_api_key_env, +) +from utils.get_env import get_google_api_key_env +from utils.get_env import get_ollama_model_env +from utils.get_env import get_custom_llm_api_key_env +from utils.get_env import get_custom_llm_url_env +from utils.get_env import get_custom_model_env from utils.llm_provider import ( get_llm_provider, is_custom_llm_selected, @@ -16,21 +26,22 @@ from utils.image_provider import ( is_dalle3_selected, ) + async def check_llm_and_image_provider_api_or_model_availability(): can_change_keys = get_can_change_keys_env() != "false" if not can_change_keys: if get_llm_provider() == LLMProvider.OPENAI: - openai_api_key = os.getenv("OPENAI_API_KEY") + openai_api_key = get_openai_api_key_env() if not openai_api_key: raise Exception("OPENAI_API_KEY must be provided") elif get_llm_provider() == LLMProvider.GOOGLE: - google_api_key = os.getenv("GOOGLE_API_KEY") + google_api_key = get_google_api_key_env() if not google_api_key: raise Exception("GOOGLE_API_KEY must be provided") elif is_ollama_selected(): - ollama_model = os.getenv("OLLAMA_MODEL") + ollama_model = get_ollama_model_env() if not ollama_model: raise Exception("OLLAMA_MODEL must be provided") @@ -45,9 +56,9 @@ async def check_llm_and_image_provider_api_or_model_availability(): print("-" * 50) elif is_custom_llm_selected(): - custom_model = os.getenv("CUSTOM_MODEL") - custom_llm_url = os.getenv("CUSTOM_LLM_URL") - custom_llm_api_key = os.getenv("CUSTOM_LLM_API_KEY") + custom_model = get_custom_model_env() + custom_llm_url = get_custom_llm_url_env() + custom_llm_api_key = get_custom_llm_api_key_env() if not custom_model: raise Exception("CUSTOM_MODEL must be provided") if not custom_llm_url: @@ -64,21 +75,21 @@ async def check_llm_and_image_provider_api_or_model_availability(): if custom_model not in models: raise Exception(f"Model {custom_model} is not available") elif is_pixels_selected(): - pexels_api_key = os.getenv("PEXELS_API_KEY") + pexels_api_key = get_pexels_api_key_env() if not pexels_api_key: raise Exception("PEXELS_API_KEY must be provided") elif is_pixabay_selected(): - pixabay_api_key = os.getenv("PIXABAY_API_KEY") + pixabay_api_key = get_pixabay_api_key_env() if not pixabay_api_key: raise Exception("PIXABAY_API_KEY must be provided") elif is_gemini_flash_selected(): - google_api_key = os.getenv("GOOGLE_API_KEY") + google_api_key = get_google_api_key_env() if not google_api_key: raise Exception("GOOGLE_API_KEY must be provided") elif is_dalle3_selected(): - openai_api_key = os.getenv("OPENAI_API_KEY") + openai_api_key = get_openai_api_key_env() if not openai_api_key: - raise Exception("OPENAI_API_KEY must be provided") \ No newline at end of file + raise Exception("OPENAI_API_KEY must be provided")