Initial commit

api/core/index/index_builder.py

@@ -0,0 +1,45 @@
+from langchain.callbacks import CallbackManager
+from llama_index import ServiceContext, PromptHelper, LLMPredictor
+from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
+from core.embedding.openai_embedding import OpenAIEmbedding
+from core.llm.llm_builder import LLMBuilder
+
+
+class IndexBuilder:
+    @classmethod
+    def get_default_service_context(cls, tenant_id: str) -> ServiceContext:
+        # set number of output tokens
+        num_output = 512
+
+        # only for verbose
+        callback_manager = CallbackManager([DifyStdOutCallbackHandler()])
+
+        llm = LLMBuilder.to_llm(
+            tenant_id=tenant_id,
+            model_name='text-davinci-003',
+            temperature=0,
+            max_tokens=num_output,
+            callback_manager=callback_manager,
+        )
+
+        llm_predictor = LLMPredictor(llm=llm)
+
+        # These parameters here will affect the logic of segmenting the final synthesized response.
+        # The number of refinement iterations in the synthesis process depends
+        # on whether the length of the segmented output exceeds the max_input_size.
+        prompt_helper = PromptHelper(
+            max_input_size=3500,
+            num_output=num_output,
+            max_chunk_overlap=20
+        )
+
+        model_credentials = LLMBuilder.get_model_credentials(
+            tenant_id=tenant_id,
+            model_name='text-embedding-ada-002'
+        )
+
+        return ServiceContext.from_defaults(
+            llm_predictor=llm_predictor,
+            prompt_helper=prompt_helper,
+            embed_model=OpenAIEmbedding(**model_credentials),
+        )

api/core/index/keyword_table/jieba_keyword_table.py

@@ -0,0 +1,159 @@
+import re
+from typing import (
+    Any,
+    Dict,
+    List,
+    Set,
+    Optional
+)
+
+import jieba.analyse
+
+from core.index.keyword_table.stopwords import STOPWORDS
+from llama_index.indices.query.base import IS
+from llama_index import QueryMode
+from llama_index.indices.base import QueryMap
+from llama_index.indices.keyword_table.base import BaseGPTKeywordTableIndex
+from llama_index.indices.keyword_table.query import BaseGPTKeywordTableQuery
+from llama_index.docstore import BaseDocumentStore
+from llama_index.indices.postprocessor.node import (
+    BaseNodePostprocessor,
+)
+from llama_index.indices.response.response_builder import ResponseMode
+from llama_index.indices.service_context import ServiceContext
+from llama_index.optimization.optimizer import BaseTokenUsageOptimizer
+from llama_index.prompts.prompts import (
+    QuestionAnswerPrompt,
+    RefinePrompt,
+    SimpleInputPrompt,
+)
+
+from core.index.query.synthesizer import EnhanceResponseSynthesizer
+
+
+def jieba_extract_keywords(
+        text_chunk: str,
+        max_keywords: Optional[int] = None,
+        expand_with_subtokens: bool = True,
+) -> Set[str]:
+    """Extract keywords with JIEBA tfidf."""
+    keywords = jieba.analyse.extract_tags(
+        sentence=text_chunk,
+        topK=max_keywords,
+    )
+
+    if expand_with_subtokens:
+        return set(expand_tokens_with_subtokens(keywords))
+    else:
+        return set(keywords)
+
+
+def expand_tokens_with_subtokens(tokens: Set[str]) -> Set[str]:
+    """Get subtokens from a list of tokens., filtering for stopwords."""
+    results = set()
+    for token in tokens:
+        results.add(token)
+        sub_tokens = re.findall(r"\w+", token)
+        if len(sub_tokens) > 1:
+            results.update({w for w in sub_tokens if w not in list(STOPWORDS)})
+
+    return results
+
+
+class GPTJIEBAKeywordTableIndex(BaseGPTKeywordTableIndex):
+    """GPT JIEBA Keyword Table Index.
+
+    This index uses a JIEBA keyword extractor to extract keywords from the text.
+
+    """
+
+    def _extract_keywords(self, text: str) -> Set[str]:
+        """Extract keywords from text."""
+        return jieba_extract_keywords(text, max_keywords=self.max_keywords_per_chunk)
+
+    @classmethod
+    def get_query_map(self) -> QueryMap:
+        """Get query map."""
+        super_map = super().get_query_map()
+        super_map[QueryMode.DEFAULT] = GPTKeywordTableJIEBAQuery
+        return super_map
+
+    def _delete(self, doc_id: str, **delete_kwargs: Any) -> None:
+        """Delete a document."""
+        # get set of ids that correspond to node
+        node_idxs_to_delete = {doc_id}
+
+        # delete node_idxs from keyword to node idxs mapping
+        keywords_to_delete = set()
+        for keyword, node_idxs in self._index_struct.table.items():
+            if node_idxs_to_delete.intersection(node_idxs):
+                self._index_struct.table[keyword] = node_idxs.difference(
+                    node_idxs_to_delete
+                )
+                if not self._index_struct.table[keyword]:
+                    keywords_to_delete.add(keyword)
+
+        for keyword in keywords_to_delete:
+            del self._index_struct.table[keyword]
+
+
+class GPTKeywordTableJIEBAQuery(BaseGPTKeywordTableQuery):
+    """GPT Keyword Table Index JIEBA Query.
+
+    Extracts keywords using JIEBA keyword extractor.
+    Set when `mode="jieba"` in `query` method of `GPTKeywordTableIndex`.
+
+    .. code-block:: python
+
+        response = index.query("<query_str>", mode="jieba")
+
+    See BaseGPTKeywordTableQuery for arguments.
+
+    """
+
+    @classmethod
+    def from_args(
+            cls,
+            index_struct: IS,
+            service_context: ServiceContext,
+            docstore: Optional[BaseDocumentStore] = None,
+            node_postprocessors: Optional[List[BaseNodePostprocessor]] = None,
+            verbose: bool = False,
+            # response synthesizer args
+            response_mode: ResponseMode = ResponseMode.DEFAULT,
+            text_qa_template: Optional[QuestionAnswerPrompt] = None,
+            refine_template: Optional[RefinePrompt] = None,
+            simple_template: Optional[SimpleInputPrompt] = None,
+            response_kwargs: Optional[Dict] = None,
+            use_async: bool = False,
+            streaming: bool = False,
+            optimizer: Optional[BaseTokenUsageOptimizer] = None,
+            # class-specific args
+            **kwargs: Any,
+    ) -> "BaseGPTIndexQuery":
+        response_synthesizer = EnhanceResponseSynthesizer.from_args(
+            service_context=service_context,
+            text_qa_template=text_qa_template,
+            refine_template=refine_template,
+            simple_template=simple_template,
+            response_mode=response_mode,
+            response_kwargs=response_kwargs,
+            use_async=use_async,
+            streaming=streaming,
+            optimizer=optimizer,
+        )
+        return cls(
+            index_struct=index_struct,
+            service_context=service_context,
+            response_synthesizer=response_synthesizer,
+            docstore=docstore,
+            node_postprocessors=node_postprocessors,
+            verbose=verbose,
+            **kwargs,
+        )
+
+    def _get_keywords(self, query_str: str) -> List[str]:
+        """Extract keywords."""
+        return list(
+            jieba_extract_keywords(query_str, max_keywords=self.max_keywords_per_query)
+        )

