Help with the minimal replicated code of async extracting nodes metadata.

[Whisht]

Say I have a basic aiohttp method for extracting the metadata of a given node.

class Fetcher:
	def __init__(self):
		self._session = None
		self.url = None
	async def fetch(self, data):
		async with self._session.get(self.url,data=data) as resp:
			return await resp.json()

And other helper functions, such as run_async_tasks:

def run_async_tasks(tasks):
    """Run a list of async tasks."""
    tasks_to_execute = tasks
    callback_manager = CallbackManager([])

    async def _gather():
        return await asyncio.gather(*tasks_to_execute)

    with callback_manager.as_trace("index_construction"):
        outputs = asyncio.run(_gather())
        return outputs

Now, I want to write a CustomMetadataExtractor to extract metadata for nodes.

class CustomMetadataExtractor():
	async def aextract(self, nodes):
			tasks = [self.fetcher.fetch(node) for node in nodes] # I have no idea about how to create a shared fetcher or ClientSession to manage the parallel requests in one single connection pool.
			return await asyncio.gather(*tasks)

And, the async aextract should be running normally in a NodeParser by a synchronous funtionget_nodes_from_documents. However, I have no idea about how to create a shared fetcher or ClientSession to manage the parallel requests in one single connection pool. It always has the Unclosed client session error and multiple sessions.

I want to maintain a single persist ClientSession for the metadata extract process in NodeParser, and it should be destroyed automatically. How could it be implemented? Thanks.

There is a similar process the async index building. But I didn't see any manual close operation for the async embed_model.embedding in VectorStoreIndex. How do you guys manage this async context?

[dosubot[bot]]

🤖

To manage a shared ClientSession across multiple requests, you can initialize the ClientSession in the __init__ method of your Fetcher class and close it in a close method. This way, the same ClientSession will be used for all requests made by the Fetcher instance. Here's how you can modify your Fetcher class:

class Fetcher:
    def __init__(self):
        self._session = aiohttp.ClientSession()

    async def fetch(self, data):
        async with self._session.get(self.url, data=data) as resp:
            return await resp.json()

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

Then, in your CustomMetadataExtractor class, you can create a Fetcher instance and use it for all requests. After all requests are done, you can close the Fetcher instance to close the ClientSession. Here's how you can modify your CustomMetadataExtractor class:

class CustomMetadataExtractor():
    def __init__(self):
        self.fetcher = Fetcher()

    async def aextract(self, nodes):
        tasks = [self.fetcher.fetch(node) for node in nodes]
        results = await asyncio.gather(*tasks)
        await self.fetcher.close()
        return results

This way, the same ClientSession (and thus the same connection pool) will be used for all requests made in the aextract method, and the ClientSession will be closed automatically after all requests are done.

Please note that you should handle exceptions properly to ensure the ClientSession is closed even if an error occurs. You can use a try/finally block for this:

class CustomMetadataExtractor():
    def __init__(self):
        self.fetcher = Fetcher()

    async def aextract(self, nodes):
        try:
            tasks = [self.fetcher.fetch(node) for node in nodes]
            results = await asyncio.gather(*tasks)
        finally:
            await self.fetcher.close()
        return results

This should solve the Unclosed client session error and the issue with multiple sessions.

Sources

• llama_index/callbacks/base.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[Whisht]

I wrote a minimal custom Async MetadataExtractor with llama_index 0.8.59. Hope this could be helpful for you if you have the same problem.