update tools view and indexing

src/core/database.py

@@ -2,7 +2,6 @@ import os
 import json
 import time
 from src.utils import hashstr, setup_logger, is_text_pdf
-from src.core.filereader import pdfreader, plainreader
 from src.models.embedding import get_embedding_model
 
 logger = setup_logger("DataBaseManager")
@@ -95,18 +94,16 @@ class DataBaseManager:
         self._save_databases()
         return self.get_databases()
 
-    def add_files(self, db_id, files):
+    def add_files(self, db_id, files, params=None):
         db = self.get_kb_by_id(db_id)
 
         if db.embed_model != self.config.embed_model:
             logger.error(f"Embed model not match, {db.embed_model} != {self.config.embed_model}")
             return {"message": f"Embed model not match, cur: {self.config.embed_model}", "status": "failed"}
 
+        # Preprocessing the files to the queue
         new_files = []
         for file in files:
-            # filenames = [f["filename"] for f in db.files]
-            # if os.path.basename(file) in filenames:
-            #     continue
             new_file = {
                 "file_id": "file_" + hashstr(file + str(time.time())),
                 "filename": os.path.basename(file),
@@ -118,25 +115,30 @@ class DataBaseManager:
             db.files.append(new_file)
             new_files.append(new_file)
 
+        from src.core.indexing import chunk_file, chunk_text
         for new_file in new_files:
             file_id = new_file["file_id"]
-            idx = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0]
+            idx = self.get_idx_by_fileid(db, file_id)
             db.files[idx]["status"] = "processing"
 
             try:
-                text = self.read_text(new_file["path"])
+                if new_file["type"] in ["txt", "docx", "md"]:
-                chunks = self.chunking(text)
+                    nodes = chunk_file(new_file["path"], params=params)
+                else:
+                    text = self.read_text(new_file["path"])
+                    nodes = chunk_text(text, params)
+
                 self.knowledge_base.add_documents(
-                    docs=chunks,
+                    docs=[node.to_dict() for node in nodes],
                     collection_name=db.metaname,
                     file_id=file_id)
 
-                idx = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0]
+                idx = self.get_idx_by_fileid(db, file_id)
                 db.files[idx]["status"] = "done"
 
             except Exception as e:
                 logger.error(f"Failed to add documents to collection {db.metaname}, {e}")
-                idx = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0]
+                idx = self.get_idx_by_fileid(db, file_id)
                 db.files[idx]["status"] = "failed"
 
             self._save_databases()
@@ -151,7 +153,7 @@ class DataBaseManager:
             db.update(self.knowledge_base.get_collection_info(db.metaname))
             return db.to_dict()
 
-    def read_text(self, file):
+    def read_text(self, file, params=None):
         support_format = [".pdf", ".txt", ".md"]
         assert os.path.exists(file), "File not found"
         logger.info(f"Try to read file {file}")
@@ -162,12 +164,14 @@ class DataBaseManager:
 
         if file.endswith(".pdf"):
             if is_text_pdf(file):
+                from src.core.filereader import pdfreader
                 return pdfreader(file)
             else:
                 from src.plugins import pdf2txt
                 return pdf2txt(file, return_text=True)
 
         elif file.endswith(".txt") or file.endswith(".md"):
+            from src.core.filereader import plainreader
             return plainreader(file)
 
         else:
@@ -176,7 +180,7 @@ class DataBaseManager:
 
     def delete_file(self, db_id, file_id):
         db = self.get_kb_by_id(db_id)
-        file_idx_to_delete = [idx for idx, f in enumerate(db.files) if f["file_id"] == file_id][0]
+        file_idx_to_delete = self.get_idx_by_fileid(db, file_id)
 
         self.knowledge_base.client.delete(
             collection_name=db.metaname,
@@ -196,7 +200,10 @@ class DataBaseManager:
         )
         return {"lines": lines}
 
-    def chunking(self, text, chunk_size=1024):
+    def chunking(self, text, params=None):
+        chunk_method = params.get("chunk_method", "fixed")
+        chunk_size = params.get("chunk_size", 500)
+
         """将文本切分成固定大小的块"""
         chunks = []
         for i in range(0, len(text), chunk_size):
@@ -219,6 +226,11 @@ class DataBaseManager:
                 return db
         return None
 