api/core/index/keyword_table/stopwords.py

@@ -0,0 +1,90 @@
+STOPWORDS = {
+    "during", "when", "but", "then", "further", "isn", "mustn't", "until", "own", "i", "couldn", "y", "only", "you've",
+    "ours", "who", "where", "ourselves", "has", "to", "was", "didn't", "themselves", "if", "against", "through", "her",
+    "an", "your", "can", "those", "didn", "about", "aren't", "shan't", "be", "not", "these", "again", "so", "t",
+    "theirs", "weren", "won't", "won", "itself", "just", "same", "while", "why", "doesn", "aren", "him", "haven",
+    "for", "you'll", "that", "we", "am", "d", "by", "having", "wasn't", "than", "weren't", "out", "from", "now",
+    "their", "too", "hadn", "o", "needn", "most", "it", "under", "needn't", "any", "some", "few", "ll", "hers", "which",
+    "m", "you're", "off", "other", "had", "she", "you'd", "do", "you", "does", "s", "will", "each", "wouldn't", "hasn't",
+    "such", "more", "whom", "she's", "my", "yours", "yourself", "of", "on", "very", "hadn't", "with", "yourselves",
+    "been", "ma", "them", "mightn't", "shan", "mustn", "they", "what", "both", "that'll", "how", "is", "he", "because",
+    "down", "haven't", "are", "no", "it's", "our", "being", "the", "or", "above", "myself", "once", "don't", "doesn't",
+    "as", "nor", "here", "herself", "hasn", "mightn", "have", "its", "all", "were", "ain", "this", "at", "after",
+    "over", "shouldn't", "into", "before", "don", "wouldn", "re", "couldn't", "wasn", "in", "should", "there",
+    "himself", "isn't", "should've", "doing", "ve", "shouldn", "a", "did", "and", "his", "between", "me", "up", "below",
+    "人民", "末##末", "啊", "阿", "哎", "哎呀", "哎哟", "唉", "俺", "俺们", "按", "按照", "吧", "吧哒", "把", "罢了", "被", "本",
+    "本着", "比", "比方", "比如", "鄙人", "彼", "彼此", "边", "别", "别的", "别说", "并", "并且", "不比", "不成", "不单", "不但",
+    "不独", "不管", "不光", "不过", "不仅", "不拘", "不论", "不怕", "不然", "不如", "不特", "不惟", "不问", "不只", "朝", "朝着",
+    "趁", "趁着", "乘", "冲", "除", "除此之外", "除非", "除了", "此", "此间", "此外", "从", "从而", "打", "待", "但", "但是", "当",
+    "当着", "到", "得", "的", "的话", "等", "等等", "地", "第", "叮咚", "对", "对于", "多", "多少", "而", "而况", "而且", "而是",
+    "而外", "而言", "而已", "尔后", "反过来", "反过来说", "反之", "非但", "非徒", "否则", "嘎", "嘎登", "该", "赶", "个", "各",
+    "各个", "各位", "各种", "各自", "给", "根据", "跟", "故", "故此", "固然", "关于", "管", "归", "果然", "果真", "过", "哈",
+    "哈哈", "呵", "和", "何", "何处", "何况", "何时", "嘿", "哼", "哼唷", "呼哧", "乎", "哗", "还是", "还有", "换句话说", "换言之",
+    "或", "或是", "或者", "极了", "及", "及其", "及至", "即", "即便", "即或", "即令", "即若", "即使", "几", "几时", "己", "既",
+    "既然", "既是", "继而", "加之", "假如", "假若", "假使", "鉴于", "将", "较", "较之", "叫", "接着", "结果", "借", "紧接着",
+    "进而", "尽", "尽管", "经", "经过", "就", "就是", "就是说", "据", "具体地说", "具体说来", "开始", "开外", "靠", "咳", "可",
+    "可见", "可是", "可以", "况且", "啦", "来", "来着", "离", "例如", "哩", "连", "连同", "两者", "了", "临", "另", "另外",
+    "另一方面", "论", "嘛", "吗", "慢说", "漫说", "冒", "么", "每", "每当", "们", "莫若", "某", "某个", "某些", "拿", "哪",
+    "哪边", "哪儿", "哪个", "哪里", "哪年", "哪怕", "哪天", "哪些", "哪样", "那", "那边", "那儿", "那个", "那会儿", "那里", "那么",
+    "那么些", "那么样", "那时", "那些", "那样", "乃", "乃至", "呢", "能", "你", "你们", "您", "宁", "宁可", "宁肯", "宁愿", "哦",
+    "呕", "啪达", "旁人", "呸", "凭", "凭借", "其", "其次", "其二", "其他", "其它", "其一", "其余", "其中", "起", "起见", "岂但",
+    "恰恰相反", "前后", "前者", "且", "然而", "然后", "然则", "让", "人家", "任", "任何", "任凭", "如", "如此", "如果", "如何",
+    "如其", "如若", "如上所述", "若", "若非", "若是", "啥", "上下", "尚且", "设若", "设使", "甚而", "甚么", "甚至", "省得", "时候",
+    "什么", "什么样", "使得", "是", "是的", "首先", "谁", "谁知", "顺", "顺着", "似的", "虽", "虽然", "虽说", "虽则", "随", "随着",
+    "所", "所以", "他", "他们", "他人", "它", "它们", "她", "她们", "倘", "倘或", "倘然", "倘若", "倘使", "腾", "替", "通过", "同",
+    "同时", "哇", "万一", "往", "望", "为", "为何", "为了", "为什么", "为着", "喂", "嗡嗡", "我", "我们", "呜", "呜呼", "乌乎",
+    "无论", "无宁", "毋宁", "嘻", "吓", "相对而言", "像", "向", "向着", "嘘", "呀", "焉", "沿", "沿着", "要", "要不", "要不然",
+    "要不是", "要么", "要是", "也", "也罢", "也好", "一", "一般", "一旦", "一方面", "一来", "一切", "一样", "一则", "依", "依照",
+    "矣", "以", "以便", "以及", "以免", "以至", "以至于", "以致", "抑或", "因", "因此", "因而", "因为", "哟", "用", "由",
+    "由此可见", "由于", "有", "有的", "有关", "有些", "又", "于", "于是", "于是乎", "与", "与此同时", "与否", "与其", "越是",
+    "云云", "哉", "再说", "再者", "在", "在下", "咱", "咱们", "则", "怎", "怎么", "怎么办", "怎么样", "怎样", "咋", "照", "照着",
+    "者", "这", "这边", "这儿", "这个", "这会儿", "这就是说", "这里", "这么", "这么点儿", "这么些", "这么样", "这时", "这些", "这样",
+    "正如", "吱", "之", "之类", "之所以", "之一", "只是", "只限", "只要", "只有", "至", "至于", "诸位", "着", "着呢", "自", "自从",
+    "自个儿", "自各儿", "自己", "自家", "自身", "综上所述", "总的来看", "总的来说", "总的说来", "总而言之", "总之", "纵", "纵令",
+    "纵然", "纵使", "遵照", "作为", "兮", "呃", "呗", "咚", "咦", "喏", "啐", "喔唷", "嗬", "嗯", "嗳", "~", "!", ".", ":",
+    "\"", "'", "(", ")", "*", "A", "白", "社会主义", "--", "..", ">>", " [", " ]", "", "<", ">", "/", "\\", "|", "-", "_",
+    "+", "=", "&", "^", "%", "#", "@", "`", ";", "$", "（", "）", "——", "—", "￥", "·", "...", "‘", "’", "〉", "〈", "…",
+    "　", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "０", "１", "２", "３", "４", "５", "６", "７", "８", "９", "二",
+    "三", "四", "五", "六", "七", "八", "九", "零", "＞", "＜", "＠", "＃", "＄", "％", "︿", "＆", "＊", "＋", "～", "｜", "［",
+    "］", "｛", "｝", "啊哈", "啊呀", "啊哟", "挨次", "挨个", "挨家挨户", "挨门挨户", "挨门逐户", "挨着", "按理", "按期", "按时",
+    "按说", "暗地里", "暗中", "暗自", "昂然", "八成", "白白", "半", "梆", "保管", "保险", "饱", "背地里", "背靠背", "倍感", "倍加",
+    "本人", "本身", "甭", "比起", "比如说", "比照", "毕竟", "必", "必定", "必将", "必须", "便", "别人", "并非", "并肩", "并没",
+    "并没有", "并排", "并无", "勃然", "不", "不必", "不常", "不大", "不但...而且", "不得", "不得不", "不得了", "不得已", "不迭",
+    "不定", "不对", "不妨", "不管怎样", "不会", "不仅...而且", "不仅仅", "不仅仅是", "不经意", "不可开交", "不可抗拒", "不力", "不了",
+    "不料", "不满", "不免", "不能不", "不起", "不巧", "不然的话", "不日", "不少", "不胜", "不时", "不是", "不同", "不能", "不要",
+    "不外", "不外乎", "不下", "不限", "不消", "不已", "不亦乐乎", "不由得", "不再", "不择手段", "不怎么", "不曾", "不知不觉", "不止",
+    "不止一次", "不至于", "才", "才能", "策略地", "差不多", "差一点", "常", "常常", "常言道", "常言说", "常言说得好", "长此下去",
+    "长话短说", "长期以来", "长线", "敞开儿", "彻夜", "陈年", "趁便", "趁机", "趁热", "趁势", "趁早", "成年", "成年累月", "成心",
+    "乘机", "乘胜", "乘势", "乘隙", "乘虚", "诚然", "迟早", "充分", "充其极", "充其量", "抽冷子", "臭", "初", "出", "出来", "出去",
+    "除此", "除此而外", "除此以外", "除开", "除去", "除却", "除外", "处处", "川流不息", "传", "传说", "传闻", "串行", "纯", "纯粹",
+    "此后", "此中", "次第", "匆匆", "从不", "从此", "从此以后", "从古到今", "从古至今", "从今以后", "从宽", "从来", "从轻", "从速",
+    "从头", "从未", "从无到有", "从小", "从新", "从严", "从优", "从早到晚", "从中", "从重", "凑巧", "粗", "存心", "达旦", "打从",
+    "打开天窗说亮话", "大", "大不了", "大大", "大抵", "大都", "大多", "大凡", "大概", "大家", "大举", "大略", "大面儿上", "大事",
+    "大体", "大体上", "大约", "大张旗鼓", "大致", "呆呆地", "带", "殆", "待到", "单", "单纯", "单单", "但愿", "弹指之间", "当场",
+    "当儿", "当即", "当口儿", "当然", "当庭", "当头", "当下", "当真", "当中", "倒不如", "倒不如说", "倒是", "到处", "到底", "到了儿",
+    "到目前为止", "到头", "到头来", "得起", "得天独厚", "的确", "等到", "叮当", "顶多", "定", "动不动", "动辄", "陡然", "都", "独",
+    "独自", "断然", "顿时", "多次", "多多", "多多少少", "多多益善", "多亏", "多年来", "多年前", "而后", "而论", "而又", "尔等",
+    "二话不说", "二话没说", "反倒", "反倒是", "反而", "反手", "反之亦然", "反之则", "方", "方才", "方能", "放量", "非常", "非得",
+    "分期", "分期分批", "分头", "奋勇", "愤然", "风雨无阻", "逢", "弗", "甫", "嘎嘎", "该当", "概", "赶快", "赶早不赶晚", "敢",
+    "敢情", "敢于", "刚", "刚才", "刚好", "刚巧", "高低", "格外", "隔日", "隔夜", "个人", "各式", "更", "更加", "更进一步", "更为",
+    "公然", "共", "共总", "够瞧的", "姑且", "古来", "故而", "故意", "固", "怪", "怪不得", "惯常", "光", "光是", "归根到底",
+    "归根结底", "过于", "毫不", "毫无", "毫无保留地", "毫无例外", "好在", "何必", "何尝", "何妨", "何苦", "何乐而不为", "何须",
+    "何止", "很", "很多", "很少", "轰然", "后来", "呼啦", "忽地", "忽然", "互", "互相", "哗啦", "话说", "还", "恍然", "会", "豁然",
+    "活", "伙同", "或多或少", "或许", "基本", "基本上", "基于", "极", "极大", "极度", "极端", "极力", "极其", "极为", "急匆匆",
+    "即将", "即刻", "即是说", "几度", "几番", "几乎", "几经", "既...又", "继之", "加上", "加以", "间或", "简而言之", "简言之",
+    "简直", "见", "将才", "将近", "将要", "交口", "较比", "较为", "接连不断", "接下来", "皆可", "截然", "截至", "藉以", "借此",
+    "借以", "届时", "仅", "仅仅", "谨", "进来", "进去", "近", "近几年来", "近来", "近年来", "尽管如此", "尽可能", "尽快", "尽量",
+    "尽然", "尽如人意", "尽心竭力", "尽心尽力", "尽早", "精光", "经常", "竟", "竟然", "究竟", "就此", "就地", "就算", "居然", "局外",
+    "举凡", "据称", "据此", "据实", "据说", "据我所知", "据悉", "具体来说", "决不", "决非", "绝", "绝不", "绝顶", "绝对", "绝非",
+    "均", "喀", "看", "看来", "看起来", "看上去", "看样子", "可好", "可能", "恐怕", "快", "快要", "来不及", "来得及", "来讲",
+    "来看", "拦腰", "牢牢", "老", "老大", "老老实实", "老是", "累次", "累年", "理当", "理该", "理应", "历", "立", "立地", "立刻",
+    "立马", "立时", "联袂", "连连", "连日", "连日来", "连声", "连袂", "临到", "另方面", "另行", "另一个", "路经", "屡", "屡次",
+    "屡次三番", "屡屡", "缕缕", "率尔", "率然", "略", "略加", "略微", "略为", "论说", "马上", "蛮", "满", "没", "没有", "每逢",
+    "每每", "每时每刻", "猛然", "猛然间", "莫", "莫不", "莫非", "莫如", "默默地", "默然", "呐", "那末", "奈", "难道", "难得", "难怪",
+    "难说", "内", "年复一年", "凝神", "偶而", "偶尔", "怕", "砰", "碰巧", "譬如", "偏偏", "乒", "平素", "颇", "迫于", "扑通",
+    "其后", "其实", "奇", "齐", "起初", "起来", "起首", "起头", "起先", "岂", "岂非", "岂止", "迄", "恰逢", "恰好", "恰恰", "恰巧",
+    "恰如", "恰似", "千", "千万", "千万千万", "切", "切不可", "切莫", "切切", "切勿", "窃", "亲口", "亲身", "亲手", "亲眼", "亲自",
+    "顷", "顷刻", "顷刻间", "顷刻之间", "请勿", "穷年累月", "取道", "去", "权时", "全都", "全力", "全年", "全然", "全身心", "然",
+    "人人", "仍", "仍旧", "仍然", "日复一日", "日见", "日渐", "日益", "日臻", "如常", "如此等等", "如次", "如今", "如期", "如前所述",
+    "如上", "如下", "汝", "三番两次", "三番五次", "三天两头", "瑟瑟", "沙沙", "上", "上来", "上去", "一个", "月", "日", "\n"
+}

