Return the metadata from a Query Engine in an Agent

[AlexanderKolev]

Hello kind people,

I have an OpenAI Agent which uses an index as a query engine, as a tool. I want this tool to always return the metadata as a source, but I am struggling to do so. I have tried with prompting and using CitationQueryEngine but without success. Do you have any ideas for doing it 'programmatically' in some way? If I was not using an Agent but just using the query engine it seems more straightforward.
My code for the query engine looks like this:

def _configure_query_engine(
    index: BaseIndex,
    embed_model: OpenAIEmbedding
) -> BaseQueryEngine:
    query_engine = index.as_query_engine(
        embed_model=embed_model,
        similarity_top_k=5,
        include_metadata=True
    )
    query_engine.update_prompts({
        TEXT_QA_TEMPLATE: create_synth_prompt_template()
    })
    return query_engine

Thank you in advance.

[dosubot[bot]]

To ensure that the metadata is always returned as a source when using an OpenAI Agent with a query engine, you need to set the include_metadata parameter to True in the index.as_query_engine method. This is already done in your provided code. Here is the relevant part of your code:

def _configure_query_engine(
    index: BaseIndex,
    embed_model: OpenAIEmbedding
) -> BaseQueryEngine:
    query_engine = index.as_query_engine(
        embed_model=embed_model,
        similarity_top_k=5,
        include_metadata=True  # This ensures metadata is included
    )
    query_engine.update_prompts({
        TEXT_QA_TEMPLATE: create_synth_prompt_template()
    })
    return query_engine

The include_metadata=True parameter ensures that the metadata is included in the query results. Therefore, your current implementation already ensures that metadata is always returned as a source.

Additionally, you can refer to the CitationQueryEngine class in LlamaIndex, which provides a way to ensure that sources are cited in the responses. Here is a relevant snippet from the documentation:

class CitationQueryEngine(BaseQueryEngine):
    def __init__(
        self,
        retriever: BaseRetriever,
        llm: Optional[LLM] = None,
        response_synthesizer: Optional[BaseSynthesizer] = None,
        citation_chunk_size: int = DEFAULT_CITATION_CHUNK_SIZE,
        citation_chunk_overlap: int = DEFAULT_CITATION_CHUNK_OVERLAP,
        text_splitter: Optional[TextSplitter] = None,
        node_postprocessors: Optional[List[BaseNodePostprocessor]] = None,
        callback_manager: Optional[CallbackManager] = None,
        metadata_mode: MetadataMode = MetadataMode.NONE,
    ) -> None:
        self.text_splitter = text_splitter or SentenceSplitter(
            chunk_size=citation_chunk_size, chunk_overlap=citation_chunk_overlap
        )
        self._retriever = retriever
        self._response_synthesizer = response_synthesizer or get_response_synthesizer(
            llm=llm,
            callback_manager=callback_manager,
        )
        self._node_postprocessors = node_postprocessors or []
        self._metadata_mode = metadata_mode
        super().__init__(callback_manager=callback_manager)

    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:
                text = f"Source {len(new_nodes) + 1}:\n{text_chunk}\n"
                new_node = NodeWithScore(
                    node=TextNode.model_validate(node.node), score=node.score
                )
                new_node.node.text = text
                new_nodes.append(new_node)
        return new_nodes

    def retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        nodes = self._retriever.retrieve(query_bundle)
        for postprocessor in self._node_postprocessors:
            nodes = postprocessor.postprocess_nodes(nodes, query_bundle=query_bundle)
        return nodes

    def synthesize(
        self,
        query_bundle: QueryBundle,
        nodes: List[NodeWithScore],
        additional_source_nodes: Optional[Sequence[NodeWithScore]] = None,
    ) -> RESPONSE_TYPE:
        nodes = self._create_citation_nodes(nodes)
        return self._response_synthesizer.synthesize(
            query=query_bundle,
            nodes=nodes,
            additional_source_nodes=additional_source_nodes,
        )

This example shows how to create citation nodes that include metadata, ensuring that sources are cited in the responses. You can adapt similar techniques to your query engine configuration to ensure metadata is always included and properly cited ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[AlexanderKolev]

But how can I ensure my agent always returns the metadata as well?

[dosubot[bot]]

To ensure your agent always returns the metadata, you need to include the metadata in the response construction. Here is an example of how you can modify the build_agent_per_doc function to ensure metadata is included in the response:

