48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
"""AsyncParallelBatchNode implementation for article summarization."""
|
|
|
|
from pocketflow import AsyncParallelBatchNode, AsyncNode
|
|
from utils import call_llm_async, load_articles, save_summaries
|
|
|
|
class LoadArticles(AsyncNode):
|
|
"""Node that loads articles to process."""
|
|
|
|
async def prep_async(self, shared):
|
|
"""Load articles from data directory."""
|
|
print("\nLoading articles...")
|
|
articles = await load_articles()
|
|
return articles
|
|
|
|
async def exec_async(self, articles):
|
|
"""No processing needed."""
|
|
return articles
|
|
|
|
async def post_async(self, shared, prep_res, exec_res):
|
|
"""Store articles in shared store."""
|
|
shared["articles"] = exec_res
|
|
print(f"Found {len(exec_res)} articles to process")
|
|
return "process"
|
|
|
|
class ParallelSummarizer(AsyncParallelBatchNode):
|
|
"""Node that summarizes articles in parallel."""
|
|
|
|
async def prep_async(self, shared):
|
|
"""Get articles from shared store."""
|
|
print("\nProcessing in parallel...")
|
|
return shared["articles"]
|
|
|
|
async def exec_async(self, article):
|
|
"""Summarize a single article (called in parallel)."""
|
|
summary = await call_llm_async(article)
|
|
return summary
|
|
|
|
async def post_async(self, shared, prep_res, summaries):
|
|
"""Store summaries and save to file."""
|
|
shared["summaries"] = summaries
|
|
|
|
print("\nSummaries generated:")
|
|
for i, summary in enumerate(summaries, 1):
|
|
print(f"{i}. {summary}")
|
|
|
|
save_summaries(summaries)
|
|
print("\nFinal report saved to: summaries.txt")
|
|
return "default" |