api/core/index/keyword_table_index.py

@@ -0,0 +1,135 @@
+import json
+from typing import List, Optional
+
+from llama_index import ServiceContext, LLMPredictor, OpenAIEmbedding
+from llama_index.data_structs import KeywordTable, Node
+from llama_index.indices.keyword_table.base import BaseGPTKeywordTableIndex
+from llama_index.indices.registry import load_index_struct_from_dict
+
+from core.docstore.dataset_docstore import DatesetDocumentStore
+from core.docstore.empty_docstore import EmptyDocumentStore
+from core.index.index_builder import IndexBuilder
+from core.index.keyword_table.jieba_keyword_table import GPTJIEBAKeywordTableIndex
+from core.llm.llm_builder import LLMBuilder
+from extensions.ext_database import db
+from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment
+
+
+class KeywordTableIndex:
+
+    def __init__(self, dataset: Dataset):
+        self._dataset = dataset
+
+    def add_nodes(self, nodes: List[Node]):
+        llm = LLMBuilder.to_llm(
+            tenant_id=self._dataset.tenant_id,
+            model_name='fake'
+        )
+
+        service_context = ServiceContext.from_defaults(
+            llm_predictor=LLMPredictor(llm=llm),
+            embed_model=OpenAIEmbedding()
+        )
+
+        dataset_keyword_table = self.get_keyword_table()
+        if not dataset_keyword_table or not dataset_keyword_table.keyword_table_dict:
+            index_struct = KeywordTable()
+        else:
+            index_struct_dict = dataset_keyword_table.keyword_table_dict
+            index_struct: KeywordTable = load_index_struct_from_dict(index_struct_dict)
+
+        # create index
+        index = GPTJIEBAKeywordTableIndex(
+            index_struct=index_struct,
+            docstore=EmptyDocumentStore(),
+            service_context=service_context
+        )
+
+        for node in nodes:
+            keywords = index._extract_keywords(node.get_text())
+            self.update_segment_keywords(node.doc_id, list(keywords))
+            index._index_struct.add_node(list(keywords), node)
+
+        index_struct_dict = index.index_struct.to_dict()
+
+        if not dataset_keyword_table:
+            dataset_keyword_table = DatasetKeywordTable(
+                dataset_id=self._dataset.id,
+                keyword_table=json.dumps(index_struct_dict)
+            )
+            db.session.add(dataset_keyword_table)
+        else:
+            dataset_keyword_table.keyword_table = json.dumps(index_struct_dict)
+
+        db.session.commit()
+
+    def del_nodes(self, node_ids: List[str]):
+        llm = LLMBuilder.to_llm(
+            tenant_id=self._dataset.tenant_id,
+            model_name='fake'
+        )
+
+        service_context = ServiceContext.from_defaults(
+            llm_predictor=LLMPredictor(llm=llm),
+            embed_model=OpenAIEmbedding()
+        )
+
+        dataset_keyword_table = self.get_keyword_table()
+        if not dataset_keyword_table or not dataset_keyword_table.keyword_table_dict:
+            return
+        else:
+            index_struct_dict = dataset_keyword_table.keyword_table_dict
+            index_struct: KeywordTable = load_index_struct_from_dict(index_struct_dict)
+
+        # create index
+        index = GPTJIEBAKeywordTableIndex(
+            index_struct=index_struct,
+            docstore=EmptyDocumentStore(),
+            service_context=service_context
+        )
+
+        for node_id in node_ids:
+            index.delete(node_id)
+
+        index_struct_dict = index.index_struct.to_dict()
+
+        if not dataset_keyword_table:
+            dataset_keyword_table = DatasetKeywordTable(
+                dataset_id=self._dataset.id,
+                keyword_table=json.dumps(index_struct_dict)
+            )
+            db.session.add(dataset_keyword_table)
+        else:
+            dataset_keyword_table.keyword_table = json.dumps(index_struct_dict)
+
+        db.session.commit()
+
+    @property
+    def query_index(self) -> Optional[BaseGPTKeywordTableIndex]:
+        docstore = DatesetDocumentStore(
+            dataset=self._dataset,
+            user_id=self._dataset.created_by,
+            embedding_model_name="text-embedding-ada-002"
+        )
+
+        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
+
+        dataset_keyword_table = self.get_keyword_table()
+        if not dataset_keyword_table or not dataset_keyword_table.keyword_table_dict:
+            return None
+
+        index_struct: KeywordTable = load_index_struct_from_dict(dataset_keyword_table.keyword_table_dict)
+
+        return GPTJIEBAKeywordTableIndex(index_struct=index_struct, docstore=docstore, service_context=service_context)
+
+    def get_keyword_table(self):
+        dataset_keyword_table = self._dataset.dataset_keyword_table
+        if dataset_keyword_table:
+            return dataset_keyword_table
+        return None
+
+    def update_segment_keywords(self, node_id: str, keywords: List[str]):
+        document_segment = db.session.query(DocumentSegment).filter(DocumentSegment.index_node_id == node_id).first()
+        if document_segment:
+            document_segment.keywords = keywords
+            db.session.commit()

