class AsyncNeo4jQuery:
    def __init__(self, neo4j):
        self.neo4j = neo4j

    async def close(self):
        await self.neo4j.close()

    async def get_document_by_id(self, doc_id: str):
        query = """
        MATCH (d:Document {doc_id: $doc_id})
        OPTIONAL MATCH (d)-[:HAS_CHUNK]->(c:Chunk)
        RETURN d, collect(c) as chunks
        """
        async with self.neo4j.session() as session:
            result = await session.run(query, doc_id=doc_id)
            return [
                record.data() for record in await result.consume().records
            ]  # 注意 result 需要 async 迭代

    async def search_chunks_by_topic(self, topic: str):
        query = """
        MATCH (c:Chunk {topic: $topic})
        OPTIONAL MATCH (c)-[:HAS_ENTITY]->(e:Entity)
        RETURN c, collect(e.name) as entities
        """
        async with self.neo4j.session() as session:
            result = await session.run(query, topic=topic)
            return [record.data() async for record in result]