diff --git a/gpt_index/indices/list/base.py b/gpt_index/indices/list/base.py index a994ab56c27ced98dc744954f7d6b04032cc3de3..1f763cf1f6db351a84a707eee6f8af232b1e367f 100644 --- a/gpt_index/indices/list/base.py +++ b/gpt_index/indices/list/base.py @@ -14,6 +14,7 @@ from gpt_index.indices.data_structs import IndexList from gpt_index.indices.list.query import GPTListIndexQuery from gpt_index.indices.utils import get_chunk_size_given_prompt, truncate_text from gpt_index.langchain_helpers.text_splitter import TokenTextSplitter +from gpt_index.prompts.base import Prompt from gpt_index.prompts.default_prompts import DEFAULT_TEXT_QA_PROMPT from gpt_index.schema import Document @@ -25,7 +26,7 @@ class GPTListIndex(BaseGPTIndex[IndexList]): self, documents: Optional[List[Document]] = None, index_struct: Optional[IndexList] = None, - text_qa_template: str = DEFAULT_TEXT_QA_PROMPT, + text_qa_template: Prompt = DEFAULT_TEXT_QA_PROMPT, ) -> None: """Initialize params.""" self.text_qa_template = text_qa_template diff --git a/gpt_index/indices/list/query.py b/gpt_index/indices/list/query.py index 5bb7bbc5d9bd2d9d3eff5cf1bcd51003c3841854..b38a47ae335f75f343e5759a1d98d132df224f55 100644 --- a/gpt_index/indices/list/query.py +++ b/gpt_index/indices/list/query.py @@ -1,4 +1,6 @@ """Default query for GPTListIndex.""" +from typing import Optional + from gpt_index.indices.base import BaseGPTIndexQuery from gpt_index.indices.data_structs import IndexList from gpt_index.indices.response_utils import give_response, refine_response @@ -18,17 +20,21 @@ class GPTListIndexQuery(BaseGPTIndexQuery[IndexList]): index_struct: IndexList, text_qa_template: Prompt = DEFAULT_TEXT_QA_PROMPT, refine_template: Prompt = DEFAULT_REFINE_PROMPT, + keyword: Optional[str] = None, ) -> None: """Initialize params.""" super().__init__(index_struct=index_struct) self.text_qa_template = text_qa_template self.refine_template = refine_template + self.keyword = keyword def query(self, query_str: str, verbose: bool = False) -> str: """Answer a query.""" print(f"> Starting query: {query_str}") response = None for node in self.index_struct.nodes: + if self.keyword is not None and self.keyword not in node.text: + continue fmt_text_chunk = truncate_text(node.text, 50) if verbose: print(f"> Searching in chunk: {fmt_text_chunk}") diff --git a/requirements.txt b/requirements.txt index a211e87c815a32dc1d0100d5375f0796bc06a3f2..2c9695e47db833dc98edb46f3106c4270d19dc6f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,6 @@ black isort mypy -# flake8 -# flake8-docstrings +flake8 +flake8-docstrings pylint