api/core/index/query/synthesizer.py

@@ -0,0 +1,79 @@
+from typing import (
+    Any,
+    Dict,
+    Optional, Sequence,
+)
+
+from llama_index.indices.response.response_synthesis import ResponseSynthesizer
+from llama_index.indices.response.response_builder import ResponseMode, BaseResponseBuilder, get_response_builder
+from llama_index.indices.service_context import ServiceContext
+from llama_index.optimization.optimizer import BaseTokenUsageOptimizer
+from llama_index.prompts.prompts import (
+    QuestionAnswerPrompt,
+    RefinePrompt,
+    SimpleInputPrompt,
+)
+from llama_index.types import RESPONSE_TEXT_TYPE
+
+
+class EnhanceResponseSynthesizer(ResponseSynthesizer):
+    @classmethod
+    def from_args(
+            cls,
+            service_context: ServiceContext,
+            streaming: bool = False,
+            use_async: bool = False,
+            text_qa_template: Optional[QuestionAnswerPrompt] = None,
+            refine_template: Optional[RefinePrompt] = None,
+            simple_template: Optional[SimpleInputPrompt] = None,
+            response_mode: ResponseMode = ResponseMode.DEFAULT,
+            response_kwargs: Optional[Dict] = None,
+            optimizer: Optional[BaseTokenUsageOptimizer] = None,
+    ) -> "ResponseSynthesizer":
+        response_builder: Optional[BaseResponseBuilder] = None
+        if response_mode != ResponseMode.NO_TEXT:
+            if response_mode == 'no_synthesizer':
+                response_builder = NoSynthesizer(
+                    service_context=service_context,
+                    simple_template=simple_template,
+                    streaming=streaming,
+                )
+            else:
+                response_builder = get_response_builder(
+                    service_context,
+                    text_qa_template,
+                    refine_template,
+                    simple_template,
+                    response_mode,
+                    use_async=use_async,
+                    streaming=streaming,
+                )
+        return cls(response_builder, response_mode, response_kwargs, optimizer)
+
+
+class NoSynthesizer(BaseResponseBuilder):
+    def __init__(
+            self,
+            service_context: ServiceContext,
+            simple_template: Optional[SimpleInputPrompt] = None,
+            streaming: bool = False,
+    ) -> None:
+        super().__init__(service_context, streaming)
+
+    async def aget_response(
+            self,
+            query_str: str,
+            text_chunks: Sequence[str],
+            prev_response: Optional[str] = None,
+            **response_kwargs: Any,
+    ) -> RESPONSE_TEXT_TYPE:
+        return "\n".join(text_chunks)
+
+    def get_response(
+            self,
+            query_str: str,
+            text_chunks: Sequence[str],
+            prev_response: Optional[str] = None,
+            **response_kwargs: Any,
+    ) -> RESPONSE_TEXT_TYPE:
+        return "\n".join(text_chunks)