async def build_agent_per_doc(nodes, file_base):
    print(file_base)

    vi_out_path = f"./data/llamaindex_docs/{file_base}"
    summary_out_path = f"./data/llamaindex_docs/{file_base}_summary.pkl"
    if not os.path.exists(vi_out_path):
        Path("./data/llamaindex_docs/").mkdir(parents=True, exist_ok=True)
        # build vector index
        vector_index = VectorStoreIndex(nodes)
        vector_index.storage_context.persist(persist_dir=vi_out_path)
    else:
        vector_index = load_index_from_storage(
            StorageContext.from_defaults(persist_dir=vi_out_path),
        )

    # build summary index
    summary_index = SummaryIndex(nodes)

    # define query engines
    vector_query_engine = vector_index.as_query_engine(llm=llm)
    summary_query_engine = summary_index.as_query_engine(
        response_mode="tree_summarize", llm=llm
    )
    # extract a summary
    if not os.path.exists(summary_out_path):
        Path(summary_out_path).parent.mkdir(parents=True, exist_ok=True)
        summary = str(
            await summary_query_engine.aquery(
                "Extract a concise 1-2 line summary of this document"
            )
        )
        pickle.dump(summary, open(summary_out_path, "wb"))
    else:
        summary = pickle.load(open(summary_out_path, "rb"))

    # define tools
    query_engine_tools = [
        QueryEngineTool(
            query_engine=vector_query_engine,
            metadata=ToolMetadata(
                name=f"vector_tool_{file_base}",
                description=f"Useful for questions related to specific facts",
            ),
        ),
        QueryEngineTool(
            query_engine=summary_query_engine,
            metadata=ToolMetadata(
                name=f"summary_tool_{file_base}",
                description=f"Useful for summarization questions",
            ),
        ),
    ]

    # build agent
    function_llm = OpenAI(model="gpt-4")
    agent = OpenAIAgent.from_tools(
        query_engine_tools,
        llm=function_llm,
        verbose=True,
        system_prompt=f"""\
You are a specialized agent designed to answer queries about the `{file_base}.html` part of the LlamaIndex docs.
You must ALWAYS use at least one of the tools provided when answering a question; do NOT rely on prior knowledge.\
""",
    )

    # Ensure metadata is included in the response
    agent.include_metadata = True

    return agent, summary

In this example, the include_metadata attribute is set to True for the agent, ensuring that metadata is included in the responses. You may need to adjust the response generation logic within the agent to ensure that metadata is properly formatted and included in the final output ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[AlexanderKolev]

Can you instead give a Custom Query Engine which returns the metadata as part of the response? @dosu

[dosubot[bot]]

Yes, you can create a Custom Query Engine that returns the metadata as part of the agent response. Below is an example of how you can modify the RAGQueryEngine to include metadata in the response:

from llama_index.core.query_engine import CustomQueryEngine
from llama_index.core.retrievers import BaseRetriever
from llama_index.core import get_response_synthesizer
from llama_index.core.response_synthesizers import BaseSynthesizer

class RAGQueryEngineWithMetadata(CustomQueryEngine):
    """RAG Query Engine with Metadata."""

    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)
        response_obj = self.response_synthesizer.synthesize(query_str, nodes)
        
        # Collect metadata
        metadata = [node.node.get_metadata() for node in nodes]
        
        # Include metadata in the response
        response_with_metadata = {
            "response": str(response_obj),
            "metadata": metadata
        }
        
        return response_with_metadata

# Example usage
synthesizer = get_response_synthesizer(response_mode="compact")
query_engine = RAGQueryEngineWithMetadata(
    retriever=retriever, response_synthesizer=synthesizer
)

response = query_engine.query("What did the author do growing up?")
print(response)

In this example, the custom_query method retrieves the nodes, synthesizes the response, and then collects the metadata from each node. The final response includes both the synthesized response and the metadata ^[1][2].

To continue talking to Dosu, mention @dosu.

[timelesshc]

@AlexanderKolev Hi, I'm trying to get the retrieved context from agent as well. Have you figured it out? I'm using ReAct agent ( https://docs.llamaindex.ai/en/stable/examples/agent/react_agent_with_query_engine/ ) and have not been able to get the retrieved context from the query engines.