Add wiki to the rag system

routers/embedding/{embeddings_manual.pkl → embeddings_manual_wiki.pkl}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9ed7475fc8ffda0d9e9deb6480b7152b53657f0fe6a6140bcb60360e425e7a01
-size 18659241
+oid sha256:3c4a71f60f1878e528b190c3c43f744611c90efdea4c2ef333962773fd2fd637
+size 19670346

routers/tool_wiki_search.py

@@ -1,28 +1,60 @@
-# routers/wiki_search.py
+# routers/tool_wiki_search.py
 
+import base64
 import os
 import pickle
 import re
 import torch
+from enum import Enum
 from typing import Dict, List
 from sentence_transformers import util
 from fastapi import APIRouter
 from fastapi.responses import PlainTextResponse
+from utils_gitea import gitea_wiki_page_get, gitea_wiki_pages_get
 
 try:
     from .embedding import EMBEDDING_CTX
 except:
     from embedding import EMBEDDING_CTX
 
-router = APIRouter()
-
 MANUAL_DIR = "D:/BlenderDev/blender-manual/manual/"
-BASE_URL = "https://docs.blender.org/manual/en/dev"
-G_data = None
+
+
+class Group(str, Enum):
+    wiki = "wiki"
+    manual = "manual"
+    all = "all"
 
 
 class _Data(dict):
-    cache_path = "routers/embedding/embeddings_manual.pkl"
+    cache_path = "routers/embedding/embeddings_manual_wiki.pkl"
+
+    def __init__(self):
+        if os.path.exists(self.cache_path):
+            with open(self.cache_path, 'rb') as file:
+                data = pickle.load(file)
+                self.update(data)
+                return
+
+        # Generate
+
+        print("Embedding Texts...")
+        for grp in list(Group)[:-1]:
+            self[grp.name] = {}
+
+            # Create a list to store the text files
+            texts = self.manual_get_texts_to_embed(
+            ) if grp == Group.manual else self.wiki_get_texts_to_embed()
+
+            self[grp]['texts'] = texts
+            self[grp]['embeddings'] = EMBEDDING_CTX.encode(texts)
+
+        with open(self.cache_path, "wb") as file:
+            # Converting the embeddings to be CPU compatible, as the virtual machine in use currently only supports the CPU.
+            for val in self.values():
+                val['embeddings'] = val['embeddings'].to(torch.device('cpu'))
+
+            pickle.dump(dict(self), file, protocol=pickle.HIGHEST_PROTOCOL)
 
     @staticmethod
     def reduce_text(text):
@@ -38,24 +70,20 @@ class _Data(dict):
         return text
 
     @classmethod
-    def parse_file_recursive(cls, filedir, filename):
-        with open(os.path.join(filedir, filename), 'r', encoding='utf-8') as file:
+    def parse_file_recursive(cls, filepath):
+        with open(filepath, 'r', encoding='utf-8') as file:
             content = file.read()
 
         parsed_data = {}
 
-        if not filename.endswith('index.rst'):
-            body = content.strip()
-        else:
+        if filepath.endswith('index.rst'):
+            filedir = os.path.dirname(filepath)
             parts = content.split(".. toctree::")
-            body = parts[0].strip()
-
             if len(parts) > 1:
                 parsed_data["toctree"] = {}
                 for part in parts[1:]:
-                    toctree_entries = part.split('\n')
-                    line = toctree_entries[0]
-                    for entry in toctree_entries[1:]:
+                    toctree_entries = part.splitlines()[1:]
+                    for entry in toctree_entries:
                         entry = entry.strip()
                         if not entry:
                             continue
@@ -67,20 +95,12 @@ class _Data(dict):
                         if not entry.endswith('.rst'):
                             continue
 
-                        if entry.endswith('/index.rst'):
-                            entry_name = entry[:-10]
-                            filedir_ = os.path.join(filedir, entry_name)
-                            filename_ = 'index.rst'
-                        else:
-                            entry_name = entry[:-4]
-                            filedir_ = filedir
-                            filename_ = entry
-
+                        entry_name = entry[:-4]  # remove '.rst'
+                        filepath_iter = os.path.join(filedir, entry)
                         parsed_data['toctree'][entry_name] = cls.parse_file_recursive(
-                            filedir_, filename_)
+                            filepath_iter)
 
