arxiv_audio_summary/vibe/main.py

58 lines
2.4 KiB
Python
Raw Normal View History

2025-03-02 03:22:35 +00:00
import argparse
import logging
from vibe.orchestrator import process_articles
from vibe.tts import text_to_speech
2025-03-02 12:47:03 +00:00
from vibe.config import DEFAULT_ARXIV_URL
2025-03-02 03:22:35 +00:00
logging.basicConfig(
level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
def main():
parser = argparse.ArgumentParser(description="vibe: Article Summarization & TTS Pipeline")
parser.add_argument("--serve", action="store_true", help="Run as a Flask server.")
parser.add_argument("--generate", action="store_true", help="Run the pipeline once and generate a summary MP3, then exit.")
parser.add_argument("--prompt", type=str, default="", help="User info for LLM filtering & summaries.")
parser.add_argument("--max-articles", type=int, default=5, help="Maximum articles to process in the pipeline.")
parser.add_argument("--new-only", action="store_true", help="Only process articles newer than cached.")
parser.add_argument("--arxiv-url", type=str, default=DEFAULT_ARXIV_URL, help="URL for fetching arXiv articles.")
parser.add_argument("--output", type=str, default="final_output.mp3", help="Output path for the generated MP3 file.")
2025-03-02 12:47:03 +00:00
# New: LLM Level
parser.add_argument("--llm-level", type=str, default="medium", choices=["low","medium","high"],
help="Desired LLM quality level: low, medium, or high. Defaults to medium.")
2025-03-02 03:22:35 +00:00
args = parser.parse_args()
if args.serve:
from vibe.server import app
logger.info("Starting Flask server.")
2025-03-02 12:47:03 +00:00
app.run(host='0.0.0.0', port='14200', debug=True)
2025-03-02 03:22:35 +00:00
elif args.generate:
logger.info("Running pipeline in CLI mode.")
user_info = args.prompt
2025-03-02 12:47:03 +00:00
final_summary = process_articles(
user_info,
arxiv_url=args.arxiv_url,
max_articles=args.max_articles,
new_only=args.new_only,
llm_level=args.llm_level
)
2025-03-02 03:22:35 +00:00
if not final_summary.strip():
logger.error("No summaries generated.")
exit(1)
try:
text_to_speech(final_summary, args.output)
logger.info(f"Generated MP3 at: {args.output}")
except Exception as e:
logger.exception("TTS conversion failed: %s", e)
exit(1)
else:
logger.info("No mode specified; defaulting to Flask server.")
from vibe.server import app
2025-03-02 12:47:03 +00:00
app.run(host='0.0.0.0', port='14200', debug=True)
2025-03-02 03:22:35 +00:00
if __name__ == "__main__":
2025-03-02 12:47:03 +00:00
main()