diff --git a/src/dreamer/consolidate.py b/src/dreamer/consolidate.py index d834df4b..bb3c047c 100644 --- a/src/dreamer/consolidate.py +++ b/src/dreamer/consolidate.py @@ -15,6 +15,7 @@ from src.utils.queue_payload import DreamPayload from src.utils.representation import ( ExplicitObservation, Representation, + flatten_message_ids, ) logger = logging.getLogger(__name__) @@ -202,7 +203,7 @@ async def _consolidate_cluster( # NOTE: other kinds of observations here in the future metadata = schemas.DocumentMetadata( - message_ids=obs.message_ids, + message_ids=flatten_message_ids(obs.message_ids), message_created_at=format_datetime_utc(obs.created_at), premises=premises, ) diff --git a/src/utils/representation.py b/src/utils/representation.py index 7a0644d0..74557496 100644 --- a/src/utils/representation.py +++ b/src/utils/representation.py @@ -8,6 +8,40 @@ from src import models from src.utils.formatting import parse_datetime_iso +def flatten_message_ids( + message_ids: list[int] | list[list[int]] | list[tuple[int, int]], +) -> list[int]: + """ + Flatten message_ids that may be in old tuple format or nested list format. + + This handles backwards compatibility with the old schema where message_ids + was list[tuple[int, int]] representing ranges, and the new schema where + it's list[int] representing individual message IDs. + + Args: + message_ids: Either a flat list of ints, nested list, or list of tuples + + Returns: + A flat list of unique message IDs, sorted + + Examples: + [1, 2, 3] -> [1, 2, 3] + [[1, 2], [3, 4]] -> [1, 2, 3, 4] + [(105, 105)] -> [105] + [[105, 105]] -> [105] + """ + result: list[int] = [] + for item in message_ids: + if isinstance(item, (list | tuple)): + # Nested list or tuple - flatten it + result.extend(item) + else: + # Already flat + result.append(item) + # Remove duplicates and sort + return sorted(set(result)) + + class ObservationMetadata(BaseModel): created_at: datetime message_ids: list[int] @@ -267,7 +301,9 @@ class Representation(BaseModel): doc.internal_metadata, doc.created_at ), content=doc.content, - message_ids=doc.internal_metadata.get("message_ids", []), + message_ids=flatten_message_ids( + doc.internal_metadata.get("message_ids", []) + ), session_name=doc.session_name, ) for doc in documents @@ -279,7 +315,9 @@ class Representation(BaseModel): doc.internal_metadata, doc.created_at ), conclusion=doc.content, - message_ids=doc.internal_metadata.get("message_ids", []), + message_ids=flatten_message_ids( + doc.internal_metadata.get("message_ids", []) + ), session_name=doc.session_name, premises=doc.internal_metadata.get("premises", []), )