-        # The '\n' at the end of the file resolves regex patterns
-        parsed_data['body'] = body + '\n'
+        parsed_data['body'] = content
 
         return parsed_data
 
@@ -221,82 +241,122 @@ class _Data(dict):
         return result
 
     @classmethod
-    def get_texts_recursive(cls, page, path=''):
+    def get_texts_recursive(cls, page, path='index'):
         result = cls.split_into_many(page['body'], path)
 
         try:
             for key in page['toctree'].keys():
                 page_child = page['toctree'][key]
                 result.extend(cls.get_texts_recursive(
-                    page_child, f'{path}/{key}'))
+                    page_child, path.replace('index', key)))
         except KeyError:
             pass
 
         return result
 
-    def _embeddings_generate(self):
-        if os.path.exists(self.cache_path):
-            with open(self.cache_path, 'rb') as file:
-                data = pickle.load(file)
-                self.update(data)
-                return self
-
-        # Generate
+    @classmethod
+    def manual_get_texts_to_embed(cls):
+        manual = cls.parse_file_recursive(
+            os.path.join(MANUAL_DIR, 'index.rst'))
+        manual['toctree']["copyright"] = cls.parse_file_recursive(
+            os.path.join(MANUAL_DIR, 'copyright.rst'))
 
-        manual = self.parse_file_recursive(MANUAL_DIR, 'index.rst')
-        manual['toctree']["copyright"] = self.parse_file_recursive(
-            MANUAL_DIR, 'copyright.rst')
+        return cls.get_texts_recursive(manual)
 
-        # Create a list to store the text files
-        texts = self.get_texts_recursive(manual)
+    @classmethod
+    def wiki_get_texts_to_embed(cls):
+        tokenizer = EMBEDDING_CTX.model.tokenizer
+        max_tokens = EMBEDDING_CTX.model.max_seq_length
 
-        print("Embedding Texts...")
-        self['texts'] = texts
-        self['embeddings'] = EMBEDDING_CTX.encode(texts)
+        texts = []
+        owner = "blender"
+        repo = "blender"
+        pages = gitea_wiki_pages_get(owner, repo)
+        for page_name in pages:
+            page_name_title = page_name["title"]
+            page = gitea_wiki_page_get(owner, repo, page_name_title)
+            prefix = f'/{page["sub_url"]}\n# {page_name_title}:'
+            text = base64.b64decode(page["content_base64"]).decode('utf-8')
+            text = text.replace(
+                'https://projects.blender.org/blender/blender', '')
+            tokens_prefix_len = len(tokenizer.tokenize(prefix))
+            tokens_so_far = tokens_prefix_len
+            text_so_far = prefix
+            text_parts = text.split('\n#')
+            for part in text_parts:
+                part = '\n#' + part
+                part_tokens_len = len(tokenizer.tokenize(part))
+                if tokens_so_far + part_tokens_len > max_tokens:
+                    texts.append(text_so_far)
+                    text_so_far = prefix
+                    tokens_so_far = tokens_prefix_len
+                text_so_far += part
+                tokens_so_far += part_tokens_len
+
+            if tokens_so_far != tokens_prefix_len:
+                texts.append(text_so_far)
+
+        return texts
+
+    def _sort_similarity(self, text_to_search, group: Group = Group.all, limit=4):
+        result = []
 
-        with open(self.cache_path, "wb") as file:
-            # Converting the embeddings to be CPU compatible, as the virtual machine in use currently only supports the CPU.
-            self['embeddings'] = self['embeddings'].to(torch.device('cpu'))
-            pickle.dump(dict(self), file, protocol=pickle.HIGHEST_PROTOCOL)
+        query_emb = EMBEDDING_CTX.encode([text_to_search])
 
-        return G_data
+        ret = {}
 
