workflow.py

#!/usr/bin/env python
"""
AI Paper Recommendation Workflow - Main Entry Point

This workflow:
1. Fetches latest papers from arXiv
2. Fetches ModelScope/魔搭社区 Papers
3. Filters papers based on user interests (limit to max_papers_per_day)
4. Downloads PDFs and extracts images
5. Parses PDF content with Chinese summary
6. Generates Obsidian notes with embedded images
"""

import os
import sys
import argparse
import logging
from datetime import datetime

# Add project root to path
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

from src.utils import load_config, setup_logging, ensure_dir, get_papers_dir
from src.fetch_papers import fetch_arxiv_papers, save_papers as save_arxiv_papers
from src.fetch_modelscope import fetch_modelscope_daily_papers, save_modelscope_papers
from src.filter_papers import filter_papers, save_filtered_papers, load_filtered_papers
from src.download_images import process_all_papers_images
from src.parse_paper import process_papers_with_pdfs
from src.generate_note import generate_all_notes


def run_workflow(config, args):
    """
    Run the complete paper recommendation workflow

    Args:
        config: Configuration dictionary
        args: Command line arguments
    """
    logger = logging.getLogger(__name__)

    # Get max papers per day from config
    max_papers = config.get('max_papers_per_day', 8)
    # Get paper distribution: 5 from arXiv + 3 from ModelScope
    arxiv_limit = config.get('arxiv_limit', 5)
    modelscope_limit = config.get('modelscope_limit', 3)

    # Step 1: Fetch arXiv papers
    if args.fetch_arxiv or args.all:
        logger.info("=" * 50)
        logger.info("Step 1: Fetching arXiv papers...")
        logger.info("=" * 50)
        papers = fetch_arxiv_papers(config, days_back=7, max_results=80)
        save_arxiv_papers(papers)
        logger.info(f"Fetched {len(papers)} arXiv papers")

    # Step 2: Fetch ModelScope/魔搭社区 Papers
    if args.fetch_modelscope or args.all:
        logger.info("=" * 50)
        logger.info("Step 2: Fetching ModelScope/魔搭社区 Papers...")
        logger.info("=" * 50)
        ms_papers = fetch_modelscope_daily_papers(config, max_papers=20)
        save_modelscope_papers(ms_papers)
        logger.info(f"Fetched {len(ms_papers)} ModelScope papers")

    # Step 3: Filter papers (5 from arXiv + 3 from ModelScope)
    if args.filter or args.all:
        logger.info("=" * 50)
        logger.info("Step 3: Filtering papers based on interests...")
        logger.info(f"         Target: {arxiv_limit} arXiv + {modelscope_limit} ModelScope")
        logger.info("=" * 50)
        filtered = filter_papers(config, arxiv_limit=arxiv_limit, modelscope_limit=modelscope_limit)

        save_filtered_papers(filtered)
        logger.info(f"Filtered to {len(filtered)} relevant papers")

        # Print top papers
        logger.info(f"\nTop {len(filtered)} relevant papers:")
        for i, paper in enumerate(filtered):
            source = paper.get('source_type', 'Unknown')
            logger.info(f"  {i+1}. [{source}] {paper.get('title', '')[:50]}...")
            logger.info(f"     Keywords: {', '.join(paper.get('matched_keywords', []))}")
            logger.info(f"     Score: {paper.get('relevance', 0):.2f}")

    # Step 4: Download PDFs and images
    if args.download or args.all:
        logger.info("=" * 50)
        logger.info("Step 4: Downloading PDFs and extracting images...")
        logger.info("=" * 50)
        papers = load_filtered_papers()

        papers = process_all_papers_images(config, papers)
        save_filtered_papers(papers)  # Save with updated image paths
        logger.info("Downloaded PDFs and extracted images")

    # Step 5: Parse PDF content with Chinese summary
    if args.parse or args.all:
        logger.info("=" * 50)
        logger.info("Step 5: Parsing PDF content with Chinese summary...")
        logger.info("=" * 50)
        papers = load_filtered_papers()

        parsed = process_papers_with_pdfs(config, papers)
        logger.info(f"Parsed {len(parsed)} papers with Chinese summaries")

    # Step 6: Generate Obsidian notes
    if args.notes or args.all:
        logger.info("=" * 50)
        logger.info("Step 6: Generating Obsidian notes...")
        logger.info("=" * 50)
        papers = load_filtered_papers()

        notes = generate_all_notes(config, papers, parse_pdfs=True)
        logger.info(f"Generated {len(notes)} Obsidian notes")

    logger.info("=" * 50)
    logger.info("Workflow complete!")
    logger.info("=" * 50)


def main():
    """Main entry point"""
    # Load configuration
    config_path = os.path.join(project_root, 'config.yaml')
    config = load_config(config_path)

    # Setup logging
    logger = setup_logging(config)

    # Parse arguments
    parser = argparse.ArgumentParser(
        description='AI Paper Recommendation Workflow'
    )
    parser.add_argument('--all', action='store_true',
                       help='Run all workflow steps')
    parser.add_argument('--fetch-arxiv', action='store_true',
                       help='Fetch arXiv papers')
    parser.add_argument('--fetch-modelscope', action='store_true',
                       help='Fetch ModelScope/魔搭社区 Papers')
    parser.add_argument('--filter', action='store_true',
                       help='Filter papers by relevance')
    parser.add_argument('--download', action='store_true',
                       help='Download PDFs and extract images')
    parser.add_argument('--parse', action='store_true',
                       help='Parse PDF content with Chinese summary')
    parser.add_argument('--notes', action='store_true',
                       help='Generate Obsidian notes')

    args = parser.parse_args()

    # If no arguments, run full workflow
    if not any([args.all, args.fetch_arxiv, args.fetch_modelscope, args.filter,
                args.download, args.parse, args.notes]):
        args.all = True

    # Print welcome message
    max_papers = config.get('max_papers_per_day', 5)

    logger.info("=" * 60)
    logger.info("AI Paper Recommendation Workflow")
    logger.info("=" * 60)
    logger.info(f"Vault: {config.get('obsidian', {}).get('vault_path', 'Not configured')}")
    logger.info(f"Keywords: {', '.join(config.get('keywords', [])[:5])}...")
    logger.info(f"Max papers per day: {max_papers}")
    logger.info("=" * 60)

    # Run workflow
    try:
        run_workflow(config, args)
    except Exception as e:
        logger.error(f"Workflow error: {e}", exc_info=True)
        sys.exit(1)


if __name__ == '__main__':
    main()