Spaces:
Runtime error
Runtime error
| """Leaf query mechanism.""" | |
| import logging | |
| from typing import Any, Dict, Optional, cast | |
| from gpt_index.data_structs.data_structs import IndexGraph, Node | |
| from gpt_index.indices.query.base import BaseGPTIndexQuery | |
| from gpt_index.indices.query.schema import QueryBundle | |
| from gpt_index.indices.response.builder import ResponseBuilder | |
| from gpt_index.indices.utils import ( | |
| extract_numbers_given_response, | |
| get_sorted_node_list, | |
| truncate_text, | |
| ) | |
| from gpt_index.prompts.default_prompts import ( | |
| DEFAULT_QUERY_PROMPT, | |
| DEFAULT_QUERY_PROMPT_MULTIPLE, | |
| ) | |
| from gpt_index.prompts.prompts import TreeSelectMultiplePrompt, TreeSelectPrompt | |
| from gpt_index.response.schema import Response | |
| class GPTTreeIndexLeafQuery(BaseGPTIndexQuery[IndexGraph]): | |
| """GPT Tree Index leaf query. | |
| This class traverses the index graph and searches for a leaf node that can best | |
| answer the query. | |
| .. code-block:: python | |
| response = index.query("<query_str>", mode="default") | |
| Args: | |
| query_template (Optional[TreeSelectPrompt]): Tree Select Query Prompt | |
| (see :ref:`Prompt-Templates`). | |
| query_template_multiple (Optional[TreeSelectMultiplePrompt]): Tree Select | |
| Query Prompt (Multiple) | |
| (see :ref:`Prompt-Templates`). | |
| child_branch_factor (int): Number of child nodes to consider at each level. | |
| If child_branch_factor is 1, then the query will only choose one child node | |
| to traverse for any given parent node. | |
| If child_branch_factor is 2, then the query will choose two child nodes. | |
| """ | |
| def __init__( | |
| self, | |
| index_struct: IndexGraph, | |
| query_template: Optional[TreeSelectPrompt] = None, | |
| query_template_multiple: Optional[TreeSelectMultiplePrompt] = None, | |
| child_branch_factor: int = 1, | |
| **kwargs: Any, | |
| ) -> None: | |
| """Initialize params.""" | |
| super().__init__(index_struct, **kwargs) | |
| self.query_template = query_template or DEFAULT_QUERY_PROMPT | |
| self.query_template_multiple = ( | |
| query_template_multiple or DEFAULT_QUERY_PROMPT_MULTIPLE | |
| ) | |
| self.child_branch_factor = child_branch_factor | |
| def _query_with_selected_node( | |
| self, | |
| selected_node: Node, | |
| query_bundle: QueryBundle, | |
| prev_response: Optional[str] = None, | |
| level: int = 0, | |
| ) -> str: | |
| """Get response for selected node. | |
| If not leaf node, it will recursively call _query on the child nodes. | |
| If prev_response is provided, we will update prev_response with the answer. | |
| """ | |
| query_str = query_bundle.query_str | |
| if len(selected_node.child_indices) == 0: | |
| response_builder = ResponseBuilder( | |
| self._prompt_helper, | |
| self._llm_predictor, | |
| self.text_qa_template, | |
| self.refine_template, | |
| ) | |
| self.response_builder.add_node_as_source(selected_node) | |
| # use response builder to get answer from node | |
| node_text, sub_response = self._get_text_from_node( | |
| query_bundle, selected_node, level=level | |
| ) | |
| if sub_response is not None: | |
| # these are source nodes from within this node (when it's an index) | |
| for source_node in sub_response.source_nodes: | |
| self.response_builder.add_source_node(source_node) | |
| cur_response = response_builder.get_response_over_chunks( | |
| query_str, [node_text], prev_response=prev_response | |
| ) | |
| cur_response = cast(str, cur_response) | |
| logging.debug(f">[Level {level}] Current answer response: {cur_response} ") | |
| else: | |
| cur_response = self._query_level( | |
| { | |
| i: self.index_struct.all_nodes[i] | |
| for i in selected_node.child_indices | |
| }, | |
| query_bundle, | |
| level=level + 1, | |
| ) | |
| if prev_response is None: | |
| return cur_response | |
| else: | |
| context_msg = "\n".join([selected_node.get_text(), cur_response]) | |
| cur_response, formatted_refine_prompt = self._llm_predictor.predict( | |
| self.refine_template, | |
| query_str=query_str, | |
| existing_answer=prev_response, | |
| context_msg=context_msg, | |
| ) | |
| logging.debug(f">[Level {level}] Refine prompt: {formatted_refine_prompt}") | |
| logging.debug(f">[Level {level}] Current refined response: {cur_response} ") | |
| return cur_response | |
| def _query_level( | |
| self, | |
| cur_nodes: Dict[int, Node], | |
| query_bundle: QueryBundle, | |
| level: int = 0, | |
| ) -> str: | |
| """Answer a query recursively.""" | |
| query_str = query_bundle.query_str | |
| cur_node_list = get_sorted_node_list(cur_nodes) | |
| if len(cur_node_list) == 1: | |
| logging.debug(f">[Level {level}] Only one node left. Querying node.") | |
| return self._query_with_selected_node( | |
| cur_node_list[0], query_bundle, level=level | |
| ) | |
| elif self.child_branch_factor == 1: | |
| query_template = self.query_template.partial_format( | |
| num_chunks=len(cur_node_list), query_str=query_str | |
| ) | |
| numbered_node_text = self._prompt_helper.get_numbered_text_from_nodes( | |
| cur_node_list, prompt=query_template | |
| ) | |
| response, formatted_query_prompt = self._llm_predictor.predict( | |
| query_template, | |
| context_list=numbered_node_text, | |
| ) | |
| else: | |
| query_template_multiple = self.query_template_multiple.partial_format( | |
| num_chunks=len(cur_node_list), | |
| query_str=query_str, | |
| branching_factor=self.child_branch_factor, | |
| ) | |
| numbered_node_text = self._prompt_helper.get_numbered_text_from_nodes( | |
| cur_node_list, prompt=query_template_multiple | |
| ) | |
| response, formatted_query_prompt = self._llm_predictor.predict( | |
| query_template_multiple, | |
| context_list=numbered_node_text, | |
| ) | |
| logging.debug( | |
| f">[Level {level}] current prompt template: {formatted_query_prompt}" | |
| ) | |
| numbers = extract_numbers_given_response(response, n=self.child_branch_factor) | |
| if numbers is None: | |
| logging.debug( | |
| f">[Level {level}] Could not retrieve response - no numbers present" | |
| ) | |
| # just join text from current nodes as response | |
| return response | |
| result_response = None | |
| for number_str in numbers: | |
| number = int(number_str) | |
| if number > len(cur_node_list): | |
| logging.debug( | |
| f">[Level {level}] Invalid response: {response} - " | |
| f"number {number} out of range" | |
| ) | |
| return response | |
| # number is 1-indexed, so subtract 1 | |
| selected_node = cur_node_list[number - 1] | |
| logging.info( | |
| f">[Level {level}] Selected node: " | |
| f"[{number}]/[{','.join([str(int(n)) for n in numbers])}]" | |
| ) | |
| debug_str = " ".join(selected_node.get_text().splitlines()) | |
| logging.debug( | |
| f">[Level {level}] Node " | |
| f"[{number}] Summary text: " | |
| f"{ truncate_text(debug_str, 100) }" | |
| ) | |
| result_response = self._query_with_selected_node( | |
| selected_node, | |
| query_bundle, | |
| prev_response=result_response, | |
| level=level, | |
| ) | |
| # result_response should not be None | |
| return cast(str, result_response) | |
| def _query(self, query_bundle: QueryBundle) -> Response: | |
| """Answer a query.""" | |
| # NOTE: this overrides the _query method in the base class | |
| logging.info(f"> Starting query: {query_bundle.query_str}") | |
| response_str = self._query_level( | |
| self.index_struct.root_nodes, | |
| query_bundle, | |
| level=0, | |
| ).strip() | |
| return Response(response_str, source_nodes=self.response_builder.get_sources()) | |