api/core/index/readers/html_parser.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+from typing import Dict
+
+from bs4 import BeautifulSoup
+from llama_index.readers.file.base_parser import BaseParser
+
+
+class HTMLParser(BaseParser):
+    """HTML parser."""
+
+    def _init_parser(self) -> Dict:
+        """Init parser."""
+        return {}
+
+    def parse_file(self, file: Path, errors: str = "ignore") -> str:
+        """Parse file."""
+        with open(file, "rb") as fp:
+            soup = BeautifulSoup(fp, 'html.parser')
+            text = soup.get_text()
+            text = text.strip() if text else ''
+
+        return text

api/core/index/readers/pdf_parser.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+from typing import Dict
+
+from flask import current_app
+from llama_index.readers.file.base_parser import BaseParser
+from pypdf import PdfReader
+
+from extensions.ext_storage import storage
+from models.model import UploadFile
+
+
+class PDFParser(BaseParser):
+    """PDF parser."""
+
+    def _init_parser(self) -> Dict:
+        """Init parser."""
+        return {}
+
+    def parse_file(self, file: Path, errors: str = "ignore") -> str:
+        """Parse file."""
+        if not current_app.config.get('PDF_PREVIEW', True):
+            return ''
+
+        plaintext_file_key = ''
+        plaintext_file_exists = False
+        if self._parser_config and 'upload_file' in self._parser_config and self._parser_config['upload_file']:
+            upload_file: UploadFile = self._parser_config['upload_file']
+            if upload_file.hash:
+                plaintext_file_key = 'upload_files/' + upload_file.tenant_id + '/' + upload_file.hash + '.plaintext'
+                try:
+                    text = storage.load(plaintext_file_key).decode('utf-8')
+                    plaintext_file_exists = True
+                    return text
+                except FileNotFoundError:
+                    pass
+
+        text_list = []
+        with open(file, "rb") as fp:
+            # Create a PDF object
+            pdf = PdfReader(fp)
+
+            # Get the number of pages in the PDF document
+            num_pages = len(pdf.pages)
+
+            # Iterate over every page
+            for page in range(num_pages):
+                # Extract the text from the page
+                page_text = pdf.pages[page].extract_text()
+                text_list.append(page_text)
+        text = "\n".join(text_list)
+
+        # save plaintext file for caching
+        if not plaintext_file_exists and plaintext_file_key:
+            storage.save(plaintext_file_key, text.encode('utf-8'))
+
+        return text

