QueryRewrite/rag2_0/tool/WikijsTool.py

import os.path

import requests
import json
import time
from pathlib import Path


class WikijsTool:
    BASE_URL = "http://10.1.16.39:8090/graphql"
    HEADERS = {
        "Authorization": "Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJhcGkiOjcsImdycCI6MSwiaWF"
                         "0IjoxNzIzMDIwNzg4LCJleHAiOjE4MTc2OTM1ODgsImF1ZCI6InVybjp3aWtpLmpzIiwiaX"
                         "NzIjoidXJuOndpa2kuanMifQ.NSfE4tB7tkN8yapAs0CgkR-Yll6wc3gO3QGKMAv-TlGxx6A-9fJRmkwhRDTVMj_yPVG6"
                         "NXVy_AZpJtLapRXFGn0cvscsRJxq3fY1KgEyt8wO99jvd8DpNHpHhAIgrtyDelmHsBD2Wb5Ib3WJFsWC6d8Yhm9dkpx6tZ"
                         "vMAlFIKOg6UodMoMIry3YWiPGLaqJPQ0gcKmcnB2tC7sPXIIZnvfb5912GVM0n-4wvWobQnb_tXQuYZf99wH_leXjC_7BK8"
                         "8JSaAmB980i3rBxfejmaJ8E6D48zRxwwPFa0veVjjzRkVqHPwAjl1CXb2HE29pGtNmSEE1kLQVqOZD_ibOwKQ"
    }

    def __init__(self):
        pass

    @staticmethod
    def init_url():
        # 获取当前文件的路径
        file_path = Path(__file__).resolve()
        file_path = os.path.join(file_path.parent, 'wikiconfig.json')
        if not os.path.exists(file_path):
            return False
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)

        if 'url' in data:
            WikijsTool.BASE_URL = data['url']

        if 'Authorization' in data:
            WikijsTool.HEADERS['Authorization'] = data['Authorization']

        return True

    @staticmethod
    def get_all_documents() -> list[dict]:
        query = """
        query Pages {
          pages {
            list {
              path
              locale
              title
              contentType
              id
              isPublished
            }
          }
        }
        """
        # 构建请求数据
        data = {
            'query': query,
        }

        # 发送 POST 请求
        response = requests.post(WikijsTool.BASE_URL, headers=WikijsTool.HEADERS, json=data)
        if response.status_code == 200:
            # 解析数据
            list_info = json.loads(response.content)['data']['pages']['list']
            return [item for item in list_info]
        else:
            raise ValueError(f"获取文档列表失败，原因：“{response.text}")

    @staticmethod
    def get_all_doc_by_path(path: str, path_is_dir: bool = True) -> list[dict]:
        list_document = WikijsTool.get_all_documents()
        all_document_list = []
        if path_is_dir:
            temp_path = path + '/'
        else:
            temp_path = path
        for document_info in list_document:
            document_path = str(document_info["path"])
            # 根据路径过滤出对应的所有文档
            if not document_path.startswith(temp_path):
                continue

            all_document_list.append(document_info)

        return all_document_list

    @staticmethod
    def search_document(query_str: str) -> list[dict]:
        graphql_query = f"""
        query Pages {{
            pages {{
                search(query: "{query_str}") {{
                    results {{
                        id
                        path
                        locale
                        title
                    }}
                }}
            }}
        }}
        """
        # 构建请求数据
        data = {
            'query': graphql_query,
        }

        # 发送 POST 请求
        response = requests.post(WikijsTool.BASE_URL, headers=WikijsTool.HEADERS, json=data)
        if response.status_code == 200:
            # 解析数据
            search_results = json.loads(response.content)['data']['pages']['search']['results']
            return search_results
        else:
            raise ValueError(f"查询文档失败，原因：“{response.text}")

    @staticmethod
    def query_doc_info(doc_id: int) -> dict:
        query = """
        query singlePages($doc_id: Int!) {
            pages {
                single(id: $doc_id) {
                    id
                    path
                    title
                    isPublished
                    content
                    contentType
                    isPrivate
                    updatedAt
                    createdAt
                }
            }
        }
        """
        # 构建请求数据
        variables = {
            'doc_id': doc_id,
        }
        data = {
            'query': query,
            'variables': variables
        }

        # 发送 POST 请求
        response = requests.post(WikijsTool.BASE_URL, headers=WikijsTool.HEADERS, json=data)
        if "errors" in response.text:
            result = json.loads(response.content)['errors'][0]['message']
            raise ValueError(f"查询文档失败，原因：“{result}")
        else:
            return json.loads(response.content)['data']['pages']['single']


WikijsTool.init_url()
if __name__ == "__main__":
    WikijsTool.query_doc_info(6448)
    print(WikijsTool.rename_directory("配网知识库/配网造价软件", "配网知识库/配网造价软件1"))