From a032ecbc55cac85b9c457a46bbdf5a5e45b5ca18 Mon Sep 17 00:00:00 2001 From: Marty Oehme Date: Thu, 5 Jun 2025 17:49:45 +0200 Subject: [PATCH] Save and load with pickle --- prophet/app.py | 42 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/prophet/app.py b/prophet/app.py index e0d2348..07d0298 100644 --- a/prophet/app.py +++ b/prophet/app.py @@ -1,6 +1,9 @@ import os +import pickle from dataclasses import dataclass from datetime import datetime +from pathlib import Path +from pprint import pprint from uuid import uuid4 import feedparser @@ -11,6 +14,8 @@ from groq import Groq BEE_FEED = "https://babylonbee.com/feed" BEE_FEED_TEST = "test/resources/feed.atom" # NOTE: Switch out when done testing +PICKLE_DIR = "/tmp/pollenprophet" + @dataclass class Original: # BadJoke: Sting @@ -45,6 +50,34 @@ def grab_latest_originals() -> list[Original]: return results +def save_new_improvements(improvements: list[Improvement]) -> None: + save_dir = Path(PICKLE_DIR) + save_dir.mkdir(parents=True, exist_ok=True) + for imp in improvements: + fname = save_dir / f"{int(imp.original.date.timestamp())}_{imp.id}" + try: + with open(fname, "wb") as f: + pickle.dump(imp, f) + print(f"Saved {fname}") + except Exception as e: + print(f"Error saving file {fname}: {e}") + + +def load_existing_improvements() -> list[Improvement]: + improvements: list[Improvement] = [] + for fname in Path(PICKLE_DIR).iterdir(): + if not fname.is_file(): + continue + + try: + with open(fname, "rb") as f: + obj: Improvement = pickle.load(f) + improvements.append(obj) + except FileNotFoundError as e: + print(f"Error loading file {fname}: {e}") + return improvements + + def improve_with_groq(original: str) -> str: client = Groq(api_key=os.getenv("GROQ_API_KEY", "NO_API_KEY_FOUND")) @@ -98,7 +131,14 @@ def improve_headline(content: str): def start() -> None: from uvicorn import run - grab_latest_originals() + orig = grab_latest_originals() + improvements = [ + Improvement(original=o, title="hithere", summary="alongsummary") + for o in orig + ] + save_new_improvements(improvements) + + pprint(load_existing_improvements()) # run("prophet.app:app", reload=True)