diff --git a/prophet/app.py b/prophet/app.py index b7839b8..2c8b296 100644 --- a/prophet/app.py +++ b/prophet/app.py @@ -1,6 +1,5 @@ import hashlib import json -import os import pickle import re from dataclasses import dataclass, field @@ -15,6 +14,8 @@ from fastapi.responses import HTMLResponse from fastapi_utils.tasks import repeat_every from groq import Groq +from prophet.config import AiConfig + BEE_FEED = "https://babylonbee.com/feed" BEE_FEED_TEST = "test/resources/feed_short.atom" # NOTE: Switch out when done testing @@ -22,6 +23,8 @@ PICKLE_DIR = "/tmp/pollenprophet" REFRESH_PERIOD = 3600 # between fetching articles, in seconds +config_ai: AiConfig = AiConfig.from_env() + @dataclass class Original: # BadJoke: Sting @@ -133,7 +136,7 @@ def improve_originals(originals: list[Original]) -> list[Improvement]: def rewrite_title_with_groq(original_content: str) -> str: - client = Groq(api_key=os.getenv("GROQ_API_KEY", "NO_API_KEY_FOUND")) + client = Groq(api_key=config_ai.API_KEY) suggestions = client.chat.completions.create( messages=[ @@ -173,7 +176,7 @@ def rewrite_title_with_groq(original_content: str) -> str: def rewrite_summary_with_groq(orig: Original, improved_title: str) -> str: - client = Groq(api_key=os.getenv("GROQ_API_KEY", "NO_API_KEY_FOUND")) + client = Groq(api_key=config_ai.API_KEY) summary = client.chat.completions.create( messages=[ diff --git a/prophet/config.py b/prophet/config.py new file mode 100644 index 0000000..55a8702 --- /dev/null +++ b/prophet/config.py @@ -0,0 +1,22 @@ +import os + +# Load environment variables from .env +from dataclasses import dataclass + +from dotenv import load_dotenv + +_ = load_dotenv() + + +@dataclass +class AiConfig: + API_KEY: str + + @classmethod + def from_env(cls) -> "AiConfig": + API_KEY = os.getenv("GROQ_API_KEY", "") + + if not API_KEY: + raise ValueError(f"{API_KEY} cannot be empty") + + return cls(**{"API_KEY": API_KEY})