diff --git a/src/typeagent/knowpro/answers.py b/src/typeagent/knowpro/answers.py index 93ecc84f..d0307929 100644 --- a/src/typeagent/knowpro/answers.py +++ b/src/typeagent/knowpro/answers.py @@ -431,15 +431,36 @@ async def get_enclosing_data_range_for_messages( messages: IMessageCollection, message_ordinals: Iterable[MessageOrdinal], ) -> DateRange | None: + """The time span of the given messages, for display in a prompt. + + It runs from the earliest to the latest message's timestamp; for a single + message the end is None. Like `get_time_range_for_conversation`, the end is + an inclusive-looking bound, so do NOT pass this range to a timestamp index + lookup: `DateRange` is half-open and the last message would fall outside it. + """ text_range = get_enclosing_text_range(message_ordinals) if not text_range: return None - return await get_enclosing_date_range_for_text_range(messages, text_range) + start_timestamp = ( + await messages.get_item(text_range.start.message_ordinal) + ).timestamp + if not start_timestamp: + return None + end_timestamp: str | None = None + if text_range.end: + # TextRange.end is exclusive; the last message is the one before it. + last_ordinal = text_range.end.message_ordinal - 1 + end_timestamp = (await messages.get_item(last_ordinal)).timestamp + return DateRange( + start=Datetime.fromisoformat(start_timestamp), + end=Datetime.fromisoformat(end_timestamp) if end_timestamp else None, + ) def get_enclosing_text_range( message_ordinals: Iterable[MessageOrdinal], ) -> TextRange | None: + """The smallest TextRange containing the given messages, or None if empty.""" start: MessageOrdinal | None = None end: MessageOrdinal | None = start for ordinal in message_ordinals: @@ -455,13 +476,17 @@ def get_enclosing_text_range( def text_range_from_message_range( start: MessageOrdinal, end: MessageOrdinal ) -> TextRange | None: + """Return the TextRange covering messages ``start`` through ``end``, inclusive. + + TextRange.end is exclusive, so it is set one past ``end``. + """ if start == end: # Point location return TextRange(start=TextLocation(start)) elif start < end: return TextRange( start=TextLocation(start), - end=TextLocation(end), + end=TextLocation(end + 1), ) else: raise ValueError(f"Expect message ordinal range: {start} <= {end}") diff --git a/tests/test_answers.py b/tests/test_answers.py index ebc47096..9fd8c2f3 100644 --- a/tests/test_answers.py +++ b/tests/test_answers.py @@ -14,12 +14,13 @@ AnswerGeneratorSettings, facets_to_merged_facets, generate_answers, + get_enclosing_data_range_for_messages, get_enclosing_date_range_for_text_range, get_enclosing_text_range, merged_facets_to_facets, text_range_from_message_range, ) -from typeagent.knowpro.interfaces import TextLocation, TextRange +from typeagent.knowpro.interfaces import Datetime, TextLocation, TextRange from typeagent.knowpro.knowledge_schema import Facet from typeagent.knowpro.search import ConversationSearchResult @@ -79,12 +80,12 @@ def test_roundtrip_through_merged(self) -> None: # --------------------------------------------------------------------------- -# Change 2: get_enclosing_date_range_for_text_range uses ordinal-1 for end +# get_enclosing_date_range_for_text_range treats TextRange.end as exclusive # --------------------------------------------------------------------------- class TestGetEnclosingDateRangeForTextRange: - """Verify the off-by-one fix: end is exclusive, so we subtract 1.""" + """TextRange.end is exclusive, so its message's timestamp is the exclusive end.""" @pytest_asyncio.fixture() async def messages(self) -> AsyncGenerator[FakeMessageCollection, None]: @@ -156,6 +157,54 @@ async def test_no_timestamp_returns_none(self) -> None: assert dr is None +class TestGetEnclosingDataRangeForMessages: + """The time range spans the first through the last source message.""" + + @pytest_asyncio.fixture() + async def messages(self) -> AsyncGenerator[FakeMessageCollection, None]: + """Five messages with ordinals 0..4; message i is at hour i.""" + coll = FakeMessageCollection() + for i in range(5): + await coll.append(FakeMessage("text", message_ordinal=i)) + yield coll + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("ordinals", "end_hour"), + [ + ([2], None), # single message: start only + ([0, 1], 1), + ([3, 1], 3), + ([0, 2, 3], 3), + ([1, 4], 4), # includes the conversation's last message + ], + ) + async def test_spans_first_through_last_source_message( + self, + messages: FakeMessageCollection, + ordinals: list[int], + end_hour: int | None, + ) -> None: + dr = await get_enclosing_data_range_for_messages(messages, ordinals) + assert dr is not None + assert dr.start.hour == min(ordinals) + assert (dr.end.hour if dr.end else None) == end_hour + for ordinal in ordinals: + timestamp = (await messages.get_item(ordinal)).timestamp + assert timestamp is not None + # The end is an inclusive display bound, not DateRange's exclusive one. + dt = Datetime.fromisoformat(timestamp) + assert dr.start <= dt + assert dr.end is None or dt <= dr.end + + @pytest.mark.asyncio + async def test_no_timestamp_returns_none(self) -> None: + coll = FakeMessageCollection() + await coll.append(FakeMessage("text")) # no message_ordinal -> no timestamp + await coll.append(FakeMessage("text")) + assert await get_enclosing_data_range_for_messages(coll, [0, 1]) is None + + # --------------------------------------------------------------------------- # Helper functions (also exercised for completeness) # --------------------------------------------------------------------------- @@ -173,7 +222,14 @@ def test_multiple_ordinals(self) -> None: assert tr is not None assert tr.start.message_ordinal == 1 assert tr.end is not None - assert tr.end.message_ordinal == 7 + assert tr.end.message_ordinal == 8 # exclusive: one past the last message + + @pytest.mark.parametrize("ordinals", [[0, 1], [3, 1, 7], [4, 4, 2]]) + def test_contains_every_source_message(self, ordinals: list[int]) -> None: + tr = get_enclosing_text_range(ordinals) + assert tr is not None + for ordinal in ordinals: + assert TextRange(start=TextLocation(ordinal)) in tr def test_empty_ordinals(self) -> None: tr = get_enclosing_text_range([]) @@ -192,7 +248,7 @@ def test_range(self) -> None: assert tr is not None assert tr.start.message_ordinal == 2 assert tr.end is not None - assert tr.end.message_ordinal == 5 + assert tr.end.message_ordinal == 6 # exclusive: one past message 5 def test_invalid_raises(self) -> None: with pytest.raises(ValueError, match="Expect message ordinal range"):