Create Config dataclass for LLM options

This commit is contained in:
Marty Oehme 2025-06-06 10:22:14 +02:00
parent 3668cc3cbd
commit f96b6413e2
Signed by: Marty
GPG key ID: 4E535BC19C61886E
2 changed files with 28 additions and 3 deletions

View file

@ -1,6 +1,5 @@
import hashlib import hashlib
import json import json
import os
import pickle import pickle
import re import re
from dataclasses import dataclass, field from dataclasses import dataclass, field
@ -15,6 +14,8 @@ from fastapi.responses import HTMLResponse
from fastapi_utils.tasks import repeat_every from fastapi_utils.tasks import repeat_every
from groq import Groq from groq import Groq
from prophet.config import AiConfig
BEE_FEED = "https://babylonbee.com/feed" BEE_FEED = "https://babylonbee.com/feed"
BEE_FEED_TEST = "test/resources/feed_short.atom" # NOTE: Switch out when done testing 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 REFRESH_PERIOD = 3600 # between fetching articles, in seconds
config_ai: AiConfig = AiConfig.from_env()
@dataclass @dataclass
class Original: # BadJoke: Sting 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: 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( suggestions = client.chat.completions.create(
messages=[ 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: 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( summary = client.chat.completions.create(
messages=[ messages=[

22
prophet/config.py Normal file
View file

@ -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})