+    def get_idx_by_fileid(self, db, file_id):
+        for idx, f in enumerate(db.files):
+            if f["file_id"] == file_id:
+                return idx
+
 
 class DataBaseLite:
     def __init__(self, name, description, db_type, dimension=None, **kwargs) -> None:

src/core/indexing.py

@@ -0,0 +1,37 @@
+import os
+from pathlib import Path
+from llama_index.core.node_parser import SimpleFileNodeParser
+from llama_index.readers.file import FlatReader
+
+from src.utils import hashstr
+
+def chunk_text(text, params=None):
+    params = params or {}
+    from llama_index.core import Document
+    from llama_index.core.node_parser import SentenceSplitter
+    chunk_size = int(params.get("chunk_size", 500))
+    chunk_overlap = int(params.get("chunk_overlap", 20))
+    splitter = SentenceSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+    )
+    doc = Document(id_=hashstr(text), text=text)
+    nodes = splitter.get_nodes_from_documents([doc])
+    return nodes
+
+def chunk_file(file, params=None):
+    parser = SimpleFileNodeParser()
+    if file.endswith(".txt"):
+        from llama_index.readers.file import FlatReader
+        docs = FlatReader().load_data(Path(file))
+    elif file.endswith(".docx"):
+        from llama_index.readers.file import DocxReader
+        docs = DocxReader().load_data(Path(file))
+    elif file.endswith(".md"):
+        from llama_index.readers.file import MarkdownReader
+        docs = MarkdownReader().load_data(Path(file))
+    else:
+        raise ValueError("Unsupported file type")
+
+    nodes = parser.get_nodes_from_documents(docs)
+    return nodes  # 返回节点

src/core/retriever.py

@@ -45,9 +45,10 @@ class Retriever:
             external_parts.extend(["图数据库信息:", db_text])
 
         # 构造查询
+        from src.utils.prompts import knowbase_qa_template
         if external_parts and len(external_parts) > 0:
             external = "\n\n".join(external_parts)
-            query = f"参考资料：\n\n\n{external}\n\n\n请根据前面的知识回答问题。\n\n问题：{query}\n\n回答："
+            query = knowbase_qa_template.format(external=external, query=query)
 
         return query
 

src/utils/prompts.py

