Wenjie Zhang
GitHub
commit
ad2e6e9f92
3
.gitignore
vendored
3
.gitignore
vendored
@ -40,4 +40,5 @@ local_neo4j/data
|
||||
local_neo4j/logs
|
||||
local_neo4j/import
|
||||
local_neo4j/plugins
|
||||
local_neo4j/conf
|
||||
local_neo4j/conf
|
||||
graphrag
|
||||
@ -1,15 +1,17 @@
|
||||
from neo4j import GraphDatabase
|
||||
from neo4j.exceptions import ServiceUnavailable, AuthError
|
||||
|
||||
# sudo ln -s /snap/core22/1586/usr/sbin/iptables /usr/sbin/iptables
|
||||
|
||||
def check_neo4j_status(uri="bolt://localhost:7687", username="neo4j", password="0123456789"):
|
||||
"""
|
||||
检查 Neo4j 数据库是否可以连接并正常工作。
|
||||
|
||||
|
||||
参数:
|
||||
uri (str): Neo4j 的 URI，默认为 "bolt://localhost:7687"
|
||||
username (str): 数据库用户名，默认为 "neo4j"
|
||||
password (str): 数据库密码，默认为 "0123456789"
|
||||
|
||||
|
||||
返回:
|
||||
str: "OK" 表示连接成功，"UNAVAILABLE" 表示服务不可用，"AUTH_FAILED" 表示认证失败。
|
||||
"""
|
||||
|
||||
@ -49,7 +49,7 @@ class Config(SimpleConfig):
|
||||
# 模型配置
|
||||
## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径
|
||||
## 如果需要自定义路径，则在 config/base.yaml 中配置 model_local_paths
|
||||
self.add_item("model_provider", default="zhipu", des="模型提供商", choices=["qianfan", "vllm", "zhipu", "deepseek", "dashscope"])
|
||||
self.add_item("model_provider", default="zhipu", des="模型提供商", choices=["openai", "qianfan", "vllm", "zhipu", "deepseek", "dashscope"])
|
||||
self.add_item("model_name", default=None, des="模型名称")
|
||||
self.add_item("embed_model", default="zhipu-embedding-3", des="Embedding 模型", choices=list(EMBED_MODEL_INFO.keys()))
|
||||
self.add_item("reranker", default="bge-reranker-v2-m3", des="Re-Ranker 模型", choices=["bge-reranker-v2-m3"])
|
||||
@ -145,6 +145,13 @@ class Config(SimpleConfig):
|
||||
|
||||
MODEL_NAMES = {
|
||||
# https://platform.deepseek.com/api-docs/zh-cn/pricing
|
||||
"openai": [
|
||||
"gpt-4",
|
||||
"gpt-4o",
|
||||
"gpt-4-0125-preview",
|
||||
"gpt-4o-mini",
|
||||
],
|
||||
|
||||
"deepseek": [
|
||||
"deepseek-chat",
|
||||
"deepseek-coder"
|
||||
|
||||
@ -15,6 +15,71 @@ logger = setup_logger("server-graphbase")
|
||||
|
||||
warnings.filterwarnings("ignore", category=UserWarning)
|
||||
|
||||
|
||||
"""
|
||||
from neo4j import GraphDatabase
|
||||
import random
|
||||
|
||||
class KnowledgeGraph:
|
||||
def __init__(self, uri, user, password):
|
||||
self._driver = GraphDatabase.driver(uri, auth=(user, password))
|
||||
|
||||
def close(self):
|
||||
self._driver.close()
|
||||
|
||||
def use_database(self, kgdb_name):
|
||||
with self._driver.session() as session:
|
||||
session.run(f"USE {kgdb_name}")
|
||||
|
||||
def get_sample_nodes(self, kgdb_name='neo4j', num=50):
|
||||
self.use_database(kgdb_name)
|
||||
selected_nodes = set()
|
||||
nodes_to_expand = set()
|
||||
result_nodes = []
|
||||
|
||||
with self._driver.session() as session:
|
||||
while len(selected_nodes) < num:
|
||||
# 如果需要扩展的节点为空，随机选择一个新节点
|
||||
if not nodes_to_expand:
|
||||
result = session.run("MATCH (n) RETURN n, rand() as r ORDER BY r LIMIT 1")
|
||||
for record in result:
|
||||
nodes_to_expand.add(record['n'].id)
|
||||
result_nodes.append({'n': record['n'], 'r': None, 'm': None})
|
||||
|
||||
# 从需要扩展的节点中随机选择一个节点
|
||||
current_node_id = random.choice(list(nodes_to_expand))
|
||||
nodes_to_expand.remove(current_node_id)
|
||||
|
||||
# 获取当前节点的邻居节点，最多5个
|
||||
result = session.run(
|
||||
f"MATCH (n)-[r]-(m) WHERE id(n) = {current_node_id} RETURN n, r, m LIMIT 5"
|
||||
)
|
||||
|
||||
for record in result:
|
||||
neighbor_node_id = record['m'].id
|
||||
if neighbor_node_id not in selected_nodes:
|
||||
selected_nodes.add(neighbor_node_id)
|
||||
nodes_to_expand.add(neighbor_node_id)
|
||||
result_nodes.append({'n': record['n'], 'r': record['r'], 'm': record['m']})
|
||||
|
||||
# 如果已经达到最大值，停止扩展
|
||||
if len(selected_nodes) >= num:
|
||||
break
|
||||
|
||||
return result_nodes[:num]
|
||||
|
||||
# 示例用法
|
||||
uri = "bolt://localhost:7687"
|
||||
user = "neo4j"
|
||||
password = "password"
|
||||
kg = KnowledgeGraph(uri, user, password)
|
||||
sample_nodes = kg.get_sample_nodes(num=50)
|
||||
for node in sample_nodes:
|
||||
print(f"Node: {node['n'].id}, Relationship: {node['r']}, Neighbor: {node['m'].id}")
|
||||
kg.close()
|
||||
|
||||
"""
|
||||
|
||||
UIE_MODEL = None
|
||||
|
||||
class GraphDatabase:
|
||||
@ -38,7 +103,7 @@ class GraphDatabase:
|
||||
self.driver.close()
|
||||
|
||||
def get_sample_nodes(self, kgdb_name='neo4j', num=50):
|
||||
"""获取指定数据库的前 num 个节点信息"""
|
||||
"""获取指定数据库的 num 个节点信息"""
|
||||
self.use_database(kgdb_name)
|
||||
def query(tx, num):
|
||||
result = tx.run("MATCH (n)-[r]->(m) RETURN n, r, m LIMIT $num", num=int(num))
|
||||
|
||||
@ -38,5 +38,6 @@ class HistoryManager():
|
||||
def __str__(self):
|
||||
history_str = ""
|
||||
for message in self.messages:
|
||||
history_str += f"{message['role']}: {message['content']}"
|
||||
msg = message["content"].replace('\n', ' ')
|
||||
history_str += f"\n{message['role']}: {msg}"
|
||||
return history_str
|
||||
@ -63,7 +63,7 @@ class KnowledgeBase:
|
||||
"id": int(random.random() * 1e12),
|
||||
"vector": vectors[i],
|
||||
"text": docs[i],
|
||||
"hash": hashstr(docs[i] + str(random.random())),
|
||||
"hash": hashstr(docs[i], with_salt=True),
|
||||
**kwargs} for i in range(len(vectors))]
|
||||
|
||||
res = self.client.insert(collection_name=collection_name, data=data)
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
from src.models.embedding import Reranker
|
||||
from src.utils.logging_config import setup_logger
|
||||
|
||||
logger = setup_logger("server-common")
|
||||
|
||||
|
||||
@ -37,12 +38,14 @@ class Retriever:
|
||||
# 解析图数据库的结果
|
||||
db_res = refs.get("graph_base").get("results", [])
|
||||
if len(db_res["nodes"]) > 0:
|
||||
db_text = '\n'.join([f"{edge['source_name']}和{edge['target_name']}的关系是{edge['type']}" for edge in db_res['edges']])
|
||||
db_text = "\n".join(
|
||||
[f"{edge['source_name']}和{edge['target_name']}的关系是{edge['type']}" for edge in db_res["edges"]]
|
||||
)
|
||||
external += f"图数据库信息: \n\n{db_text}"
|
||||
|
||||
# 构造查询
|
||||
if len(external) > 0:
|
||||
query = f"以下是参考资料：\n\n\n{external}\n\n\n请根据前面的知识回答：{query}"
|
||||
query = f"参考资料：\n\n\n{external}\n\n\n请根据前面的知识回答问题。\n\n问题：{query}\n\n回答："
|
||||
|
||||
return query
|
||||
|
||||
@ -70,42 +73,53 @@ class Retriever:
|
||||
kb_res = []
|
||||
final_res = []
|
||||
if not refs["meta"].get("db_name") or not self.config.enable_knowledge_base:
|
||||
return {"results": final_res, "all_results": kb_res, "rw_query": query, "message": "Knowledge base is disabled"}
|
||||
return {
|
||||
"results": final_res,
|
||||
"all_results": kb_res,
|
||||
"rw_query": query,
|
||||
"message": "Knowledge base is disabled",
|
||||
}
|
||||
|
||||
rw_query = self.rewrite_query(query, history, refs)
|
||||
|
||||
db_name = refs["meta"]["db_name"]
|
||||
kb = self.dbm.metaname2db[db_name]
|
||||
limit = refs["meta"].get("queryCount", 10)
|
||||
|
||||
kb_res = self.dbm.knowledge_base.search(rw_query, db_name, limit=limit)
|
||||
for r in kb_res:
|
||||
max_query_count = refs["meta"].get("maxQueryCount", 10)
|
||||
rerank_threshold = refs["meta"].get("rerankThreshold", 0.1)
|
||||
distance_threshold = refs["meta"].get("distanceThreshold", 0)
|
||||
top_k = refs["meta"].get("topK", 5)
|
||||
|
||||
all_kb_res = self.dbm.knowledge_base.search(rw_query, db_name, limit=max_query_count)
|
||||
for r in all_kb_res:
|
||||
r["file"] = kb.id2file(r["entity"]["file_id"])
|
||||
|
||||
# use distance threshold to filter results
|
||||
kb_res = [r for r in all_kb_res if r["distance"] > distance_threshold]
|
||||
|
||||
if self.config.enable_reranker:
|
||||
RERANK_THRESHOLD = 0.001
|
||||
for r in kb_res:
|
||||
r["rerank_score"] = self.reranker.compute_score([query, r["entity"]["text"]], normalize=True)
|
||||
r["rerank_score"] = self.reranker.compute_score([rw_query, r["entity"]["text"]], normalize=True)
|
||||
kb_res.sort(key=lambda x: x["rerank_score"], reverse=True)
|
||||
final_res = [_res for _res in kb_res if _res["rerank_score"] > RERANK_THRESHOLD]
|
||||
kb_res = [_res for _res in kb_res if _res["rerank_score"] > rerank_threshold]
|
||||
|
||||
else:
|
||||
final_res = kb_res[:5]
|
||||
kb_res = kb_res[:top_k]
|
||||
|
||||
return {"results": final_res, "all_results": kb_res, "rw_query": rw_query}
|
||||
return {"results": kb_res, "all_results": all_kb_res, "rw_query": rw_query}
|
||||
|
||||
def rewrite_query(self, query, history, refs):
|
||||
"""重写查询"""
|
||||
rewrite_query_span = refs["meta"].get("rewrite_query", None)
|
||||
if rewrite_query_span is None or rewrite_query_span == "OFF":
|
||||
rewrite_query_span = refs["meta"].get("rewriteQuery", "off")
|
||||
if rewrite_query_span == "off":
|
||||
rewritten_query = query
|
||||
else:
|
||||
from src.utils.prompts import rewritten_query_prompt_template
|
||||
history_query = [entry['content'] for entry in history if entry['role'] == 'user'] if history else ""
|
||||
|
||||
history_query = [entry["content"] for entry in history if entry["role"] == "user"] if history else ""
|
||||
rewritten_query_prompt = rewritten_query_prompt_template.format(history=history_query, query=query)
|
||||
rewritten_query = self.model.predict(rewritten_query_prompt).content
|
||||
|
||||
if rewrite_query_span == "HyDE":
|
||||
if rewrite_query_span == "hyde":
|
||||
hy_doc = self.model.predict(rewritten_query).content
|
||||
rewritten_query = f"{rewritten_query} {hy_doc}"
|
||||
|
||||
@ -118,54 +132,68 @@ class Retriever:
|
||||
entities = []
|
||||
if refs["meta"].get("use_graph"):
|
||||
from src.utils.prompts import entity_extraction_prompt_template
|
||||
|
||||
entity_extraction_prompt = entity_extraction_prompt_template.format(text=query)
|
||||
entities = self.model.predict(entity_extraction_prompt).content.split(",")
|
||||
entities = [entity for entity in entities if all(char.isalnum() or char in '汉字' for char in entity)]
|
||||
entities = [entity for entity in entities if all(char.isalnum() or char in "汉字" for char in entity)]
|
||||
|
||||
return entities
|
||||
|
||||
def foramt_general_results(self, results):
|
||||
logger.debug(f"Formatting general results: {results}")
|
||||
def _extract_relationship_info(self, relationship, source_name, target_name):
|
||||
"""
|
||||
提取关系信息并返回格式化的节点和边信息
|
||||
"""
|
||||
rel_id = relationship.element_id
|
||||
nodes = relationship.nodes
|
||||
if len(nodes) != 2:
|
||||
return None, None
|
||||
|
||||
source, target = nodes
|
||||
source_id = source.element_id
|
||||
target_id = target.element_id
|
||||
|
||||
relationship_type = relationship._properties.get("type", "unknown")
|
||||
if relationship_type == "unknown":
|
||||
relationship_type = relationship.type
|
||||
|
||||
edge_info = {
|
||||
"id": rel_id,
|
||||
"type": relationship_type,
|
||||
"source_id": source_id,
|
||||
"target_id": target_id,
|
||||
"source_name": source_name,
|
||||
"target_name": target_name,
|
||||
}
|
||||
|
||||
node_info = [
|
||||
{"id": source_id, "name": source_name},
|
||||
{"id": target_id, "name": target_name},
|
||||
]
|
||||
|
||||
return node_info, edge_info
|
||||
|
||||
def format_general_results(self, results):
|
||||
formatted_results = {"nodes": [], "edges": []}
|
||||
|
||||
for item in results:
|
||||
relationship = item[1]
|
||||
rel_id = relationship.element_id
|
||||
nodes = relationship.nodes
|
||||
if len(nodes) != 2:
|
||||
source_name = item[0]._properties.get("name", "unknown")
|
||||
target_name = item[2]._properties.get("name", "unknown") if len(item) > 2 else "unknown"
|
||||
|
||||
node_info, edge_info = self._extract_relationship_info(relationship, source_name, target_name)
|
||||
if node_info is None or edge_info is None:
|
||||
continue
|
||||
|
||||
source, target = nodes
|
||||
for node in node_info:
|
||||
if node["id"] not in [n["id"] for n in formatted_results["nodes"]]:
|
||||
formatted_results["nodes"].append(node)
|
||||
|
||||
source_id = source.element_id
|
||||
target_id = target.element_id
|
||||
source_name = source._properties.get('name', 'unknown')
|
||||
target_name = target._properties.get('name', 'unknown')
|
||||
|
||||
if source_id not in formatted_results["nodes"]:
|
||||
formatted_results["nodes"].append({"id": source_id, "name": source_name})
|
||||
if target_id not in formatted_results["nodes"]:
|
||||
formatted_results["nodes"].append({"id": target_id, "name": target_name})
|
||||
|
||||
relationship_type = relationship._properties.get('type', 'unknown')
|
||||
if relationship_type == 'unknown':
|
||||
relationship_type = relationship.type
|
||||
|
||||
formatted_results["edges"].append({
|
||||
"id": rel_id,
|
||||
"type": relationship_type,
|
||||
"source_id": source_id,
|
||||
"target_id": target_id,
|
||||
"source_name": source_name,
|
||||
"target_name": target_name
|
||||
})
|
||||
formatted_results["edges"].append(edge_info)
|
||||
|
||||
return formatted_results
|
||||
|
||||
def format_query_results(self, results):
|
||||
logger.debug(f"Formatting query results: {results}")
|
||||
formatted_results = {"nodes": [], "edges": []}
|
||||
|
||||
node_dict = {}
|
||||
|
||||
for item in results:
|
||||
@ -173,35 +201,18 @@ class Retriever:
|
||||
continue
|
||||
|
||||
relationship = item[1][0]
|
||||
rel_id = relationship.element_id
|
||||
nodes = relationship.nodes
|
||||
if len(nodes) != 2:
|
||||
source_name = item[0]
|
||||
target_name = item[2] if len(item) > 2 else "unknown"
|
||||
|
||||
node_info, edge_info = self._extract_relationship_info(relationship, source_name, target_name)
|
||||
if node_info is None or edge_info is None:
|
||||
continue
|
||||
|
||||
source, target = nodes
|
||||
for node in node_info:
|
||||
if node["id"] not in node_dict:
|
||||
node_dict[node["id"]] = node
|
||||
|
||||
source_id = source.element_id
|
||||
target_id = target.element_id
|
||||
source_name = item[0]
|
||||
target_name = item[2] if len(item) > 2 else 'unknown'
|
||||
|
||||
if source_id not in node_dict:
|
||||
node_dict[source_id] = {"id": source_id, "name": source_name}
|
||||
if target_id not in node_dict:
|
||||
node_dict[target_id] = {"id": target_id, "name": target_name}
|
||||
|
||||
relationship_type = relationship._properties.get('type', 'unknown')
|
||||
if relationship_type == 'unknown':
|
||||
relationship_type = relationship.type
|
||||
|
||||
formatted_results["edges"].append({
|
||||
"id": rel_id,
|
||||
"type": relationship_type,
|
||||
"source_id": source_id,
|
||||
"target_id": target_id,
|
||||
"source_name": source_name,
|
||||
"target_name": target_name
|
||||
})
|
||||
formatted_results["edges"].append(edge_info)
|
||||
|
||||
formatted_results["nodes"] = list(node_dict.values())
|
||||
|
||||
@ -210,4 +221,4 @@ class Retriever:
|
||||
def __call__(self, query, history, meta):
|
||||
refs = self.retrieval(query, history, meta)
|
||||
query = self.construct_query(query, refs, meta)
|
||||
return query, refs
|
||||
return query, refs
|
||||
|
||||
@ -38,7 +38,7 @@ python -m vllm.entrypoints.openai.api_server \
|
||||
|`bge-large-zh-v1.5`|`BAAI/bge-large-zh-v1.5`|`bge-large-zh-v1.5`（*修改为本地路径）|
|
||||
|`zhipu`|`embedding-2`|`ZHIPUAPI` (`.env`)|
|
||||
|
||||
### 3. 重排序模型支持
|
||||
|
||||
|
||||
|
||||
例如（`config/base.yaml`）：
|
||||
@ -52,3 +52,7 @@ model_name: null # for default
|
||||
model_local_paths:
|
||||
bge-large-zh-v1.5: /models/bge-large-zh-v1.5
|
||||
```
|
||||
|
||||
### 3. 重排序模型支持
|
||||
|
||||
目前仅支持 `BAAI/bge-reranker-v2-m3`。
|
||||
|
||||
@ -28,6 +28,10 @@ def select_model(config):
|
||||
from src.models.chat_model import DashScope
|
||||
return DashScope(model_name)
|
||||
|
||||
elif model_provider == "openai":
|
||||
from src.models.chat_model import OpenModel
|
||||
return OpenModel(model_name)
|
||||
|
||||
elif model_provider is None:
|
||||
raise ValueError("Model provider not specified, please modify `model_provider` in `src/config/base.yaml`")
|
||||
else:
|
||||
|
||||
@ -39,6 +39,14 @@ class OpenAIBase():
|
||||
return response.choices[0].message
|
||||
|
||||
|
||||
class OpenModel(OpenAIBase):
|
||||
def __init__(self, model_name=None):
|
||||
model_name = model_name or "gpt-4o-mini"
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
base_url = os.getenv("OPENAI_API_BASE")
|
||||
super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)
|
||||
|
||||
|
||||
class DeepSeek(OpenAIBase):
|
||||
def __init__(self, model_name=None):
|
||||
model_name = model_name or "deepseek-chat"
|
||||
|
||||
@ -1,12 +1,16 @@
|
||||
import os
|
||||
import uuid
|
||||
from FlagEmbedding import FlagModel, FlagReranker
|
||||
|
||||
from src.config import EMBED_MODEL_INFO, RERANKER_LIST
|
||||
from src.utils.logging_config import setup_logger
|
||||
from src.utils import hashstr
|
||||
|
||||
|
||||
logger = setup_logger("EmbeddingModel")
|
||||
|
||||
GLOBAL_EMBED_STATE = {}
|
||||
|
||||
|
||||
class EmbeddingModel(FlagModel):
|
||||
def __init__(self, model_info, config, **kwargs):
|
||||
@ -47,9 +51,21 @@ class ZhipuEmbedding:
|
||||
def predict(self, message):
|
||||
data = []
|
||||
|
||||
if len(message) > 10:
|
||||
global GLOBAL_EMBED_STATE
|
||||
task_id = hashstr(message)
|
||||
logger.info(f"Creating new state for process {task_id}")
|
||||
GLOBAL_EMBED_STATE[task_id] = {
|
||||
'status': 'in-progress',
|
||||
'total': len(message),
|
||||
'progress': 0
|
||||
}
|
||||
|
||||
for i in range(0, len(message), 10):
|
||||
if len(message) > 10:
|
||||
logger.info(f"Encoding {i} to {i+10} with {len(message)} messages")
|
||||
GLOBAL_EMBED_STATE[task_id]['progress'] = i
|
||||
|
||||
group_msg = message[i:i+10]
|
||||
response = self.client.embeddings.create(
|
||||
model=self.model_info.default_path,
|
||||
@ -58,6 +74,10 @@ class ZhipuEmbedding:
|
||||
|
||||
data.extend([a.embedding for a in response.data])
|
||||
|
||||
if len(message) > 10:
|
||||
GLOBAL_EMBED_STATE[task_id]['progress'] = len(message)
|
||||
GLOBAL_EMBED_STATE[task_id]['status'] = 'completed'
|
||||
|
||||
return data
|
||||
|
||||
def encode(self, message):
|
||||
|
||||
@ -16,7 +16,7 @@ logger = setup_logger("OneKE")
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
MODEL_NAME_OR_PATH = os.path.join(os.getenv('MODEL_ROOT_DIR'), 'OneKE')
|
||||
MODEL_NAME_OR_PATH = os.path.join(os.getenv('MODEL_ROOT_DIR', './'), 'OneKE')
|
||||
|
||||
instruction_mapper = {
|
||||
'NERzh': "你是专门进行实体抽取的专家。请从input中抽取出符合schema定义的实体，不存在的实体类型返回空列表。请按照JSON字符串的格式回答。",
|
||||
|
||||
@ -1,25 +1,53 @@
|
||||
import os
|
||||
import uuid
|
||||
import fitz # fitz就是pip install PyMuPDF
|
||||
import cv2
|
||||
from copy import deepcopy
|
||||
from tqdm import tqdm
|
||||
from argparse import ArgumentParser
|
||||
from src.utils import logger
|
||||
|
||||
GOLBAL_STATE = {}
|
||||
|
||||
|
||||
def pdf2txt(pdf_path, return_text=False):
|
||||
if not os.path.exists(pdf_path):
|
||||
raise FileNotFoundError(f"File not found: {pdf_path}")
|
||||
|
||||
# Importing these modules here to avoid unnecessary imports in other files
|
||||
from paddleocr import PPStructure, save_structure_res
|
||||
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
|
||||
output_dir = os.path.join('tmp', 'pdf2txt', os.path.basename(pdf_path).split('.')[0])
|
||||
filename = os.path.basename(pdf_path).split('.')[0]
|
||||
output_dir = os.path.join('saves', 'data', 'pdf2txt', filename)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
table_engine = PPStructure(recovery=True, lang='ch')
|
||||
|
||||
imgs = convert_imgs(pdf_path, output_dir)
|
||||
respath = os.path.join(output_dir, 'res.txt')
|
||||
respath = os.path.join(output_dir, f'{filename}.txt')
|
||||
|
||||
global GOLBAL_STATE
|
||||
task_id = str(uuid.uuid4())
|
||||
if task_id in GOLBAL_STATE:
|
||||
logger.info(f"Reusing previous state for process {task_id}")
|
||||
return GOLBAL_STATE[task_id]
|
||||
else:
|
||||
logger.info(f"Creating new state for process {task_id}")
|
||||
GOLBAL_STATE[task_id] = {
|
||||
'pdf_path': pdf_path,
|
||||
'return_text': return_text,
|
||||
'output_dir': output_dir,
|
||||
'status': 'in-progress',
|
||||
'total': len(imgs),
|
||||
'progress': 0
|
||||
}
|
||||
|
||||
|
||||
text = []
|
||||
for img_name in tqdm(imgs, desc='to txt', ncols=100):
|
||||
img = cv2.imread(img_name)
|
||||
result = table_engine(img)
|
||||
GOLBAL_STATE[task_id]['progress'] += 1
|
||||
|
||||
save_structure_res(result, output_dir, "structure_result")
|
||||
|
||||
@ -43,6 +71,9 @@ def pdf2txt(pdf_path, return_text=False):
|
||||
whole_text = ''.join(text)
|
||||
with open(respath, 'w', encoding='utf-8') as f:
|
||||
f.write(whole_text)
|
||||
logger.info(f"Extracted text saved to {respath}")
|
||||
|
||||
GOLBAL_STATE[task_id]['status'] = 'completed'
|
||||
|
||||
if return_text:
|
||||
return whole_text
|
||||
@ -72,6 +103,13 @@ def convert_imgs(pdf_path, output_dir):
|
||||
|
||||
return imgs
|
||||
|
||||
def get_state(task_id):
|
||||
return GOLBAL_STATE.get(task_id, {})
|
||||
|
||||
if __name__ == "__main__":
|
||||
pdf_path = r'saves/data/uploads/2e04d5_保健食品.pdf'
|
||||
print(pdf2txt(pdf_path))
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument('--pdf-path', type=str, required=True, help='Path to the PDF file')
|
||||
parser.add_argument('--return-text', action='store_true', help='Return the extracted text')
|
||||
args = parser.parse_args()
|
||||
|
||||
pdf2txt(args.pdf_path, args.return_text)
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import time
|
||||
import random
|
||||
from src.utils.logging_config import setup_logger, logger
|
||||
|
||||
def is_text_pdf(pdf_path):
|
||||
@ -15,7 +16,7 @@ def hashstr(input_string, length=8, with_salt=False):
|
||||
import hashlib
|
||||
# 添加时间戳作为干扰
|
||||
if with_salt:
|
||||
input_string += str(time.time())
|
||||
input_string += str(time.time() + random.random())
|
||||
|
||||
hash = hashlib.md5(str(input_string).encode()).hexdigest()
|
||||
return hash[:length]
|
||||
@ -5,10 +5,10 @@ from datetime import datetime
|
||||
|
||||
DATETIME = datetime.now().strftime('%Y-%m-%d-%H%M%S')
|
||||
# DATETIME = "debug" # 为了方便，调试的时候输出到 debug.log 文件
|
||||
LOG_FILE = f'log/project-{DATETIME}.log'
|
||||
LOG_FILE = f'saves/log/project-{DATETIME}.log'
|
||||
|
||||
def setup_logger(name, level=logging.DEBUG, console=False):
|
||||
os.makedirs("log", exist_ok=True)
|
||||
os.makedirs("saves/log", exist_ok=True)
|
||||
|
||||
"""Function to setup logger with the given name and log file."""
|
||||
logger = logging.getLogger(name)
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
import itertools
|
||||
import json
|
||||
from flask import Blueprint, jsonify, request, Response
|
||||
|
||||
@ -89,6 +90,7 @@ def restart():
|
||||
def get_log():
|
||||
from src.utils.logging_config import LOG_FILE
|
||||
with open(LOG_FILE, 'r') as f:
|
||||
log = f.read()
|
||||
log_lines = itertools.islice(f, 1000)
|
||||
log = ''.join(log_lines)
|
||||
|
||||
return jsonify({"log": log})
|
||||
@ -135,7 +135,7 @@ def get_graph_nodes():
|
||||
|
||||
logger.debug(f"Get graph nodes in {kgdb_name} with {num} nodes")
|
||||
result = startup.dbm.graph_base.get_sample_nodes(kgdb_name, num)
|
||||
return jsonify({'result': startup.retriever.foramt_general_results(result), 'message': 'success'}), 200
|
||||
return jsonify({'result': startup.retriever.format_general_results(result), 'message': 'success'}), 200
|
||||
|
||||
@db.route('/graph/add', methods=['POST'])
|
||||
def add_graph_entity():
|
||||
|
||||
BIN
web/public/icon-narrow.png
Normal file
BIN
web/public/icon-narrow.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 12 KiB |
BIN
web/public/icon.png
Normal file
BIN
web/public/icon.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 7.0 KiB |
@ -30,10 +30,10 @@
|
||||
</a>
|
||||
<template #overlay>
|
||||
<a-menu>
|
||||
<a-menu-item v-for="(db, index) in opts.databases" :key="index" @click="meta.selectedKB=index">
|
||||
<a-menu-item v-for="(db, index) in opts.databases" :key="index" @click="useDatabase(index)">
|
||||
<a href="javascript:;" >{{ db.name }}</a>
|
||||
</a-menu-item>
|
||||
<a-menu-item @click="meta.selectedKB = null">
|
||||
<a-menu-item @click="useDatabase(null)">
|
||||
<a href="javascript:;">不使用</a>
|
||||
</a-menu-item>
|
||||
</a-menu>
|
||||
@ -54,10 +54,10 @@
|
||||
</a>
|
||||
<template #overlay>
|
||||
<a-menu>
|
||||
<a-menu-item v-for="(db, index) in opts.databases" :key="index" @click="meta.selectedKB=index">
|
||||
<a-menu-item v-for="(db, index) in opts.databases" :key="index" @click="useDatabase(index)">
|
||||
<a href="javascript:;">{{ db.name }}</a>
|
||||
</a-menu-item>
|
||||
<a-menu-item @click="meta.selectedKB = null">
|
||||
<a-menu-item @click="useDatabase(null)">
|
||||
<a href="javascript:;">不使用</a>
|
||||
</a-menu-item>
|
||||
</a-menu>
|
||||
@ -71,15 +71,15 @@
|
||||
<div class="flex-center" @click="meta.use_web = !meta.use_web" v-if="configStore.config.enable_search_engine">
|
||||
搜索引擎（Bing） <div @click.stop><a-switch v-model:checked="meta.use_web" /></div>
|
||||
</div>
|
||||
<div class="flex-center" @click="meta.rewrite_query = !meta.rewrite_query" v-if="configStore.config.enable_reranker">
|
||||
重写查询 <div @click.stop><a-switch v-model:checked="meta.rewrite_query" /></div>
|
||||
</div>
|
||||
<div class="flex-center" @click="meta.rewrite_query = !meta.rewrite_query">
|
||||
<div class="flex-center" @click="meta.stream = !meta.stream">
|
||||
流式输出 <div @click.stop><a-switch v-model:checked="meta.stream" /></div>
|
||||
</div>
|
||||
<div class="flex-center" @click="meta.summary_title = !meta.summary_title">
|
||||
总结对话标题 <div @click.stop><a-switch v-model:checked="meta.summary_title" /></div>
|
||||
</div>
|
||||
<div class="flex-center" v-if="configStore.config.enable_knowledge_base">
|
||||
重写查询 <a-segmented v-model:value="meta.rewriteQuery" :options="['off', 'on', 'hyde']"/>
|
||||
</div>
|
||||
<div class="flex-center">
|
||||
最大历史轮数 <a-input-number id="inputNumber" v-model:value="meta.history_round" :min="1" :max="50" />
|
||||
</div>
|
||||
@ -162,6 +162,7 @@ import { onClickOutside } from '@vueuse/core'
|
||||
import { Marked } from 'marked';
|
||||
import { markedHighlight } from 'marked-highlight';
|
||||
import { useConfigStore } from '@/stores/config'
|
||||
import { message } from 'ant-design-vue'
|
||||
import RefsComponent from '@/components/RefsComponent.vue'
|
||||
import hljs from 'highlight.js';
|
||||
import 'highlight.js/styles/github.css';
|
||||
@ -197,7 +198,7 @@ const meta = reactive(JSON.parse(localStorage.getItem('meta')) || {
|
||||
use_graph: false,
|
||||
use_web: false,
|
||||
graph_name: "neo4j",
|
||||
rewrite_query: true,
|
||||
rewriteQuery: "off",
|
||||
selectedKB: null,
|
||||
stream: true,
|
||||
summary_title: true,
|
||||
@ -229,6 +230,17 @@ const renderMarkdown = (message) => {
|
||||
}
|
||||
}
|
||||
|
||||
const useDatabase = (index) => {
|
||||
const selected = opts.databases[index]
|
||||
console.log(selected)
|
||||
if (index != null && configStore.config.embed_model != selected.embed_model) {
|
||||
console.log(selected.embed_model, configStore.config.embed_model)
|
||||
message.error(`所选知识库的向量模型（${selected.embed_model}）与当前向量模型（${configStore.config.embed_model}) 不匹配，请重新选择`)
|
||||
} else {
|
||||
meta.selectedKB = index
|
||||
}
|
||||
}
|
||||
|
||||
const handleKeyDown = (e) => {
|
||||
if (e.key === 'Enter' && !e.shiftKey) {
|
||||
e.preventDefault()
|
||||
@ -517,7 +529,7 @@ watch(
|
||||
border-radius: 12px;
|
||||
padding: 12px;
|
||||
z-index: 11;
|
||||
width: 250px;
|
||||
width: 280px;
|
||||
|
||||
.flex-center {
|
||||
display: flex;
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
<template>
|
||||
<div class="log-viewer">
|
||||
<a-button @click="fetchLogs">刷新</a-button>
|
||||
<a-button @click="fetchLogs" :loading="state.fetching">刷新</a-button>
|
||||
<div ref="logContainer" class="log-container">
|
||||
<pre v-if="logs">{{ logs }}</pre>
|
||||
</div>
|
||||
@ -9,17 +9,21 @@
|
||||
</template>
|
||||
|
||||
<script setup>
|
||||
import { ref, onMounted, onActivated, nextTick } from 'vue';
|
||||
import { ref, onMounted, onActivated, nextTick, reactive } from 'vue';
|
||||
|
||||
// 定义一个 ref 来存储日志数据和错误信息
|
||||
const logs = ref('');
|
||||
const error = ref('');
|
||||
const state = reactive({
|
||||
fetching: false,
|
||||
})
|
||||
|
||||
// 定义一个 ref 来获取日志容器的 DOM 元素
|
||||
const logContainer = ref(null);
|
||||
|
||||
// 定义一个方法来获取日志数据
|
||||
const fetchLogs = async () => {
|
||||
state.fetching = true;
|
||||
try {
|
||||
// 清空之前的错误信息
|
||||
error.value = '';
|
||||
@ -45,41 +49,51 @@
|
||||
} catch (err) {
|
||||
// 处理请求错误
|
||||
error.value = `Error: ${err.message}`;
|
||||
} finally {
|
||||
state.fetching = false;
|
||||
}
|
||||
};
|
||||
|
||||
// 组件第一次挂载时自动获取日志并滚动到底部
|
||||
onMounted(() => {
|
||||
fetchLogs();
|
||||
setInterval(fetchLogs, 5000); // 每 5 秒自动刷新一次日志
|
||||
// setInterval(fetchLogs, 5000); // 每 5 秒自动刷新一次日志
|
||||
});
|
||||
|
||||
// 组件重新激活时（从 keep-alive 缓存恢复）自动获取日志
|
||||
onActivated(() => {
|
||||
fetchLogs();
|
||||
setTimeout(fetchLogs, 1000);
|
||||
});
|
||||
|
||||
</script>
|
||||
|
||||
<style scoped>
|
||||
.log-viewer {}
|
||||
.log-viewer {
|
||||
|
||||
}
|
||||
|
||||
.error {
|
||||
color: red;
|
||||
}
|
||||
|
||||
.log-container {
|
||||
max-height: 80vh; /* 设置最大高度 */
|
||||
height: calc(100vh - 150px); /* 设置最大高度 */
|
||||
overflow-y: auto; /* 启用垂直滚动 */
|
||||
background-color: #f0f0f0;
|
||||
margin: 20px 0;
|
||||
padding: 10px;
|
||||
background-color: #f6f7f7;
|
||||
padding-bottom: 0;
|
||||
border-radius: 5px;
|
||||
margin-top: 10px;
|
||||
white-space: pre-wrap; /* 使日志内容自动换行 */
|
||||
word-wrap: break-word;
|
||||
font-size: small;
|
||||
|
||||
background: #0C0C0C;
|
||||
color: #D1D1D1;
|
||||
|
||||
pre {
|
||||
margin: 0;
|
||||
padding: 10px;
|
||||
height: 100%;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
@ -6,11 +6,11 @@
|
||||
<span class="filetag item"
|
||||
v-for="(results, filename) in message.groupedResults"
|
||||
:key="filename"
|
||||
@click="openDetail = true"
|
||||
@click="toggleDrawer(filename)"
|
||||
>
|
||||
<FileTextOutlined /> {{ filename }}
|
||||
<a-drawer
|
||||
v-model:open="openDetail"
|
||||
v-model:open="openDetail[filename]"
|
||||
title="检索详情"
|
||||
width="800"
|
||||
:contentWrapperStyle="{ maxWidth: '100%'}"
|
||||
@ -19,7 +19,7 @@
|
||||
rootClassName="root"
|
||||
>
|
||||
<div class="fileinfo">
|
||||
<p><strong>文件名:</strong> {{ results[0].file.filename }}</p>
|
||||
<p><strong>文件名:</strong> {{ filename }}</p>
|
||||
<p><strong>文件类型:</strong> {{ results[0].file.type }}</p>
|
||||
<p><strong>创建时间:</strong> {{ new Date(results[0].file.created_at * 1000).toLocaleString() }}</p>
|
||||
</div>
|
||||
@ -47,7 +47,7 @@
|
||||
</template>
|
||||
|
||||
<script setup>
|
||||
import { ref, computed } from 'vue'
|
||||
import { ref, computed, reactive } from 'vue'
|
||||
import {
|
||||
GlobalOutlined,
|
||||
FileTextOutlined,
|
||||
@ -59,7 +59,18 @@ const props = defineProps({
|
||||
|
||||
const message = ref(props.message)
|
||||
|
||||
const openDetail = ref(false)
|
||||
// 使用 reactive 创建一个响应式对象来存储每个文件的抽屉状态
|
||||
const openDetail = reactive({})
|
||||
|
||||
// 初始化 openDetail 对象
|
||||
for (const filename in message.value.groupedResults) {
|
||||
openDetail[filename] = false
|
||||
}
|
||||
|
||||
const toggleDrawer = (filename) => {
|
||||
openDetail[filename] = !openDetail[filename]
|
||||
}
|
||||
|
||||
const showRefs = computed(() => message.value.role=='received' && message.value.status=='finished')
|
||||
</script>
|
||||
|
||||
|
||||
@ -54,8 +54,8 @@ console.log(route)
|
||||
|
||||
<template>
|
||||
<div class="app-layout">
|
||||
<div class="debug-panel">
|
||||
<div class="shown-btn" @click="showDebug=!showDebug"><BugOutlined /></div>
|
||||
<div class="debug-panel" @click="showDebug=!showDebug">
|
||||
<div class="shown-btn"><BugOutlined /></div>
|
||||
<a-drawer
|
||||
v-model:open="showDebug"
|
||||
title="调试面板"
|
||||
|
||||
@ -16,22 +16,41 @@
|
||||
<p class="description">{{ database.description }}</p>
|
||||
<div class="tags">
|
||||
<a-tag color="blue" v-if="database.embed_model">{{ database.embed_model }}</a-tag>
|
||||
<a-tag color="green" v-if="database.dimension">{{ database.dimension }}</a-tag>
|
||||
</div>
|
||||
<a-divider/>
|
||||
<div class="query-params" v-if="state.curPage == 'query-test'">
|
||||
<p style="text-align: center; margin: 0;"><strong>参数配置</strong></p>
|
||||
<!-- <p style="text-align: center; margin: 0;"><strong>参数配置</strong></p> -->
|
||||
<div class="params-item">
|
||||
<p>检索数量：</p>
|
||||
<a-input-number size="small" v-model:value="meta.queryCount" :min="1" :max="20" />
|
||||
<a-input-number size="small" v-model:value="meta.maxQueryCount" :min="1" :max="20" />
|
||||
</div>
|
||||
<div class="params-item">
|
||||
<p>TopK：</p>
|
||||
<a-input-number size="small" v-model:value="meta.topK" :min="1" :max="meta.maxQueryCount" />
|
||||
</div>
|
||||
<div class="params-item">
|
||||
<p>过滤低质量：</p>
|
||||
<a-switch v-model:checked="meta.filter" />
|
||||
</div>
|
||||
<div class="params-item">
|
||||
<p>重写查询</p>
|
||||
<a-segmented v-model:value="meta.rewrite_query" :options="['OFF', 'ON', 'HyDE']" />
|
||||
|
||||
<div class="params-item w100" v-if="configStore.config.enable_reranker">
|
||||
<p>重排序阈值：</p>
|
||||
<a-slider v-model:value="meta.rerankThreshold" :min="0" :max="1" :step="0.01" />
|
||||
</div>
|
||||
<div class="params-item w100">
|
||||
<p>距离阈值：</p>
|
||||
<a-slider v-model:value="meta.distanceThreshold" :min="0" :max="1" :step="0.01" />
|
||||
</div>
|
||||
<div class="params-item col">
|
||||
<p>重写查询<small>（修改后需重新检索）</small>：</p>
|
||||
<a-segmented v-model:value="meta.rewriteQuery" :options="rewriteQueryOptions" >
|
||||
<template #label="{ payload }">
|
||||
<div style="padding: 4px 4px">
|
||||
<p>{{ payload.title }}</p>
|
||||
<small>{{ payload.subTitle }}</small>
|
||||
</div>
|
||||
</template>
|
||||
</a-segmented>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@ -42,7 +61,7 @@
|
||||
<a-tab-pane key="add">
|
||||
<template #tab><span><CloudUploadOutlined />添加文件</span></template>
|
||||
<div class="db-info-container">
|
||||
<h3>向知识库中添加文件</h3>
|
||||
<!-- <h3>向知识库中添加文件</h3> -->
|
||||
<div class="upload">
|
||||
<a-upload-dragger
|
||||
class="upload-dragger"
|
||||
@ -117,7 +136,7 @@
|
||||
<a-tab-pane key="query-test" force-render>
|
||||
<template #tab><span><SearchOutlined />检索测试</span></template>
|
||||
<div class="db-info-container">
|
||||
<h3>检索测试</h3>
|
||||
<!-- <h3>检索测试</h3> -->
|
||||
<div class="query-action">
|
||||
<a-textarea
|
||||
v-model:value="queryText"
|
||||
@ -131,7 +150,11 @@
|
||||
</a-button>
|
||||
</div>
|
||||
<div class="query-test" v-if="queryResult">
|
||||
<div class="query-card" v-for="(result, idx) in (meta.filter ? queryResult.results : queryResult.all_results)" :key="idx">
|
||||
<div class="results-overview">
|
||||
<p>总数：{{ queryResult.all_results.length }}，过滤后：{{ filteredResults.length }} （TopK：{{ meta.topK }}）</p>
|
||||
<p>重写后查询：{{ queryResult.rw_query }}</p>
|
||||
</div>
|
||||
<div class="query-card" v-for="(result, idx) in (filteredResults)" :key="idx">
|
||||
<p>
|
||||
<strong>#{{ idx + 1 }} </strong>
|
||||
<span>{{ result.file.filename }} </span>
|
||||
@ -149,9 +172,10 @@
|
||||
</template>
|
||||
|
||||
<script setup>
|
||||
import { onMounted, reactive, ref, watch } from 'vue';
|
||||
import { onMounted, reactive, ref, watch, toRaw } from 'vue';
|
||||
import { message, Modal } from 'ant-design-vue';
|
||||
import { useRoute, useRouter } from 'vue-router';
|
||||
import { useConfigStore } from '@/stores/config'
|
||||
import {
|
||||
ReadFilled,
|
||||
LeftOutlined,
|
||||
@ -177,6 +201,8 @@ const selectedFile = ref(null);
|
||||
// 查询测试
|
||||
const queryText = ref('');
|
||||
const queryResult = ref(null)
|
||||
const filteredResults = ref([])
|
||||
const configStore = useConfigStore()
|
||||
|
||||
const state = reactive({
|
||||
loading: false,
|
||||
@ -189,11 +215,44 @@ const state = reactive({
|
||||
});
|
||||
|
||||
const meta = reactive({
|
||||
queryCount: 10,
|
||||
filter: false,
|
||||
rewrite_query: 'OFF',
|
||||
maxQueryCount: 30,
|
||||
filter: true,
|
||||
rewriteQuery: 'off',
|
||||
rerankThreshold: 0.1,
|
||||
distanceThreshold: 0.5,
|
||||
topK: 10,
|
||||
});
|
||||
|
||||
const rewriteQueryOptions = ref([
|
||||
{ value: 'off', payload: { title: 'off', subTitle: '不启用' } },
|
||||
{ value: 'on', payload: { title: 'on', subTitle: '启用重写' } },
|
||||
{ value: 'hyde', payload: { title: 'hyde', subTitle: '伪文档生成' } },
|
||||
])
|
||||
|
||||
const filterQueryResults = () => {
|
||||
if (!queryResult.value || !queryResult.value.all_results) {
|
||||
return;
|
||||
}
|
||||
|
||||
let results = toRaw(queryResult.value.all_results);
|
||||
console.log("results", results);
|
||||
|
||||
if (meta.filter) {
|
||||
results = results.filter(r => r.distance >= meta.distanceThreshold);
|
||||
console.log("before", results);
|
||||
if (configStore.config.enable_reranker) {
|
||||
results = results
|
||||
.filter(r => r.rerank_score >= meta.rerankThreshold)
|
||||
.sort((a, b) => b.rerank_score - a.rerank_score);
|
||||
}
|
||||
console.log("after", results);
|
||||
|
||||
results = results.slice(0, meta.topK);
|
||||
}
|
||||
|
||||
filteredResults.value = results;
|
||||
}
|
||||
|
||||
const onQuery = () => {
|
||||
console.log(queryText.value)
|
||||
state.searchLoading = true
|
||||
@ -214,6 +273,7 @@ const onQuery = () => {
|
||||
.then(data => {
|
||||
console.log(data)
|
||||
queryResult.value = data
|
||||
filterQueryResults()
|
||||
})
|
||||
.catch(error => {
|
||||
console.error(error)
|
||||
@ -348,7 +408,6 @@ const getDatabaseInfo = () => {
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
const deleteFile = (fileId) => {
|
||||
console.log(fileId)
|
||||
state.lock = true
|
||||
@ -410,36 +469,12 @@ const addDocumentByFile = () => {
|
||||
}
|
||||
|
||||
const columns = [
|
||||
{
|
||||
title: '文件ID',
|
||||
dataIndex: 'file_id',
|
||||
key: 'file_id',
|
||||
},
|
||||
{
|
||||
title: '文件名',
|
||||
dataIndex: 'filename',
|
||||
key: 'filename',
|
||||
},
|
||||
{
|
||||
title: '上传时间',
|
||||
dataIndex: 'created_at',
|
||||
key: 'created_at',
|
||||
},
|
||||
{
|
||||
title: '状态',
|
||||
dataIndex: 'status',
|
||||
key: 'status',
|
||||
},
|
||||
{
|
||||
title: '类型',
|
||||
dataIndex: 'type',
|
||||
key: 'type',
|
||||
},
|
||||
{
|
||||
title: '操作',
|
||||
key: 'action',
|
||||
dataIndex: 'file_id',
|
||||
}
|
||||
{ title: '文件ID', dataIndex: 'file_id', key: 'file_id' },
|
||||
{ title: '文件名', dataIndex: 'filename', key: 'filename' },
|
||||
{ title: '上传时间', dataIndex: 'created_at', key: 'created_at' },
|
||||
{ title: '状态', dataIndex: 'status', key: 'status' },
|
||||
{ title: '类型', dataIndex: 'type', key: 'type' },
|
||||
{ title: '操作', key: 'action', dataIndex: 'file_id' }
|
||||
];
|
||||
|
||||
watch(() => route.params.database_id, (newId) => {
|
||||
@ -450,6 +485,12 @@ watch(() => route.params.database_id, (newId) => {
|
||||
}
|
||||
);
|
||||
|
||||
// 检测到 meta 变化时重新查询
|
||||
watch(() => meta, () => {
|
||||
filterQueryResults()
|
||||
}, { deep: true })
|
||||
|
||||
|
||||
onMounted(() => {
|
||||
getDatabaseInfo();
|
||||
// const refreshInterval = setInterval(() => {
|
||||
@ -457,6 +498,8 @@ onMounted(() => {
|
||||
// }, 10000);
|
||||
})
|
||||
|
||||
|
||||
|
||||
</script>
|
||||
|
||||
<style lang="less" scoped>
|
||||
@ -546,15 +589,27 @@ onMounted(() => {
|
||||
}
|
||||
}
|
||||
.params-item.col {
|
||||
align-items: flex-start;
|
||||
flex-direction: column;
|
||||
width: 100%;
|
||||
height: auto;
|
||||
}
|
||||
|
||||
.params-item.w100 > *{
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
.ant-slider {
|
||||
margin: 6px 0px;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
.atab-container {
|
||||
padding: 12px 16px;
|
||||
padding: 0;
|
||||
width: 100%;
|
||||
max-height: 100%;
|
||||
overflow: auto;
|
||||
}
|
||||
|
||||
.db-info-container {
|
||||
@ -675,8 +730,10 @@ onMounted(() => {
|
||||
color: white;
|
||||
padding: 2px 4px;
|
||||
border-radius: 4px;
|
||||
font-size: small;
|
||||
font-size: 10px;
|
||||
font-weight: bold;
|
||||
opacity: 0.8;
|
||||
user-select: none;
|
||||
}
|
||||
|
||||
.pdf {
|
||||
@ -711,3 +768,24 @@ onMounted(() => {
|
||||
|
||||
}
|
||||
</style>
|
||||
|
||||
<style lang="less">
|
||||
.atab-container div.ant-tabs-nav {
|
||||
background: var(--main-light-5);
|
||||
padding: 1rem;
|
||||
padding-bottom: 0;
|
||||
}
|
||||
|
||||
.atab-container .ant-tabs-content-holder {
|
||||
padding: 0 1rem;
|
||||
}
|
||||
|
||||
.params-item.col .ant-segmented {
|
||||
width: 100%;
|
||||
|
||||
div.ant-segmented-group {
|
||||
display: flex;
|
||||
justify-content: space-around;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
@ -10,18 +10,18 @@
|
||||
</div>
|
||||
<div class="graph-container layout-container" v-else>
|
||||
<div class="info">
|
||||
<h2>图数据库 {{ graph?.database_name }}</h2>
|
||||
<h2>图数据库 {{ graphInfo?.database_name }}</h2>
|
||||
<p>
|
||||
<span v-if="state.graphloading">加载中</span>
|
||||
<span class="green-dot" v-if="graph?.status == 'open'"></span>
|
||||
<span v-if="state.loadingGraphInfo">加载中</span>
|
||||
<span class="green-dot" v-if="graphInfo?.status == 'open'"></span>
|
||||
<span class="red-dot" v-else></span>
|
||||
<span>{{ graph?.status }}</span> ·
|
||||
<span>共 {{ graph?.entity_count }} 实体，{{ graph?.relationship_count }} 个关系</span>
|
||||
<span>{{ graphInfo?.status }}</span> ·
|
||||
<span>共 {{ graphInfo?.entity_count }} 实体，{{ graphInfo?.relationship_count }} 个关系</span>
|
||||
</p>
|
||||
</div>
|
||||
<div class="actions">
|
||||
<div class="actions-left">
|
||||
<a-button @click="state.showModal = true">上传文件</a-button>
|
||||
<a-button @click="state.showModal = true"><UploadOutlined /> 上传文件</a-button>
|
||||
<a-modal
|
||||
:open="state.showModal" title="上传文件"
|
||||
@ok="addDocumentByFile"
|
||||
@ -42,15 +42,15 @@
|
||||
>
|
||||
<p class="ant-upload-text">点击或者把文件拖拽到这里上传</p>
|
||||
<p class="ant-upload-hint">
|
||||
目前仅支持上传文本文件，如 .pdf, .txt, .md。且同名文件无法重复添加。
|
||||
目前仅支持上传 jsonl 文件。且同名文件无法重复添加。
|
||||
</p>
|
||||
</a-upload-dragger>
|
||||
</div>
|
||||
</a-modal>
|
||||
<input v-model="sampleNodeCount">
|
||||
<a-button @click="loadSampleNodes">确定</a-button>
|
||||
<a-button @click="loadSampleNodes" :loading="state.fetching">获取节点</a-button>
|
||||
</div>
|
||||
<div class="action-right">
|
||||
<div class="actions-right">
|
||||
<input
|
||||
v-model="state.searchInput"
|
||||
placeholder="输入要查询的实体"
|
||||
@ -66,7 +66,8 @@
|
||||
</a-button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="main" id="container" ref="container"></div>
|
||||
<div class="main" id="container" ref="container" v-show="graphData.nodes.length > 0"></div>
|
||||
<a-empty v-show="graphData.nodes.length === 0" style="padding: 4rem 0;"/>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
@ -75,21 +76,23 @@ import { Graph } from "@antv/g6";
|
||||
import { computed, onMounted, reactive, ref } from 'vue';
|
||||
import { message } from "ant-design-vue";
|
||||
import { useConfigStore } from '@/stores/config';
|
||||
import { UploadOutlined } from '@ant-design/icons-vue';
|
||||
|
||||
const configStore = useConfigStore()
|
||||
|
||||
let graphInstance
|
||||
const graph = ref(null)
|
||||
const graphInfo = ref(null)
|
||||
const container = ref(null);
|
||||
const fileList = ref([]);
|
||||
const sampleNodeCount = ref(100);
|
||||
const subgraph = reactive({
|
||||
const graphData = reactive({
|
||||
nodes: [],
|
||||
edges: [],
|
||||
});
|
||||
|
||||
const state = reactive({
|
||||
graphloading: false,
|
||||
fetching: false,
|
||||
loadingGraphInfo: false,
|
||||
searchInput: '',
|
||||
searchLoading: false,
|
||||
showModal: false,
|
||||
@ -98,27 +101,27 @@ const state = reactive({
|
||||
})
|
||||
|
||||
|
||||
const loadGraph = () => {
|
||||
state.graphloading = true
|
||||
const loadGraphInfo = () => {
|
||||
state.loadingGraphInfo = true
|
||||
fetch('/api/database/graph', {
|
||||
method: "GET",
|
||||
})
|
||||
.then(response => response.json())
|
||||
.then(data => {
|
||||
console.log(data)
|
||||
graph.value = data.graph
|
||||
state.graphloading = false
|
||||
graphInfo.value = data.graph
|
||||
state.loadingGraphInfo = false
|
||||
})
|
||||
.catch(error => {
|
||||
console.error(error)
|
||||
message.error(error.message)
|
||||
state.graphloading = false
|
||||
state.loadingGraphInfo = false
|
||||
})
|
||||
}
|
||||
|
||||
const graphData = computed(() => {
|
||||
const getGraphData = () => {
|
||||
return {
|
||||
nodes: subgraph.nodes.map(node => {
|
||||
nodes: graphData.nodes.map(node => {
|
||||
return {
|
||||
id: node.id,
|
||||
data: {
|
||||
@ -126,7 +129,7 @@ const graphData = computed(() => {
|
||||
},
|
||||
}
|
||||
}),
|
||||
edges: subgraph.edges.map(edge => {
|
||||
edges: graphData.edges.map(edge => {
|
||||
return {
|
||||
source: edge.source_id,
|
||||
target: edge.target_id,
|
||||
@ -136,7 +139,7 @@ const graphData = computed(() => {
|
||||
}
|
||||
}),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
const addDocumentByFile = () => {
|
||||
state.precessing = true
|
||||
@ -159,6 +162,7 @@ const addDocumentByFile = () => {
|
||||
};
|
||||
|
||||
const loadSampleNodes = () => {
|
||||
state.fetching = true
|
||||
fetch(`/api/database/graph/nodes?kgdb_name=neo4j&num=${sampleNodeCount.value}`)
|
||||
.then((res) => {
|
||||
if (res.ok) {
|
||||
@ -168,17 +172,15 @@ const loadSampleNodes = () => {
|
||||
}
|
||||
})
|
||||
.then((data) => {
|
||||
subgraph.nodes = data.result.nodes
|
||||
subgraph.edges = data.result.edges
|
||||
console.log(data)
|
||||
console.log(subgraph)
|
||||
setTimeout(() => {
|
||||
randerGraph()
|
||||
}, 500)
|
||||
graphData.nodes = data.result.nodes
|
||||
graphData.edges = data.result.edges
|
||||
console.log(graphData)
|
||||
randerGraph()
|
||||
})
|
||||
.catch((error) => {
|
||||
message.error(error.message);
|
||||
})
|
||||
.finally(() => state.fetching = false)
|
||||
}
|
||||
|
||||
const onSearch = () => {
|
||||
@ -187,6 +189,12 @@ const onSearch = () => {
|
||||
return
|
||||
}
|
||||
|
||||
const cur_embed_model = configStore.config.embed_model
|
||||
if (cur_embed_model !== 'zhipu-embedding-3') {
|
||||
message.error('当前不支持实体检索，请在设置中选择向量模型为 zhipu-embedding-3')
|
||||
return
|
||||
}
|
||||
|
||||
state.searchLoading = true
|
||||
fetch(`/api/database/graph/node?entity_name=${state.searchInput}`)
|
||||
.then((res) => {
|
||||
@ -197,10 +205,13 @@ const onSearch = () => {
|
||||
}
|
||||
})
|
||||
.then((data) => {
|
||||
subgraph.nodes = data.result.nodes
|
||||
subgraph.edges = data.result.edges
|
||||
graphData.nodes = data.result.nodes
|
||||
graphData.edges = data.result.edges
|
||||
if (graphData.nodes.length === 0) {
|
||||
message.info('未找到相关实体')
|
||||
}
|
||||
console.log(data)
|
||||
console.log(subgraph)
|
||||
console.log(graphData)
|
||||
randerGraph()
|
||||
})
|
||||
.catch((error) => {
|
||||
@ -210,54 +221,58 @@ const onSearch = () => {
|
||||
};
|
||||
|
||||
const randerGraph = () => {
|
||||
graphInstance.setData(graphData.value);
|
||||
|
||||
if (graphInstance) {
|
||||
graphInstance.destroy();
|
||||
}
|
||||
|
||||
initGraph();
|
||||
graphInstance.setData(getGraphData());
|
||||
graphInstance.render();
|
||||
}
|
||||
|
||||
const initGraph = () => {
|
||||
graphInstance = new Graph({
|
||||
container: container.value,
|
||||
width: container.value.offsetWidth,
|
||||
height: container.value.offsetHeight,
|
||||
autoFit: true,
|
||||
autoResize: true,
|
||||
layout: {
|
||||
type: 'd3-force',
|
||||
preventOverlap: true,
|
||||
kr: 20,
|
||||
collide: {
|
||||
strength: 1.0,
|
||||
},
|
||||
},
|
||||
node: {
|
||||
type: 'circle',
|
||||
style: {
|
||||
labelText: (d) => d.data.label,
|
||||
size: 70,
|
||||
},
|
||||
palette: {
|
||||
field: 'label',
|
||||
color: 'tableau',
|
||||
},
|
||||
},
|
||||
edge: {
|
||||
type: 'line',
|
||||
style: {
|
||||
labelText: (d) => d.data.label,
|
||||
labelBackground: '#fff',
|
||||
endArrow: true,
|
||||
},
|
||||
},
|
||||
behaviors: ['drag-element', 'zoom-canvas', 'drag-canvas'],
|
||||
});
|
||||
window.addEventListener('resize', randerGraph);
|
||||
}
|
||||
|
||||
onMounted(() => {
|
||||
loadGraph();
|
||||
loadGraphInfo();
|
||||
loadSampleNodes();
|
||||
setTimeout(() => {
|
||||
if (state.showPage) {
|
||||
graphInstance = new Graph({
|
||||
container: container.value,
|
||||
width: container.value.offsetWidth,
|
||||
height: container.value.offsetHeight,
|
||||
autoFit: true,
|
||||
autoResize: true,
|
||||
layout: {
|
||||
type: 'd3-force',
|
||||
preventOverlap: true,
|
||||
kr: 100,
|
||||
collide: {
|
||||
strength: 0.5,
|
||||
},
|
||||
},
|
||||
node: {
|
||||
type: 'circle',
|
||||
style: {
|
||||
labelText: (d) => d.data.label,
|
||||
size: 40,
|
||||
},
|
||||
palette: {
|
||||
field: 'label',
|
||||
color: 'tableau',
|
||||
},
|
||||
},
|
||||
edge: {
|
||||
type: 'line',
|
||||
style: {
|
||||
labelText: (d) => d.data.label,
|
||||
labelBackground: '#fff',
|
||||
},
|
||||
},
|
||||
behaviors: ['drag-element', 'zoom-canvas', 'drag-canvas'],
|
||||
});
|
||||
graphInstance.setData(graphData.value);
|
||||
graphInstance.render();
|
||||
window.addEventListener('resize', randerGraph);
|
||||
}
|
||||
}, 400)
|
||||
});
|
||||
|
||||
|
||||
@ -303,7 +318,7 @@ const handleDrop = (event) => {
|
||||
justify-content: space-between;
|
||||
margin-bottom: 20px;
|
||||
|
||||
.actions-left {
|
||||
.actions-left, .actions-right {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 10px;
|
||||
@ -311,7 +326,6 @@ const handleDrop = (event) => {
|
||||
|
||||
input {
|
||||
width: 100px;
|
||||
margin-right: 10px;
|
||||
border-radius: 8px;
|
||||
padding: 4px 12px;
|
||||
border: 2px solid var(--main-300);
|
||||
@ -324,6 +338,7 @@ const handleDrop = (event) => {
|
||||
}
|
||||
|
||||
button {
|
||||
border-width: 2px;
|
||||
height: 40px;
|
||||
box-shadow: none;
|
||||
}
|
||||
@ -343,8 +358,9 @@ const handleDrop = (event) => {
|
||||
margin: 20px 0;
|
||||
border-radius: 16px;
|
||||
width: 100%;
|
||||
height: 800px;
|
||||
height: calc(100vh - 200px);
|
||||
resize: horizontal;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.database-empty {
|
||||
|
||||
@ -42,7 +42,8 @@
|
||||
</a-select>
|
||||
</div>
|
||||
<div class="card">
|
||||
<span class="label">{{ items?.embed_model.des }}
|
||||
<span class="label">
|
||||
{{ items?.embed_model.des }}
|
||||
<a-button small v-if="needRestart.embed_model" @click="sendRestart">
|
||||
<ReloadOutlined />需要重启
|
||||
</a-button>
|
||||
@ -58,7 +59,8 @@
|
||||
</a-select>
|
||||
</div>
|
||||
<div class="card">
|
||||
<span class="label">{{ items?.reranker.des }}
|
||||
<span class="label">
|
||||
{{ items?.reranker.des }}
|
||||
<a-button small v-if="needRestart.reranker" @click="sendRestart">
|
||||
<ReloadOutlined />需要重启
|
||||
</a-button>
|
||||
@ -89,7 +91,11 @@
|
||||
/>
|
||||
</div>
|
||||
<div class="card">
|
||||
<span class="label">{{ items?.enable_knowledge_graph.des }}</span>
|
||||
<span class="label">{{ items?.enable_knowledge_graph.des }}
|
||||
<a-button small v-if="needRestart.enable_knowledge_graph" @click="sendRestart">
|
||||
<ReloadOutlined />需要重启
|
||||
</a-button>
|
||||
</span>
|
||||
<a-switch
|
||||
:checked="configStore.config.enable_knowledge_graph"
|
||||
@change="handleChange('enable_knowledge_graph', !configStore.config.enable_knowledge_graph)"
|
||||
@ -194,6 +200,13 @@ const sendRestart = () => {
|
||||
|
||||
.label {
|
||||
margin-right: 20px;
|
||||
|
||||
button {
|
||||
margin-left: 10px;
|
||||
height: 24px;
|
||||
padding: 0 8px;
|
||||
font-size: smaller;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user