-    def _sort_similarity(self, text_to_search, limit):
-        results = []
+        for grp in list(Group)[:-1]:
+            if group in {grp, Group.all}:
+                ret[grp] = util.semantic_search(
+                    query_emb, self[grp]['embeddings'], top_k=limit, score_function=util.dot_score)
 
-        query_emb = EMBEDDING_CTX.encode([text_to_search])
-        ret = util.semantic_search(
-            query_emb, self['embeddings'], top_k=limit, score_function=util.dot_score)
+        score_best = 0.0
+        group_best = None
+        for grp, val in ret.items():
+            score_curr = val[0][0]['score']
+            if score_curr > score_best:
+                score_best = score_curr
+                group_best = grp
 
-        texts = self['texts']
-        for score in ret[0]:
+        texts = self[group_best]['texts']
+        for score in ret[group_best][0]:
             corpus_id = score['corpus_id']
             text = texts[corpus_id]
-            results.append(text)
+            result.append(text)
 
-        return results
+        return result, group_best
 
 
 G_data = _Data()
 
+router = APIRouter()
 
-@router.get("/wiki_search", response_class=PlainTextResponse)
-def wiki_search(query: str = "") -> str:
-    data = G_data._embeddings_generate()
-    texts = G_data._sort_similarity(query, 5)
 
-    result = f'BASE_URL: {BASE_URL}\n'
+@router.get("/wiki_search", response_class=PlainTextResponse)
+def wiki_search(query: str = "", group: Group = Group.all) -> str:
+    base_url = {
+        Group.wiki: "https://projects.blender.org/blender/blender",
+        Group.manual: "https://docs.blender.org/manual/en/dev"
+    }
+    texts, group_best = G_data._sort_similarity(query, group)
+
+    result = f'BASE_URL: {base_url[group_best]}\n'
     for text in texts:
-        index = text.find('#')
-        result += f'''---
+        if group_best == Group.wiki:
+            result += f'''---
+{text}
+'''
+        else:
+            index = text.find('#')
+            result += f'''---
 {text[:index] + '.html'}
 {text[index:]}
-
 '''
     return result
 
 
 if __name__ == '__main__':
-    tests = ["Set Snap Base", "Building the Manual", "Bisect Object"]
-    result = wiki_search(tests[0])
+    tests = ["Set Snap Base", "Building the Manual",
+             "Bisect Object", "Who are the Triagers"]
+    result = wiki_search(tests[1], Group.all)
     print(result)

routers/utils_gitea.py

@@ -30,7 +30,7 @@ def url_json_get(url, data=None):
 def url_json_get_all_pages(url, item_filter=None, limit=50, exclude=set(), verbose=False):
     assert limit <= 50, "50 is the maximum limit of items per page"
 
-    url_for_page = f"{url}&limit={limit}&page="
+    url_for_page = f"{url}?limit={limit}&page="
 
     with urllib.request.urlopen(url_for_page + '1') as response:
         headers_first = response.info()
@@ -82,7 +82,6 @@ def gitea_fetch_issues(owner, repo, state='all', labels='', issue_attr_filter=No
     if since:
         query_params['since'] = since
 
-    BASE_API_URL = "https://projects.blender.org/api/v1"
     base_url = f"{BASE_API_URL}/repos/{owner}/{repo}/issues"
     encoded_query_params = urllib.parse.urlencode(query_params)
     issues_url = f"{base_url}?{encoded_query_params}"
@@ -108,3 +107,20 @@ def gitea_issues_body_updated_at_get(issues, verbose=True):
         all_results = [future.result() for future in as_completed(futures)]
 
     return all_results
+
+
+def gitea_wiki_page_get(owner, repo, page_name, verbose=True):
+    """
+    Get a wiki page.
+    """
+    encoded_page_name = urllib.parse.quote(page_name, safe='')
+    base_url = f"{BASE_API_URL}/repos/{owner}/{repo}/wiki/page/{encoded_page_name}"
+    return url_json_get(base_url)
+
+
+def gitea_wiki_pages_get(owner, repo, verbose=True):
+    """
+    Get all wiki pages.
+    """
+    base_url = f"{BASE_API_URL}/repos/{owner}/{repo}/wiki/pages"
+    return url_json_get_all_pages(base_url)