Avoid that query engine takes same chunk several times in sequence

[aiwalter]

I am using a CitationQueryEngine and this is called several times in sequence to complete a task. For each task I want to make sure that a given chunk is only used/cited once. My idea was to remove the chunk from the index as soon as it is used by the query engine to avoid it will be used again. Is there a smarter way to track which chunks have been used and avoid using again? How could I implement this? @dosubot

[dosubot[bot]]

Hey @aiwalter! Great to see you diving into another challenge 🚀

Based on your requirements, I would suggest implementing a caching mechanism or a hash map to keep track of the chunks that have already been used. This approach would allow for constant time checks to see if a chunk has been used before, significantly reducing the overhead associated with searching through and modifying a list of chunks each time a new chunk is processed.

In the context of the CitationQueryEngine, this could be implemented by adding an attribute to the class that stores the identifiers (or hashes) of the chunks that have been used. Before using a new chunk, the engine would check this cache to ensure the chunk hasn't been used before. This method is more efficient than removing used chunks from a list because it avoids the need to repeatedly search through and modify the list, which can become increasingly inefficient as the list grows.

Here's a conceptual example of how this could be implemented in the CitationQueryEngine class:

class CitationQueryEngine(BaseQueryEngine):
    def __init__(self, ...):
        ...
        self.used_chunks_cache = set()  # Initialize an empty set to store used chunk identifiers

    def _create_citation_nodes(self, nodes: List[NodeWithScore]) -> List[NodeWithScore]:
        new_nodes: List[NodeWithScore] = []
        for node in nodes:
            text_chunks = self.text_splitter.split_text(
                node.node.get_content(metadata_mode=self._metadata_mode)
            )

            for text_chunk in text_chunks:
                chunk_hash = hash(text_chunk)  # Create a hash of the chunk for identification
                if chunk_hash not in self.used_chunks_cache:  # Check if the chunk has been used
                    self.used_chunks_cache.add(chunk_hash)  # Mark the chunk as used
                    text = f"Source {len(new_nodes)+1}:\n{text_chunk}\n"

                    new_node = NodeWithScore(
                        node=TextNode.parse_obj(node.node), score=node.score
                    )
                    new_node.node.text = text
                    new_nodes.append(new_node)
        return new_nodes

This example uses a set to store the hashes of used chunks, providing constant time complexity for add and check operations. This approach is significantly more efficient than removing items from a list, especially as the number of processed chunks grows.

You can find more details about the CitationQueryEngine class in the LlamaIndex codebase.

Sources

• llama-index-core/llama_index/core/query_engine/citation_query_engine.py
• llama-index-legacy/llama_index/legacy/query_engine/citation_query_engine.py

About Dosu

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.

[aiwalter]

@dosubot, how would that look like if I want to remove any used chunk from the index? Additionally the CitationQueryEngine

[dosubot]

To directly remove used chunks from the index in the CitationQueryEngine, you'll need to implement a few key changes. Here's a structured approach:

1. Modify the Retriever: The retriever used with CitationQueryEngine should support a mechanism to exclude or remove nodes that have been processed. This could involve adjusting its internal storage or query logic to filter out nodes based on a list of identifiers for nodes that have already been used.

2. Track Processed Nodes: Enhance the CitationQueryEngine to maintain a set of identifiers for nodes that have been processed. This set will be used to determine which nodes should be excluded in subsequent retrievals.

3. Adjust Retrieval Methods: Update the retrieve and aretrieve methods in CitationQueryEngine to use the list of processed node identifiers to exclude these nodes from being retrieved again. This might require passing this list to the retriever and modifying the retriever's logic to filter accordingly.

Here's an example of how you might start implementing these changes:

class CitationQueryEngine(BaseQueryEngine):
    def __init__(self, ...):
        ...
        self.used_nodes = set()  # Set to track used nodes

    def retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        nodes = self._retriever.retrieve(query_bundle, exclude=self.used_nodes)

        for postprocessor in self._node_postprocessors:
            nodes = postprocessor.postprocess_nodes(nodes, query_bundle=query_bundle)

        # Mark nodes as used
        for node in nodes:
            self.used_nodes.add(node.node.id)

        return nodes

    async def aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        nodes = await self._retriever.aretrieve(query_bundle, exclude=self.used_nodes)

        for postprocessor in self._node_postprocessors:
            nodes = postprocessor.postprocess_nodes(nodes, query_bundle=query_bundle)

        # Mark nodes as used
        for node in nodes:
            self.used_nodes.add(node.node.id)

        return nodes

This assumes that nodes have a unique identifier (node.node.id) and that the retriever's retrieve and aretrieve methods can exclude nodes based on this identifier. You'll need to adapt this to fit your specific implementation, including how nodes are identified and how the retriever supports excluding nodes.

Currently, the BaseRetriever class does not inherently support excluding nodes based on user-defined criteria. It implements deduplication based on node uniqueness rather than exclusion. You may need to extend or modify the BaseRetriever class to add this functionality.

For more details, refer to the source code of the CitationQueryEngine and BaseRetriever classes in the LlamaIndex repository:

• CitationQueryEngine Source
• BaseRetriever Source

If you have further questions or need clarification, don't hesitate to ask.

Sources

• llama-index-core/llama_index/core/query_engine/citation_query_engine.py
• llama-index-legacy/llama_index/legacy/core/base_retriever.py

About Dosu

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.