backend/utils.py

from __future__ import annotations

from typing import Iterable, List, Optional, Tuple

from backend.schemas import ChatMessagePayload

ChatPairs = List[Tuple[Optional[str], Optional[str]]]


def messages_to_pairs(messages: Iterable[ChatMessagePayload]) -> ChatPairs:
    pairs: ChatPairs = []
    pending_user: Optional[str] = None
    for msg in messages:
        if msg.role == "user":
            if pending_user is not None:
                pairs.append((pending_user, None))
            pending_user = msg.content
        elif msg.role == "assistant":
            if pending_user is not None:
                pairs.append((pending_user, msg.content))
                pending_user = None
            else:
                pairs.append((None, msg.content))
    if pending_user is not None:
        pairs.append((pending_user, None))
    return pairs