@@ -1,3 +1,19 @@
+system_prompt = """
+"""
+
+
+knowbase_qa_template = """
+请利用查询到的资料回答问题，回答问题时，不要过度的分点作答。如果非要分点作答，可以使用 一、二、等：
+
+<参考资料>：
+{external}
+</参考资料>
+
+<问题>
+{query}
+</问题>"
+"""
+
 rewritten_query_prompt_template = """
 <指令>根据提供的历史信息对问题进行优化和改写，返回的问题必须符合以下内容要求和格式要求。严格不能出现禁止内容<指令>
 <禁止>1.绝对不能自己编造无关内容,若不能改写或无需改写直接返回原本问题

src/views/__init__.py

@@ -2,6 +2,7 @@ from flask import Flask
 from flask_cors import CORS
 from src.views.common_view import common
 from src.views.database_view import db
+from src.views.tools_view import tools
 
 
 def create_app():
@@ -10,5 +11,6 @@ def create_app():
 
     app.register_blueprint(common)
     app.register_blueprint(db)
+    app.register_blueprint(tools)
 
     return app

src/views/tools_view.py

@@ -0,0 +1,45 @@
+import os
+import json
+import threading
+from flask import Blueprint, jsonify, request, Response
+
+from src.utils import setup_logger, hashstr
+from src.core.startup import startup
+
+tools = Blueprint('tools', __name__, url_prefix="/tools")
+
+logger = setup_logger("server-tools")
+
+
+@tools.route("/", methods=["GET"])
+def route_index():
+    tools = [
+        {
+            "name": "text_chunking",
+            "title": "Text Chunking",
+            "description": "Chunking text into smaller pieces for better understanding.",
+            "url": "/tools/text_chunking",
+            "method": "POST",
+            "params": [
+                {
+                    "name": "text",
+                    "type": "string",
+                    "description": "Text to be chunked."
+                },
+                {
+                    "name": "chunk_size",
+                    "type": "int",
+                }
+            ]
+        }
+    ]
+
+    return jsonify(tools)
+
+
+@tools.route("/text_chunking", methods=["POST"])
+def text_chunking():
+    from src.core.indexing import chunk_text
+    text = request.json.get("text")
+    nodes = chunk_text(text, params=request.json)
+    return jsonify({"nodes": [node.to_dict() for node in nodes]})

web/src/assets/base.css

@@ -1,5 +1,7 @@
 /* color palette from <https://github.com/vuejs/theme> */
+/* https://material-foundation.github.io/material-theme-builder/ */
 :root {
+  --main-1000: #002A36;
   --main-900: #003A51;
   --main-800: #004F69;
   --main-700: #00637F;
@@ -12,26 +14,22 @@
   --main-50: #CDF5FF;
   --main-25: #E6FAFF;
   --main-10: #F5FDFF;
+  --main-5: #FAFCFD;
 
-  --c-white: #ffffff;
+  --gray-2000: #0C1214;
-  --c-white-soft: #f8f8f8;
+  --gray-1000: #171C1F;
-  --c-white-mute: #f2f2f2;
+  --gray-900: #212729;
+  --gray-800: #42484A;
+  --gray-700: #616161;
+  --gray-600: #8C9194;
+  --gray-500: #A7ACAF;
+  --gray-400: #C2C7CA;
+  --gray-300: #DEE3E6;
+  --gray-200: #EDF1F4;
+  --gray-100: #F5FAFD;
+  --gray-50: #F9FDFF;
 
-  --c-black: #202428;
+  --main-color: #1c6586;
-  --c-black-soft: #222222;
-  --c-black-mute: #282828;
-
-  --c-black-light-1: #333333;
-  --c-black-light-2: #454545;
-  --c-black-light-3: #666666;
-  --c-black-light-4: #999999;
-
-  --c-text-light-1: var(--c-black);
-  --c-text-dark-1: #0D0D0D;
-  --c-text-dark-2: #b8b8b8;
-  --color-text: var(--c-black);
-
-  --main-color: var(--main-700);
   --main-color-dark: #004d5c;
   --main-light-1: #0076AB;
   --main-light-2: #DAEAED;
@@ -39,8 +37,13 @@
   --main-light-4: #F2F5F5;
   --main-light-5: #F7FAFB;
   --main-light-6: #FAFDFD;
+  --secondry-color: #4e616d;
+  --error-color: #ba1a1a;
+
+  --bg-sider: var(--main-5);
+  --color-text: var(--c-black);
+
   --min-width: 400px;
-  --error-color: #f50a0d;
 }
 
 *,
@@ -55,7 +58,7 @@
 body {
   display: flow-root;
   min-height: 100vh;
-  color: var(--color-text);
+  color: var(--gray-900);
   line-height: 1.6;
   font-family: 'Roboto', 'Noto Sans SC', 'HarmonyOS Sans SC', -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Helvetica Neue', Arial, sans-serif;
   font-size: 15px;

web/src/components/TextChunkingComponent.vue

@@ -0,0 +1,211 @@
+<template>
+  <div class="text-chunking-container">
+    <div class="sidebar">
+      <div class="additional-params">
+        <h4>相关参数</h4>
+        <a-form
+          :model="params"
+          name="basic"
+          autocomplete="off"
+          layout="vertical"
+          @finish="chunkText"
+        >
+          <a-form-item label="Chunk Size" name="chunkSize" >
+            <a-input v-model:value="params.chunkSize" />
+          </a-form-item>
+          <a-form-item label="Chunk Overlap" name="chunkOverlap" >
+            <a-input v-model:value="params.chunkOverlap" />
+          </a-form-item>
+
+          <a-form-item>
+            <a-button type="primary" html-type="submit">Chunk Text</a-button>
+          </a-form-item>
+        </a-form>
+        <!-- Future parameters can be added here -->
+      </div>
+    </div>
+    <div class="result-container">
+      <div class="input-container">
+        <textarea v-model="text" placeholder="Enter text to chunk"></textarea>
+        <div class="infos">
+          <!-- <span>字数: {{ wordCount }}</span> -->
+          <span>字符数: {{ charCount }}</span>
+          <span>Token 数：{{ estimatedTokenCount }}</span>
+        </div>
+      </div>
+      <div id="result-cards" class="result-cards">
+        <div v-for="(chunk, index) in chunks" :key="index" class="chunk">
+          <p>{{ chunk }}</p>
+        </div>
+      </div>
+    </div>
+  </div>
+</template>
+
+<script setup>
+import { reactive, ref, computed } from 'vue';
+import { message } from 'ant-design-vue'
+
+const text = ref('');
+const state = reactive({
+  loading: true,
+})
+
+const params = reactive({
+  chunkSize: 500,
+  chunkOverlap: 20
+})
+const chunks = ref([]);
+
+const wordCount = computed(() => text.value.split(/\s+/).filter(Boolean).length);
+const charCount = computed(() => text.value.length)
+const estimatedTokenCount = computed(() => {
+  // 将文本分割成字符数组
+  const chars = text.value.split('');
+  let tokenCount = 0;
+  for (let char of chars) {
+    if (/[\u4e00-\u9fff]/.test(char)) {
+      tokenCount += 1; // 对于中文字符，通常计为一个 token
+    }
+    else if (/[a-zA-Z]/.test(char)) {
+      tokenCount += 0.25; // 对于英文单词，大约每 4 个字符算作一个 token
+    }
+    else {
+      tokenCount += 0.5; // 对于中文字符，通常计为一个 token
+    }
+  }
+  return Math.ceil(tokenCount);
+})
+
+const chunkText = async () => {
+  if (text.value.length === 0) {
+    message.error("请输入文本")
+    return
+  }
+  try {
+    state.loading = true
+    const response = await fetch('/api/tools/text_chunking', {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({
+        text: text.value,
+        chunk_size: params.chunkSize,
+        chunk_overlap: params.chunkOverlap
+      })
+    });
+
+    if (!response.ok) {
+      throw new Error(`HTTP error! Status: ${response.status}`);
+    }
+
+    const data = await response.json();
+    chunks.value = data.nodes.map(node => node.text);
+    state.loading = false
+  } catch (error) {
+    console.error('Error chunking text:', error);
+    state.loading = false
+  }
+};
+</script>
+
+<style lang="less" scoped>
+.text-chunking-container {
+  display: flex;
+  border-radius: 8px;
+  font-family: 'Arial', sans-serif;
+
+  .sidebar {
+    position: sticky;
+    top: 0;
+    width: 350px; /* 初始宽度 */
+    background-color: var(--bg-sider);
+    border-right: 1px solid var(--main-light-3);
+    padding: 20px;
+    min-width: 250px;
+    flex: 1;
+
+    .additional-params {
+      h4 {
+        font-size: 1.2em;
+        margin-bottom: 10px;
+      }
+    }
+  }
+
+  .result-container {
+    flex: 3;
+    padding: 20px;
+
+    .input-container {
+      display: flex;
+      flex-direction: column;
+      margin-bottom: 15px;
+
+      textarea {
+        resize: vertical;
+        height: 300px;
+        padding: 1rem;
+        border: 1px solid var(--gray-300);
+        border-radius: 8px;
+        font-size: 1rem;
+        transition: border-color 0.3s;
+        background-color: var(--gray-50);
+
+        &:focus {
+          border-color: var(--main-color);
+          outline: none;
+        }
+      }
+
+      .infos {
+        padding: 10px; /* Padding for spacing */
+        margin-top: 10px; /* Space above the infos */
+        font-size: 1em; /* Font size */
+        color: var(--gray-800); /* Text color */
+        display: flex;
+        gap: 16px;
+      }
+
+    }
+
+    .result-cards {
+      column-count: 3; /* 设置列数 */
+      column-gap: 10px; /* 列之间的间距 */
+    }
+
+    .chunk {
+      background-color: #f9fcff;
+      border-radius: 4px;
+      padding: 10px;
+      margin-bottom: 10px;
+      box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
+      break-inside: avoid; /* 避免元素被分割到不同列 */
+    }
+  }
+}
+
+
+@media (max-width: 768px) {
+  #result-cards {
+    column-count: 1;
+  }
+}
+
+@media (min-width: 769px) and (max-width: 1200px) {
+  #result-cards {
+    column-count: 2;
+  }
+}
+
+@media (min-width: 1201px) and (max-width: 1900px) {
+  #result-cards {
+    column-count: 3;
+  }
+}
+
+@media (min-width: 1901px){
+  #result-cards {
+    column-count: 4;
+  }
+}
+</style>

web/src/views/ToolsView.vue

@@ -7,12 +7,14 @@
     </HeaderComponent>
 
     <div class="tools-grid">
-      <div v-for="tool in tools" :key="tool.id" class="tool-card" @click="navigateToTool(tool.id)">
+      <div v-for="tool in tools" :key="tool.id" class="tool-card" @click="navigateToTool(tool.name)">
-        <div class="tool-icon">
+        <div class="tool-header">
-          <component :is="tool.icon" />
+          <div class="tool-icon">
+            <component :is="iconMap[tool.name]" />
+          </div>
+          <h3>{{ tool.title }}</h3>
         </div>
         <div class="tool-info">
-          <h3>{{ tool.name }}</h3>
           <p>{{ tool.description }}</p>
         </div>
       </div>
@@ -21,39 +23,42 @@
 </template>
 
 <script setup>
-import { ref } from 'vue';
+import { onMounted, reactive, ref } from 'vue';
 import { useRouter } from 'vue-router';
 import { CalculatorOutlined, FileSearchOutlined, TranslationOutlined } from '@ant-design/icons-vue';
 import HeaderComponent from '@/components/HeaderComponent.vue';
-import { Button as AButton } from 'ant-design-vue';
 
 const router = useRouter();
+const tools = ref([]);
+const iconMap = ref({
+  "text_chunking": FileSearchOutlined
+})
 
-const tools = ref([
+const state = reactive({
-  {
+  loadingTools: true,
-    id: 'calculator',
+})
-    name: '计算器',
-    description: '进行各种数学计算',
-    icon: CalculatorOutlined,
-  },
-  {
-    id: 'file-analyzer',
-    name: '文件分析器',
-    description: '分析各种文件的内容和结构',
-    icon: FileSearchOutlined,
-  },
-  {
-    id: 'translator',
-    name: '翻译工具',
-    description: '在不同语言之间进行翻译',
-    icon: TranslationOutlined,
-  },
-  // 可以继续添加更多工具...
-]);
 
-const navigateToTool = (toolId) => {
+const getTools = () => {
-  router.push({ name: 'ToolsComponent', params: { id: toolId } });
+  state.loadingTools = true
+  fetch('/api/tools/')
+    .then(response => response.json())
+    .then(data => {
+      tools.value = data;
+      state.loadingTools = false;
+    })
+    .catch(error => {
+      console.error('Error fetching tools:', error);
+      state.loadingTools = false;
+    });
 };
+
+const navigateToTool = (toolName) => {
+  router.push(`/tools/${toolName}`);
+};
+
+onMounted(() => {
+  getTools();
+});
 </script>
 
 <style scoped lang="less">
@@ -65,38 +70,44 @@ const navigateToTool = (toolId) => {
   display: grid;
   grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
   gap: 20px;
-  padding: 24px;
-}
-
-.tool-card {
-  background-color: #fff;
-  border-radius: 8px;
   padding: 20px;
-  box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
-  cursor: pointer;
-  transition: all 0.3s ease;
 
-  &:hover {
+  .tool-card {
-    transform: translateY(-5px);
+    display: flex;
-    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
+    flex-direction: column;
-  }
+    background-color: white;
+    border: 1px solid var(--gray-300);
+    border-radius: 8px;
+    padding: 20px;
+    transition: transform 0.3s ease, box-shadow 0.3s ease;
+    cursor: pointer;
 
-  .tool-icon {
+    &:hover {
-    font-size: 24px;
+      transform: translateY(-3px);
-    margin-bottom: 10px;
+      box-shadow: 0 5px 15px rgba(0,0,0,0.1);
-    color: #1890ff;
-  }
-
-  .tool-info {
-    h3 {
-      margin: 0 0 10px;
-      font-size: 18px;
     }
 
-    p {
+    .tool-header {
-      margin: 0;
+      display: flex;
-      color: #666;
+      align-items: center;
-      font-size: 14px;
+      margin-bottom: 15px;
+      font-size: 1rem;
+
+      .tool-icon {
+        margin-right: 10px;
+      }
+
+      h3 {
+        margin: 0;
+      }
+    }
+
+    .tool-info {
+      flex-grow: 1;
+
+      p {
+        margin: 0;
+      }
     }
   }
 }