fleet-memory/hindsight-api/tests/test_fact_ordering.py

"""
Test that facts from the same conversation maintain temporal ordering.

This ensures that when multiple facts are extracted from a long conversation,
their relative order is preserved via time offsets, allowing retrieval to
distinguish between things said earlier vs later.
"""
import pytest
from datetime import datetime, timezone
from hindsight_api import MemoryEngine, RequestContext
from hindsight_api.engine.memory_engine import Budget
import os


@pytest.mark.asyncio
async def test_fact_ordering_within_conversation(memory, request_context):
    bank_id = "test_ordering_agent"

    # Get/create agent (auto-creates with defaults)
    await memory.get_bank_profile(bank_id, request_context=request_context)

    # Update disposition to match Marcus
    await memory.update_bank_disposition(bank_id, {
        "skepticism": 3,
        "literalism": 3,
        "empathy": 3
    }, request_context=request_context)

    # A conversation where Marcus changes his position
    conversation = """
Marcus: I think the Rams will win 27-24. Their defense is really strong.
Jamie: I disagree, I think Niners will win.
Marcus: Actually, after thinking about it more, I'm changing my prediction to Rams by 3 points only.
Jamie: That's more reasonable.
Marcus: Yeah, I realized I was being too optimistic about their defense.
"""

    base_event_date = datetime(2024, 11, 14, 10, 0, 0, tzinfo=timezone.utc)

    # Store the conversation
    await memory.retain_async(
        bank_id=bank_id,
        content=conversation,
        context="podcast discussion about NFL game",
        event_date=base_event_date,
        document_id="test_conv_1",
        request_context=request_context,
    )

    # Search for all facts about Marcus's predictions
    results = await memory.recall_async(
        bank_id=bank_id,
        query="Marcus prediction Rams",
        fact_type=['experience', 'world'],
        budget=Budget.LOW,
        max_tokens=8192,
        request_context=request_context,
    )

    print(f"\n=== Retrieved {len(results.results)} facts ===")
    for i, result in enumerate(results.results):
        print(f"{i+1}. [{result.mentioned_at}] {result.text[:100]}")

    # Get all facts (Marcus's predictions/statements)
    agent_facts = results.results

    print(f"\n=== Agent facts (Marcus's statements) ===")
    for i, fact in enumerate(agent_facts):
        print(f"{i+1}. [{fact.mentioned_at}] {fact.text}")

    # Check that agent facts have different timestamps
    if len(agent_facts) >= 2:
        # Parse timestamps
        timestamps = [datetime.fromisoformat(f.mentioned_at.replace('Z', '+00:00')) for f in agent_facts]

        # Verify timestamps are different (have time offsets)
        unique_timestamps = set(timestamps)
        assert len(unique_timestamps) == len(timestamps), \
            f"Expected unique timestamps for each fact, but got duplicates: {timestamps}"

        # Sort facts by timestamp for ordering check
        # Note: recall returns by relevance, not time order
        sorted_facts = sorted(agent_facts, key=lambda f: datetime.fromisoformat(f.mentioned_at.replace('Z', '+00:00')))
        sorted_timestamps = [datetime.fromisoformat(f.mentioned_at.replace('Z', '+00:00')) for f in sorted_facts]

        # Verify sorted timestamps are in ascending order
        for i in range(len(sorted_timestamps) - 1):
            assert sorted_timestamps[i] < sorted_timestamps[i + 1], \
                f"Facts should have sequential timestamps. Fact {i} ({sorted_timestamps[i]}) >= Fact {i+1} ({sorted_timestamps[i+1]})"

        # Verify reasonable time spacing (should be ~10 seconds apart)
        time_diffs = [(sorted_timestamps[i+1] - sorted_timestamps[i]).total_seconds() for i in range(len(sorted_timestamps) - 1)]
        print(f"\n=== Time differences between facts: {time_diffs} seconds ===")

        # Each fact should be 10+ seconds apart (allowing for some flexibility)
        for diff in time_diffs:
            assert diff >= 5, f"Expected at least 5 seconds between facts, got {diff}"

        # Update agent_facts to be sorted for subsequent checks
        agent_facts = sorted_facts
        timestamps = sorted_timestamps

        print(f"\n✅ All {len(agent_facts)} agent facts have properly ordered timestamps")

    # Verify that facts capture the key information
    # Note: LLM may merge related predictions into single facts
    agent_texts = [f.text.lower() for f in agent_facts]
    all_text = " ".join(agent_texts)

    # Look for evidence of the predictions being captured (may be merged or separate)
    has_prediction_info = '27' in all_text or 'rams' in all_text or 'prediction' in all_text

    assert has_prediction_info, "Facts should contain information about Marcus's predictions"
    print(f"\n✅ Facts capture prediction information")

    # Cleanup
    await memory.delete_bank(bank_id, request_context=request_context)

    print(f"\n✅ Test passed: Fact ordering within conversation is preserved")


@pytest.mark.asyncio
async def test_multiple_documents_ordering(memory, request_context):

    bank_id = "test_multi_doc_agent"

    await memory.get_bank_profile(bank_id, request_context=request_context)  # Auto-creates with defaults

    # Two separate conversations with same base time
    base_time = datetime(2024, 11, 14, 10, 0, 0, tzinfo=timezone.utc)

    conv1 = """
Alice: I prefer React for this project.
Bob: Why React?
Alice: It has better tooling and I'm more familiar with it.
"""

    conv2 = """
Alice: Actually, I'm thinking Vue might be better.
Bob: What changed your mind?
Alice: I reconsidered the team's experience level.
"""

    # Store both conversations with batch
    await memory.retain_batch_async(
        bank_id=bank_id,
        contents=[
            {"content": conv1, "context": "project discussion 1", "event_date": base_time},
            {"content": conv2, "context": "project discussion 2", "event_date": base_time}
        ],
        request_context=request_context,
    )

    # Search for Alice's preferences
    results = await memory.recall_async(
        bank_id=bank_id,
        query="Alice preference React Vue",
        fact_type=['experience', 'world'],
        budget=Budget.LOW,
        max_tokens=8192,
        request_context=request_context,
    )

    print(f"\n=== Retrieved {len(results.results)} agent facts ===")
    agent_facts = results.results

    for i, fact in enumerate(agent_facts):
        print(f"{i+1}. [{fact.mentioned_at}] {fact.text[:80]}")

    # Each conversation's facts should have different timestamps
    if len(agent_facts) >= 2:
        timestamps = [datetime.fromisoformat(f.mentioned_at.replace('Z', '+00:00')) for f in agent_facts]
        unique_timestamps = set(timestamps)

        assert len(unique_timestamps) >= 2, \
            f"Expected multiple unique timestamps across conversations, got: {len(unique_timestamps)}"

        print(f"\n✅ Facts from {len(agent_facts)} statements have {len(unique_timestamps)} unique timestamps")

    # Cleanup
    await memory.delete_bank(bank_id, request_context=request_context)

    print(f"\n✅ Test passed: Multiple documents maintain separate ordering")