import asyncio

import cognee


async def main():
    await cognee.forget(everything=True)
    text = """
    In 1998 the project launched. In 2001 version 1.0 shipped. In 2004 the team merged
    with another group. In 2010 support for v1 ended.
    """

    await cognee.remember(
        text,
        dataset_name="timeline_demo",
        temporal_cognify=True,
        self_improvement=False,
    )

    from cognee import SearchType

    # Before / after queries
    result = await cognee.recall(
        query_type=SearchType.TEMPORAL, query_text="What happened before 2000?", top_k=15
    )

    assert result != []

    result = await cognee.recall(
        query_type=SearchType.TEMPORAL, query_text="What happened after 2010?", top_k=15
    )

    assert result != []

    # Between queries
    result = await cognee.recall(
        query_type=SearchType.TEMPORAL, query_text="Events between 2001 and 2004", top_k=15
    )

    assert result != []

    # Scoped descriptions
    result = await cognee.recall(
        query_type=SearchType.TEMPORAL,
        query_text="Key project milestones between 1998 and 2010",
        top_k=15,
    )

    assert result != []

    result = await cognee.recall(
        query_type=SearchType.TEMPORAL,
        query_text="What happened after 2004?",
        datasets=["timeline_demo"],
        top_k=15,
    )

    assert result != []


if __name__ == "__main__":
    asyncio.run(main())