omi/backend/routers/wrapped.py at main · ghzhost/omi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
"""
Wrapped 2025 API endpoints.

Provides generation and retrieval of yearly recap data.
"""

from typing import Any, Dict, Optional

from utils.executors import llm_executor

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel

import database.wrapped as wrapped_db
from database.wrapped import WrappedStatus
from utils.other import endpoints as auth
from utils.wrapped.generate_2025 import generate_wrapped_2025
import logging

logger = logging.getLogger(__name__)

router = APIRouter()


# Response models
class WrappedStatusResponse(BaseModel):
    status: str
    year: int = 2025
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    progress: Optional[Dict[str, Any]] = None


class GenerateWrappedResponse(BaseModel):
    status: str
    message: str


def _run_wrapped_generation(uid: str, year: int):
    """Run wrapped generation in background executor."""
    try:
        generate_wrapped_2025(uid, year)
    except Exception as e:
        logger.error(f"Error in wrapped generation for user {uid}: {e}")
        wrapped_db.update_wrapped_status(uid, year, WrappedStatus.ERROR, error=str(e))


@router.get('/v1/wrapped/{year}', response_model=WrappedStatusResponse, tags=['wrapped'])
def get_wrapped_status(year: int, uid: str = Depends(auth.get_current_user_uid)):
    """
    Get the status and result of wrapped generation for a given year.

    Returns:
        - status: not_generated, processing, done, or error
        - result: The wrapped payload (only when status=done)
        - error: Error message (only when status=error)
        - progress: Progress info (only when status=processing)
    """
    # For now, only support 2025
    if year != 2025:
        raise HTTPException(status_code=400, detail="Only year 2025 is currently supported")

    wrapped = wrapped_db.get_wrapped(uid, year)

    if not wrapped:
        return WrappedStatusResponse(
            status=WrappedStatus.NOT_GENERATED,
            year=year,
        )

    return WrappedStatusResponse(
        status=wrapped.get('status', WrappedStatus.NOT_GENERATED),
        year=year,
        result=wrapped.get('result'),
        error=wrapped.get('error'),
        progress=wrapped.get('progress'),
    )


@router.post('/v1/wrapped/{year}/generate', response_model=GenerateWrappedResponse, tags=['wrapped'])
def generate_wrapped(
    year: int, uid: str = Depends(auth.with_rate_limit(auth.get_current_user_uid, "wrapped:generate"))
):
    """
    Start wrapped generation for a given year.

    This is idempotent:
    - If already done: returns done status (no regeneration in v1)
    - If already processing: returns processing status
    - If error or not generated: starts generation
    - If processing but stuck (no heartbeat for 15 min): restarts generation
    """
    # For now, only support 2025
    if year != 2025:
        raise HTTPException(status_code=400, detail="Only year 2025 is currently supported")

    wrapped = wrapped_db.get_wrapped(uid, year)

    # Already done - no regeneration in v1
    if wrapped and wrapped.get('status') == WrappedStatus.DONE:
        return GenerateWrappedResponse(
            status=WrappedStatus.DONE,
            message="Your Wrapped 2025 is already generated",
        )

    # Already processing - check if stuck
    if wrapped and wrapped.get('status') == WrappedStatus.PROCESSING:
        if wrapped_db.is_wrapped_stuck(wrapped):
            # Restart stuck job
            wrapped_db.reset_wrapped_for_regeneration(uid, year)
            llm_executor.submit(_run_wrapped_generation, uid, year)
            return GenerateWrappedResponse(
                status=WrappedStatus.PROCESSING,
                message="Restarting stuck generation...",
            )
        else:
            return GenerateWrappedResponse(
                status=WrappedStatus.PROCESSING,
                message="Generation is already in progress",
            )

    # Error or not generated - start fresh
    if wrapped and wrapped.get('status') == WrappedStatus.ERROR:
        wrapped_db.reset_wrapped_for_regeneration(uid, year)
    else:
        wrapped_db.create_wrapped(uid, year)

    # Start generation in background
    llm_executor.submit(_run_wrapped_generation, uid, year)

    return GenerateWrappedResponse(
        status=WrappedStatus.PROCESSING,
        message="Starting Wrapped 2025 generation...",
    )