kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

update tools view and indexing

Browse Source
This commit is contained in:
Wenjie Zhang 2024-09-25 13:46:23 +08:00
parent 012f77aafc
commit 6f81b061d2
9 changed files with 429 additions and 91 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
src/core/database.py
View File

@ -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"])
chunks = self.chunking(text)
if new_file["type"] in ["txt", "docx", "md"]:
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:

37
src/core/indexing.py Normal file
View File

@ -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 # 返回节点

3
src/core/retriever.py
View File

@ -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

16
src/utils/prompts.py
View File

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

2
src/views/__init__.py
View File

@ -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

45
src/views/tools_view.py Normal file
View File

@ -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]})

43
web/src/assets/base.css
View File

@ -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;
--c-white-soft: #f8f8f8;
--c-white-mute: #f2f2f2;
--gray-2000: #0C1214;
--gray-1000: #171C1F;
--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;
--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: #1c6586;
--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;

211
web/src/components/TextChunkingComponent.vue Normal file
View File

@ -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>

123
web/src/views/ToolsView.vue
View File

@ -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 class="tool-icon">
<component :is="tool.icon" />
<div v-for="tool in tools" :key="tool.id" class="tool-card" @click="navigateToTool(tool.name)">
<div class="tool-header">
<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([
{
id: 'calculator',
name: '计算器',
description: '进行各种数学计算',
icon: CalculatorOutlined,
},
{
id: 'file-analyzer',
name: '文件分析器',
description: '分析各种文件的内容和结构',
icon: FileSearchOutlined,
},
{
id: 'translator',
name: '翻译工具',
description: '在不同语言之间进行翻译',
icon: TranslationOutlined,
},
// 可以继续添加更多工具...
]);
const state = reactive({
loadingTools: true,
})
const navigateToTool = (toolId) => {
router.push({ name: 'ToolsComponent', params: { id: toolId } });
const getTools = () => {
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 {
transform: translateY(-5px);
box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
}
.tool-card {
display: flex;
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 {
font-size: 24px;
margin-bottom: 10px;
color: #1890ff;
}
.tool-info {
h3 {
margin: 0 0 10px;
font-size: 18px;
&:hover {
transform: translateY(-3px);
box-shadow: 0 5px 15px rgba(0,0,0,0.1);
}
p {
margin: 0;
color: #666;
font-size: 14px;
.tool-header {
display: flex;
align-items: center;
margin-bottom: 15px;
font-size: 1rem;
.tool-icon {
margin-right: 10px;
}
h3 {
margin: 0;
}
}
.tool-info {
flex-grow: 1;
p {
margin: 0;
}
}
}
}