import asyncio

import cognee
from cognee.api.v1.search import SearchType
from cognee.modules.pipelines import Task, run_tasks
from cognee.tasks.temporal_awareness import (
    build_graph_with_temporal_awareness,
    search_graph_with_temporal_awareness,
)

text_list = [
    "Kamala Harris is the Attorney General of California. She was previously "
    "the district attorney for San Francisco.",
    "As AG, Harris was in office from January 3, 2011 – January 3, 2017",
]


async def main():
    tasks = [
        Task(build_graph_with_temporal_awareness, text_list=text_list),
        Task(
            search_graph_with_temporal_awareness, query="Who was the California Attorney General?"
        ),
    ]

    pipeline = run_tasks(tasks)

    async for result in pipeline:
        print(result)


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