Pipeline Data Flow

from dataclasses import dataclass, field
from datetime import datetime
from typing import Dict, List, Any, Optional

@dataclass
class PipelineData:
    """In-memory state passed between pipeline steps. Not persisted to database."""
    
    # ===== INPUT DATA (from API request) =====
    task_id: str
    """Celery task ID - used for correlation in Logfire"""
    
    user_id: str
    """User ID from JWT token - for database writes"""
    
    email_template: str
    """Template string with placeholders like {{name}}, {{research}}"""
    
    recipient_name: str
    """Full name of professor/recipient (e.g., 'Dr. Jane Smith')"""
    
    recipient_interest: str
    """Research area or interest (e.g., 'machine learning')"""
    
    # ===== STEP 1 OUTPUTS (TemplateParser) =====
    search_terms: List[str] = field(default_factory=list)
    """
    Search queries extracted from template and recipient info.
    Example: ["Dr. Jane Smith machine learning", "Jane Smith publications"]
    """
    
    template_type: TemplateType | None = None
    """Required after Step 1 - set by TemplateParser (RESEARCH, BOOK, GENERAL)"""
    
    template_analysis: Dict[str, Any] = field(default_factory=dict)
    """
    Required after Step 1 - parsing details (placeholders, tone, etc.).
    Always present after TemplateParser.
    """
    
    # ===== STEP 2 OUTPUTS (WebScraper) =====
    scraped_content: str = ""
    """
    Cleaned and summarized content from web scraping.
    Limited to ~3000 chars to avoid LLM context limits.
    """
    
    scraped_urls: List[str] = field(default_factory=list)
    """URLs that were successfully scraped"""
    
    scraped_page_contents: Dict[str, str] = field(default_factory=dict)
    """Mapping of URL -> raw cleaned content (per page)."""
    
    scraping_metadata: Dict[str, Any] = field(default_factory=dict)
    """
    Scraping stats: total_urls_tried, successful_scrapes, failed_urls
    """
    
    # ===== STEP 3 OUTPUTS (ArxivEnricher) =====
    arxiv_papers: List[Dict[str, Any]] = field(default_factory=list)
    """
    Papers fetched from ArXiv API (only if template_type == RESEARCH).
    Each dict has: {title, abstract, year, url, authors}
    Limited to top 5 most relevant papers.
    """
    
    enrichment_metadata: Dict[str, Any] = field(default_factory=dict)
    """
    Enrichment stats: papers_found, search metadata
    """
    
    # ===== STEP 4 OUTPUTS (EmailComposer) =====
    final_email: str = ""
    """
    Final composed email (ready to send).
    Set by EmailComposer step.
    """
    
    composition_metadata: Dict[str, Any] = field(default_factory=dict)
    """
    Composition stats: llm_tokens_used, validation_attempts, quality_score
    """
    
    is_confident: bool = False
    """
    Whether the email composer had sufficient context to write a
    quality personalized email (vs generic fallback).
    """
    
    # ===== METADATA (for final DB write) =====
    metadata: Dict[str, Any] = field(default_factory=dict)
    """
    Metadata that will be stored in emails.metadata JSONB column.
    EmailComposer populates this before DB write.
    """
    
    # ===== TRANSIENT DATA (logged to Logfire, not persisted) =====
    started_at: datetime = field(default_factory=datetime.utcnow)
    """Pipeline start time"""
    
    step_timings: Dict[str, float] = field(default_factory=dict)
    """
    Duration of each step in seconds.
    Example: {"template_parser": 1.2, "web_scraper": 3.5, ...}
    """
    
    errors: List[str] = field(default_factory=list)
    """
    Non-fatal errors encountered during execution.
    Fatal errors raise exceptions and terminate pipeline.
    """
    
    # ===== HELPER METHODS =====
    def total_duration(self) -> float:
        """Calculate total pipeline execution time in seconds"""
        return (datetime.utcnow() - self.started_at).total_seconds()
    
    def add_timing(self, step_name: str, duration: float) -> None:
        """Record step timing"""
        self.step_timings[step_name] = duration
    
    def add_error(self, step_name: str, error_message: str) -> None:
        """Record non-fatal error"""
        self.errors.append(f"{step_name}: {error_message}")