api/core/index/vector_index.py

@@ -0,0 +1,136 @@
+import json
+import logging
+from typing import List, Optional
+
+from llama_index.data_structs import Node
+from requests import ReadTimeout
+from sqlalchemy.exc import IntegrityError
+from tenacity import retry, stop_after_attempt, retry_if_exception_type
+
+from core.index.index_builder import IndexBuilder
+from core.vector_store.base import BaseGPTVectorStoreIndex
+from extensions.ext_vector_store import vector_store
+from extensions.ext_database import db
+from models.dataset import Dataset, Embedding
+
+
+class VectorIndex:
+
+    def __init__(self, dataset: Dataset):
+        self._dataset = dataset
+
+    def add_nodes(self, nodes: List[Node], duplicate_check: bool = False):
+        if not self._dataset.index_struct_dict:
+            index_id = "Vector_index_" + self._dataset.id.replace("-", "_")
+            self._dataset.index_struct = json.dumps(vector_store.to_index_struct(index_id))
+            db.session.commit()
+
+        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
+
+        index = vector_store.get_index(
+            service_context=service_context,
+            index_struct=self._dataset.index_struct_dict
+        )
+
+        if duplicate_check:
+            nodes = self._filter_duplicate_nodes(index, nodes)
+
+        embedding_queue_nodes = []
+        embedded_nodes = []
+        for node in nodes:
+            node_hash = node.doc_hash
+
+            # if node hash in cached embedding tables, use cached embedding
+            embedding = db.session.query(Embedding).filter_by(hash=node_hash).first()
+            if embedding:
+                node.embedding = embedding.get_embedding()
+                embedded_nodes.append(node)
+            else:
+                embedding_queue_nodes.append(node)
+
+        if embedding_queue_nodes:
+            embedding_results = index._get_node_embedding_results(
+                embedding_queue_nodes,
+                set(),
+            )
+
+            # pre embed nodes for cached embedding
+            for embedding_result in embedding_results:
+                node = embedding_result.node
+                node.embedding = embedding_result.embedding
+
+                try:
+                    embedding = Embedding(hash=node.doc_hash)
+                    embedding.set_embedding(node.embedding)
+                    db.session.add(embedding)
+                    db.session.commit()
+                except IntegrityError:
+                    db.session.rollback()
+                    continue
+                except:
+                    logging.exception('Failed to add embedding to db')
+                    continue
+
+                embedded_nodes.append(node)
+
+        self.index_insert_nodes(index, embedded_nodes)
+
+    @retry(reraise=True, retry=retry_if_exception_type(ReadTimeout), stop=stop_after_attempt(3))
+    def index_insert_nodes(self, index: BaseGPTVectorStoreIndex, nodes: List[Node]):
+        index.insert_nodes(nodes)
+
+    def del_nodes(self, node_ids: List[str]):
+        if not self._dataset.index_struct_dict:
+            return
+
+        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
+
+        index = vector_store.get_index(
+            service_context=service_context,
+            index_struct=self._dataset.index_struct_dict
+        )
+
+        for node_id in node_ids:
+            self.index_delete_node(index, node_id)
+
+    @retry(reraise=True, retry=retry_if_exception_type(ReadTimeout), stop=stop_after_attempt(3))
+    def index_delete_node(self, index: BaseGPTVectorStoreIndex, node_id: str):
+        index.delete_node(node_id)
+
+    def del_doc(self, doc_id: str):
+        if not self._dataset.index_struct_dict:
+            return
+
+        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
+
+        index = vector_store.get_index(
+            service_context=service_context,
+            index_struct=self._dataset.index_struct_dict
+        )
+
+        self.index_delete_doc(index, doc_id)
+
+    @retry(reraise=True, retry=retry_if_exception_type(ReadTimeout), stop=stop_after_attempt(3))
+    def index_delete_doc(self, index: BaseGPTVectorStoreIndex, doc_id: str):
+        index.delete(doc_id)
+
+    @property
+    def query_index(self) -> Optional[BaseGPTVectorStoreIndex]:
+        if not self._dataset.index_struct_dict:
+            return None
+
+        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
+
+        return vector_store.get_index(
+            service_context=service_context,
+            index_struct=self._dataset.index_struct_dict
+        )
+
+    def _filter_duplicate_nodes(self, index: BaseGPTVectorStoreIndex, nodes: List[Node]) -> List[Node]:
+        for node in nodes:
+            node_id = node.doc_id
+            exists_duplicate_node = index.exists_by_node_id(node_id)
+            if exists_duplicate_node:
+                nodes.remove(node)
+
+        return nodes