refactor(database): 文件结构优化，顺便修复了 chroma 的 embedding 错误
This commit is contained in:
Wenjie Zhang
parent
f51e7e0557
commit
770d098f05
@ -11,7 +11,7 @@ from sqlalchemy.orm import Session
|
||||
from pydantic import BaseModel
|
||||
|
||||
from src import executor, config
|
||||
from src.core import HistoryManager
|
||||
from src.knowledge import HistoryManager
|
||||
from src.agents import agent_manager
|
||||
from src.models import select_model
|
||||
from src.utils.logging_config import logger
|
||||
|
||||
@ -10,11 +10,11 @@ from src.config import Config # noqa: E402
|
||||
config = Config()
|
||||
|
||||
# 导入知识库相关模块
|
||||
from src.core.kb_factory import KnowledgeBaseFactory # noqa: E402
|
||||
from src.core.kb_manager import KnowledgeBaseManager # noqa: E402
|
||||
from src.core.lightrag_kb import LightRagKB # noqa: E402
|
||||
from src.core.chroma_kb import ChromaKB # noqa: E402
|
||||
from src.core.milvus_kb import MilvusKB # noqa: E402
|
||||
from src.knowledge.kb_factory import KnowledgeBaseFactory # noqa: E402
|
||||
from src.knowledge.kb_manager import KnowledgeBaseManager # noqa: E402
|
||||
from src.knowledge.lightrag_kb import LightRagKB # noqa: E402
|
||||
from src.knowledge.chroma_kb import ChromaKB # noqa: E402
|
||||
from src.knowledge.milvus_kb import MilvusKB # noqa: E402
|
||||
|
||||
# 注册知识库类型
|
||||
KnowledgeBaseFactory.register("lightrag", LightRagKB, {
|
||||
@ -37,5 +37,5 @@ KnowledgeBaseFactory.register("milvus", MilvusKB, {
|
||||
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
||||
knowledge_base = KnowledgeBaseManager(work_dir)
|
||||
|
||||
from src.core import GraphDatabase # noqa: E402
|
||||
from src.knowledge import GraphDatabase # noqa: E402
|
||||
graph_base = GraphDatabase()
|
||||
|
||||
@ -6,81 +6,19 @@ from pathlib import Path
|
||||
from typing import Optional, Dict, List, Any
|
||||
from datetime import datetime
|
||||
|
||||
try:
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
from chromadb.api.types import EmbeddingFunction, Documents, Embeddings
|
||||
except ImportError:
|
||||
chromadb = None
|
||||
EmbeddingFunction = None
|
||||
Documents = None
|
||||
Embeddings = None
|
||||
import chromadb
|
||||
from chromadb.config import Settings
|
||||
from chromadb.api.types import EmbeddingFunction, Documents, Embeddings
|
||||
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
|
||||
|
||||
from src.core.knowledge_base import KnowledgeBase
|
||||
|
||||
|
||||
from src.knowledge.knowledge_base import KnowledgeBase
|
||||
from src.knowledge.kb_utils import split_text_into_chunks, prepare_item_metadata, get_embedding_config
|
||||
from src.utils import logger, hashstr
|
||||
from src import config
|
||||
|
||||
|
||||
if EmbeddingFunction is not None:
|
||||
class OpenAIEmbeddingFunction(EmbeddingFunction):
|
||||
"""
|
||||
符合 ChromaDB 0.4.16+ 接口的 OpenAI 兼容嵌入函数
|
||||
"""
|
||||
|
||||
def __init__(self, model: str, api_key: str, base_url: str):
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url.replace("/embeddings", "")
|
||||
|
||||
def __call__(self, input: Documents) -> Embeddings:
|
||||
"""
|
||||
生成文档嵌入向量
|
||||
|
||||
Args:
|
||||
input: 文档列表（字符串列表）
|
||||
|
||||
Returns:
|
||||
Embeddings: 嵌入向量列表
|
||||
"""
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
from lightrag.llm.openai import openai_embed
|
||||
|
||||
# 确保输入是列表格式
|
||||
if isinstance(input, str):
|
||||
texts = [input]
|
||||
else:
|
||||
texts = list(input)
|
||||
|
||||
# 在新线程中运行异步函数，避免事件循环冲突
|
||||
def run_embedding():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(
|
||||
openai_embed(
|
||||
texts=texts,
|
||||
model=self.model,
|
||||
api_key=self.api_key,
|
||||
base_url=self.base_url,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
# 使用线程池执行异步函数
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(run_embedding)
|
||||
embeddings = future.result()
|
||||
|
||||
return embeddings
|
||||
else:
|
||||
# 如果 ChromaDB 没有安装，提供一个空的类
|
||||
class OpenAIEmbeddingFunction:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
class ChromaKB(KnowledgeBase):
|
||||
"""基于 ChromaDB 的向量知识库实现"""
|
||||
|
||||
@ -136,7 +74,10 @@ class ChromaKB(KnowledgeBase):
|
||||
|
||||
try:
|
||||
# 尝试获取现有集合
|
||||
collection = self.chroma_client.get_collection(name=collection_name)
|
||||
collection = self.chroma_client.get_collection(
|
||||
name=collection_name,
|
||||
embedding_function=embedding_function
|
||||
)
|
||||
logger.info(f"Retrieved existing collection: {collection_name}")
|
||||
|
||||
# 检查现有集合的配置是否匹配当前的 embed_info
|
||||
@ -144,6 +85,7 @@ class ChromaKB(KnowledgeBase):
|
||||
collection_metadata = collection.metadata or {}
|
||||
current_model = collection_metadata.get("embedding_model", "unknown")
|
||||
|
||||
logger.debug(f"Collection {collection_name} uses model '{current_model}', but expected '{expected_model}'.")
|
||||
# 如果模型不匹配，删除现有集合并重新创建
|
||||
if current_model != expected_model:
|
||||
logger.warning(f"Collection {collection_name} uses model '{current_model}', but expected '{expected_model}'. Recreating collection.")
|
||||
@ -173,22 +115,12 @@ class ChromaKB(KnowledgeBase):
|
||||
|
||||
def _get_embedding_function(self, embed_info: Dict):
|
||||
"""获取 embedding 函数"""
|
||||
if embed_info:
|
||||
model = embed_info["name"]
|
||||
api_key = os.getenv(embed_info["api_key"], embed_info["api_key"])
|
||||
base_url = embed_info["base_url"]
|
||||
else:
|
||||
from src.models import select_embedding_model
|
||||
default_model = select_embedding_model(config.embed_model)
|
||||
model = default_model.model
|
||||
api_key = default_model.api_key
|
||||
base_url = default_model.base_url
|
||||
config_dict = get_embedding_config(embed_info)
|
||||
|
||||
# 返回符合 ChromaDB 0.4.16+ 接口的 EmbeddingFunction 实例
|
||||
return OpenAIEmbeddingFunction(
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
base_url=base_url
|
||||
model_name=config_dict["model"],
|
||||
api_key=config_dict["api_key"],
|
||||
api_base=config_dict["base_url"].replace('/embeddings', '')
|
||||
)
|
||||
|
||||
async def _get_chroma_collection(self, db_id: str):
|
||||
@ -214,60 +146,15 @@ class ChromaKB(KnowledgeBase):
|
||||
|
||||
def _split_text_into_chunks(self, text: str, file_id: str, filename: str) -> List[Dict]:
|
||||
"""将文本分割成块"""
|
||||
chunks = []
|
||||
chunks = split_text_into_chunks(text, file_id, filename, self.chunk_size, self.chunk_overlap)
|
||||
|
||||
# 简单的分块策略：按段落和长度分割
|
||||
paragraphs = text.split('\n\n')
|
||||
|
||||
current_chunk = ""
|
||||
chunk_index = 0
|
||||
|
||||
for paragraph in paragraphs:
|
||||
paragraph = paragraph.strip()
|
||||
if not paragraph:
|
||||
continue
|
||||
|
||||
# 如果当前块加上新段落会超过限制，保存当前块
|
||||
if len(current_chunk) + len(paragraph) > self.chunk_size and current_chunk:
|
||||
chunks.append({
|
||||
"id": f"{file_id}_chunk_{chunk_index}",
|
||||
"content": current_chunk.strip(),
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": chunk_index,
|
||||
"metadata": {
|
||||
"source": filename,
|
||||
"chunk_id": f"{file_id}_chunk_{chunk_index}",
|
||||
"full_doc_id": file_id
|
||||
}
|
||||
})
|
||||
|
||||
# 开始新块，包含重叠内容
|
||||
if len(current_chunk) > self.chunk_overlap:
|
||||
current_chunk = current_chunk[-self.chunk_overlap:] + "\n\n" + paragraph
|
||||
else:
|
||||
current_chunk = paragraph
|
||||
chunk_index += 1
|
||||
else:
|
||||
if current_chunk:
|
||||
current_chunk += "\n\n" + paragraph
|
||||
else:
|
||||
current_chunk = paragraph
|
||||
|
||||
# 添加最后一块
|
||||
if current_chunk.strip():
|
||||
chunks.append({
|
||||
"id": f"{file_id}_chunk_{chunk_index}",
|
||||
"content": current_chunk.strip(),
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": chunk_index,
|
||||
"metadata": {
|
||||
"source": filename,
|
||||
"chunk_id": f"{file_id}_chunk_{chunk_index}",
|
||||
"full_doc_id": file_id
|
||||
}
|
||||
})
|
||||
# 为 ChromaDB 添加特定的 metadata 格式
|
||||
for chunk in chunks:
|
||||
chunk["metadata"] = {
|
||||
"source": chunk["source"],
|
||||
"chunk_id": chunk["chunk_id"],
|
||||
"full_doc_id": file_id
|
||||
}
|
||||
|
||||
return chunks
|
||||
|
||||
@ -285,35 +172,17 @@ class ChromaKB(KnowledgeBase):
|
||||
processed_items_info = []
|
||||
|
||||
for item in items:
|
||||
# 根据内容类型生成不同的ID和文件名
|
||||
if content_type == "file":
|
||||
file_path = Path(item)
|
||||
file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}"
|
||||
file_type = file_path.suffix.lower().replace(".", "")
|
||||
filename = file_path.name
|
||||
item_path = str(file_path)
|
||||
else: # URL
|
||||
file_id = f"url_{hashstr(item + str(time.time()), 6)}"
|
||||
file_type = "url"
|
||||
filename = f"webpage_{hashstr(item, 6)}.md"
|
||||
item_path = item
|
||||
# 准备文件元数据
|
||||
metadata = prepare_item_metadata(item, content_type, db_id)
|
||||
file_id = metadata["file_id"]
|
||||
filename = metadata["filename"]
|
||||
item_path = metadata["path"]
|
||||
|
||||
# 添加文件记录
|
||||
file_record = {
|
||||
"database_id": db_id,
|
||||
"filename": filename,
|
||||
"path": item_path,
|
||||
"file_type": file_type,
|
||||
"status": "processing",
|
||||
"created_at": time.time()
|
||||
}
|
||||
file_record = metadata.copy()
|
||||
self.files_meta[file_id] = file_record
|
||||
self._save_metadata()
|
||||
|
||||
# 添加 file_id 到返回数据
|
||||
file_record = file_record.copy()
|
||||
file_record["file_id"] = file_id
|
||||
|
||||
try:
|
||||
# 根据内容类型处理内容
|
||||
if content_type == "file":
|
||||
@ -1,5 +1,5 @@
|
||||
from typing import Dict, Type, Any
|
||||
from src.core.knowledge_base import KnowledgeBase, KBNotFoundError
|
||||
from src.knowledge.knowledge_base import KnowledgeBase, KBNotFoundError
|
||||
from src.utils import logger
|
||||
|
||||
|
||||
@ -4,8 +4,8 @@ import time
|
||||
from typing import Dict, Optional, List, Any
|
||||
from datetime import datetime
|
||||
|
||||
from src.core.knowledge_base import KnowledgeBase, KBNotFoundError, KBOperationError
|
||||
from src.core.kb_factory import KnowledgeBaseFactory
|
||||
from src.knowledge.knowledge_base import KnowledgeBase, KBNotFoundError, KBOperationError
|
||||
from src.knowledge.kb_factory import KnowledgeBaseFactory
|
||||
from src.utils import logger
|
||||
|
||||
|
||||
137
src/knowledge/kb_utils.py
Normal file
137
src/knowledge/kb_utils.py
Normal file
@ -0,0 +1,137 @@
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Any
|
||||
from src.utils import hashstr, get_docker_safe_url, logger
|
||||
from src import config
|
||||
|
||||
|
||||
def split_text_into_chunks(text: str, file_id: str, filename: str,
|
||||
chunk_size: int = 1000, chunk_overlap: int = 200) -> List[Dict]:
|
||||
"""
|
||||
将文本分割成块
|
||||
|
||||
Args:
|
||||
text: 要分割的文本
|
||||
file_id: 文件ID
|
||||
filename: 文件名
|
||||
chunk_size: 块大小
|
||||
chunk_overlap: 块重叠大小
|
||||
|
||||
Returns:
|
||||
List[Dict]: 分割后的文本块列表
|
||||
"""
|
||||
chunks = []
|
||||
|
||||
# 简单的分块策略：按段落和长度分割
|
||||
paragraphs = text.split('\n\n')
|
||||
|
||||
current_chunk = ""
|
||||
chunk_index = 0
|
||||
|
||||
for paragraph in paragraphs:
|
||||
paragraph = paragraph.strip()
|
||||
if not paragraph:
|
||||
continue
|
||||
|
||||
# 如果当前块加上新段落会超过限制，保存当前块
|
||||
if len(current_chunk) + len(paragraph) > chunk_size and current_chunk:
|
||||
chunks.append({
|
||||
"id": f"{file_id}_chunk_{chunk_index}",
|
||||
"content": current_chunk.strip(),
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": chunk_index,
|
||||
"source": filename,
|
||||
"chunk_id": f"{file_id}_chunk_{chunk_index}"
|
||||
})
|
||||
|
||||
# 开始新块，包含重叠内容
|
||||
if len(current_chunk) > chunk_overlap:
|
||||
current_chunk = current_chunk[-chunk_overlap:] + "\n\n" + paragraph
|
||||
else:
|
||||
current_chunk = paragraph
|
||||
chunk_index += 1
|
||||
else:
|
||||
if current_chunk:
|
||||
current_chunk += "\n\n" + paragraph
|
||||
else:
|
||||
current_chunk = paragraph
|
||||
|
||||
# 添加最后一块
|
||||
if current_chunk.strip():
|
||||
chunks.append({
|
||||
"id": f"{file_id}_chunk_{chunk_index}",
|
||||
"content": current_chunk.strip(),
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": chunk_index,
|
||||
"source": filename,
|
||||
"chunk_id": f"{file_id}_chunk_{chunk_index}"
|
||||
})
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def prepare_item_metadata(item: str, content_type: str, db_id: str) -> Dict:
|
||||
"""
|
||||
准备文件或URL的元数据
|
||||
|
||||
Args:
|
||||
item: 文件路径或URL
|
||||
content_type: 内容类型 ('file' 或 'url')
|
||||
db_id: 数据库ID
|
||||
|
||||
Returns:
|
||||
Dict: 包含元数据的字典
|
||||
"""
|
||||
if content_type == "file":
|
||||
file_path = Path(item)
|
||||
file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}"
|
||||
file_type = file_path.suffix.lower().replace(".", "")
|
||||
filename = file_path.name
|
||||
item_path = str(file_path)
|
||||
else: # URL
|
||||
file_id = f"url_{hashstr(item + str(time.time()), 6)}"
|
||||
file_type = "url"
|
||||
filename = f"webpage_{hashstr(item, 6)}.md"
|
||||
item_path = item
|
||||
|
||||
return {
|
||||
"database_id": db_id,
|
||||
"filename": filename,
|
||||
"path": item_path,
|
||||
"file_type": file_type,
|
||||
"status": "processing",
|
||||
"created_at": time.time(),
|
||||
"file_id": file_id
|
||||
}
|
||||
|
||||
|
||||
def get_embedding_config(embed_info: Dict) -> Dict:
|
||||
"""
|
||||
获取嵌入模型配置
|
||||
|
||||
Args:
|
||||
embed_info: 嵌入信息字典
|
||||
|
||||
Returns:
|
||||
Dict: 标准化的嵌入配置
|
||||
"""
|
||||
config_dict = {}
|
||||
|
||||
if embed_info:
|
||||
config_dict['model'] = embed_info["name"]
|
||||
config_dict['api_key'] = os.getenv(embed_info["api_key"], embed_info["api_key"])
|
||||
config_dict['base_url'] = embed_info["base_url"]
|
||||
config_dict['dimension'] = embed_info.get("dimension", 1024)
|
||||
else:
|
||||
from src.models import select_embedding_model
|
||||
default_model = select_embedding_model(config.embed_model)
|
||||
config_dict['model'] = default_model.model
|
||||
config_dict['api_key'] = default_model.api_key
|
||||
config_dict['base_url'] = default_model.base_url
|
||||
config_dict['dimension'] = getattr(default_model, 'dimension', 1024)
|
||||
|
||||
logger.debug(f"Embedding config: {config_dict}")
|
||||
return config_dict
|
||||
@ -388,7 +388,7 @@ class KnowledgeBase(ABC):
|
||||
|
||||
if file_ext == '.pdf':
|
||||
# 使用 OCR 处理 PDF
|
||||
from src.core.indexing import parse_pdf_async
|
||||
from src.knowledge.indexing import parse_pdf_async
|
||||
text = await parse_pdf_async(str(file_path_obj), params=params)
|
||||
return f"# {file_path_obj.name}\n\n{text}"
|
||||
|
||||
@ -10,7 +10,8 @@ from lightrag.llm.openai import openai_complete_if_cache, openai_embed
|
||||
from lightrag.utils import EmbeddingFunc, setup_logger
|
||||
from lightrag.kg.shared_storage import initialize_pipeline_status
|
||||
|
||||
from src.core.knowledge_base import KnowledgeBase
|
||||
from src.knowledge.knowledge_base import KnowledgeBase
|
||||
from src.knowledge.kb_utils import split_text_into_chunks, prepare_item_metadata, get_embedding_config
|
||||
from src import config
|
||||
from src.utils import logger, hashstr, get_docker_safe_url
|
||||
|
||||
@ -126,27 +127,16 @@ class LightRagKB(KnowledgeBase):
|
||||
|
||||
def _get_embedding_func(self, embed_info: Dict):
|
||||
"""获取 embedding 函数"""
|
||||
if embed_info:
|
||||
model = embed_info["name"]
|
||||
api_key = os.getenv(embed_info["api_key"], embed_info["api_key"])
|
||||
base_url = get_docker_safe_url(embed_info["base_url"])
|
||||
dimension = embed_info["dimension"]
|
||||
else:
|
||||
from src.models import select_embedding_model
|
||||
default_model = select_embedding_model(config.embed_model)
|
||||
model = default_model.model
|
||||
api_key = default_model.api_key
|
||||
base_url = default_model.base_url
|
||||
dimension = default_model.dimension
|
||||
config_dict = get_embedding_config(embed_info)
|
||||
|
||||
return EmbeddingFunc(
|
||||
embedding_dim=dimension,
|
||||
embedding_dim=config_dict["dimension"],
|
||||
max_token_size=4096,
|
||||
func=lambda texts: openai_embed(
|
||||
texts=texts,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
base_url=base_url.replace("/embeddings", ""),
|
||||
model=config_dict["model"],
|
||||
api_key=config_dict["api_key"],
|
||||
base_url=config_dict["base_url"].replace("/embeddings", ""),
|
||||
),
|
||||
)
|
||||
|
||||
@ -164,35 +154,17 @@ class LightRagKB(KnowledgeBase):
|
||||
processed_items_info = []
|
||||
|
||||
for item in items:
|
||||
# 根据内容类型生成不同的ID和文件名
|
||||
if content_type == "file":
|
||||
file_path = Path(item)
|
||||
file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}"
|
||||
file_type = file_path.suffix.lower().replace(".", "")
|
||||
filename = file_path.name
|
||||
item_path = str(file_path)
|
||||
else: # URL
|
||||
file_id = f"url_{hashstr(item + str(time.time()), 6)}"
|
||||
file_type = "url"
|
||||
filename = f"webpage_{hashstr(item, 6)}.md"
|
||||
item_path = item
|
||||
# 准备文件元数据
|
||||
metadata = prepare_item_metadata(item, content_type, db_id)
|
||||
file_id = metadata["file_id"]
|
||||
filename = metadata["filename"]
|
||||
item_path = metadata["path"]
|
||||
|
||||
# 添加文件记录
|
||||
file_record = {
|
||||
"database_id": db_id,
|
||||
"filename": filename,
|
||||
"path": item_path,
|
||||
"file_type": file_type,
|
||||
"status": "processing",
|
||||
"created_at": time.time()
|
||||
}
|
||||
file_record = metadata.copy()
|
||||
self.files_meta[file_id] = file_record
|
||||
self._save_metadata()
|
||||
|
||||
# 添加 file_id 到返回数据
|
||||
file_record = file_record.copy()
|
||||
file_record["file_id"] = file_id
|
||||
|
||||
try:
|
||||
# 根据内容类型处理内容
|
||||
if content_type == "file":
|
||||
@ -18,7 +18,8 @@ except ImportError:
|
||||
utility = None
|
||||
Collection = None
|
||||
|
||||
from src.core.knowledge_base import KnowledgeBase
|
||||
from src.knowledge.knowledge_base import KnowledgeBase
|
||||
from src.knowledge.kb_utils import split_text_into_chunks, prepare_item_metadata, get_embedding_config
|
||||
from src.utils import logger, hashstr
|
||||
from src import config
|
||||
|
||||
@ -171,16 +172,10 @@ class MilvusKB(KnowledgeBase):
|
||||
|
||||
def _get_embedding_function(self, embed_info: Dict):
|
||||
"""获取 embedding 函数"""
|
||||
if embed_info:
|
||||
model = embed_info["name"]
|
||||
api_key = os.getenv(embed_info["api_key"], embed_info["api_key"])
|
||||
base_url = embed_info["base_url"]
|
||||
else:
|
||||
from src.models import select_embedding_model
|
||||
default_model = select_embedding_model(config.embed_model)
|
||||
model = default_model.model
|
||||
api_key = default_model.api_key
|
||||
base_url = default_model.base_url
|
||||
config_dict = get_embedding_config(embed_info)
|
||||
model = config_dict["model"]
|
||||
api_key = config_dict["api_key"]
|
||||
base_url = config_dict["base_url"]
|
||||
|
||||
# 返回同步的嵌入函数
|
||||
def embedding_function(texts):
|
||||
@ -232,56 +227,7 @@ class MilvusKB(KnowledgeBase):
|
||||
|
||||
def _split_text_into_chunks(self, text: str, file_id: str, filename: str) -> List[Dict]:
|
||||
"""将文本分割成块"""
|
||||
chunks = []
|
||||
|
||||
# 简单的分块策略：按段落和长度分割
|
||||
paragraphs = text.split('\n\n')
|
||||
|
||||
current_chunk = ""
|
||||
chunk_index = 0
|
||||
|
||||
for paragraph in paragraphs:
|
||||
paragraph = paragraph.strip()
|
||||
if not paragraph:
|
||||
continue
|
||||
|
||||
# 如果当前块加上新段落会超过限制，保存当前块
|
||||
if len(current_chunk) + len(paragraph) > self.chunk_size and current_chunk:
|
||||
chunks.append({
|
||||
"id": f"{file_id}_chunk_{chunk_index}",
|
||||
"content": current_chunk.strip(),
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": chunk_index,
|
||||
"source": filename,
|
||||
"chunk_id": f"{file_id}_chunk_{chunk_index}"
|
||||
})
|
||||
|
||||
# 开始新块，包含重叠内容
|
||||
if len(current_chunk) > self.chunk_overlap:
|
||||
current_chunk = current_chunk[-self.chunk_overlap:] + "\n\n" + paragraph
|
||||
else:
|
||||
current_chunk = paragraph
|
||||
chunk_index += 1
|
||||
else:
|
||||
if current_chunk:
|
||||
current_chunk += "\n\n" + paragraph
|
||||
else:
|
||||
current_chunk = paragraph
|
||||
|
||||
# 添加最后一块
|
||||
if current_chunk.strip():
|
||||
chunks.append({
|
||||
"id": f"{file_id}_chunk_{chunk_index}",
|
||||
"content": current_chunk.strip(),
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": chunk_index,
|
||||
"source": filename,
|
||||
"chunk_id": f"{file_id}_chunk_{chunk_index}"
|
||||
})
|
||||
|
||||
return chunks
|
||||
return split_text_into_chunks(text, file_id, filename, self.chunk_size, self.chunk_overlap)
|
||||
|
||||
async def add_content(self, db_id: str, items: List[str],
|
||||
params: Optional[Dict] = None) -> List[Dict]:
|
||||
@ -300,28 +246,15 @@ class MilvusKB(KnowledgeBase):
|
||||
processed_items_info = []
|
||||
|
||||
for item in items:
|
||||
# 根据内容类型生成不同的ID和文件名
|
||||
if content_type == "file":
|
||||
file_path = Path(item)
|
||||
file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}"
|
||||
file_type = file_path.suffix.lower().replace(".", "")
|
||||
filename = file_path.name
|
||||
item_path = str(file_path)
|
||||
else: # URL
|
||||
file_id = f"url_{hashstr(item + str(time.time()), 6)}"
|
||||
file_type = "url"
|
||||
filename = f"webpage_{hashstr(item, 6)}.md"
|
||||
item_path = item
|
||||
# 准备文件元数据
|
||||
metadata = prepare_item_metadata(item, content_type, db_id)
|
||||
file_id = metadata["file_id"]
|
||||
filename = metadata["filename"]
|
||||
item_path = metadata["path"]
|
||||
|
||||
# 添加文件记录
|
||||
file_record = {
|
||||
"database_id": db_id,
|
||||
"filename": filename,
|
||||
"path": item_path,
|
||||
"file_type": file_type,
|
||||
"status": "processing",
|
||||
"created_at": time.time()
|
||||
}
|
||||
file_record = metadata.copy()
|
||||
del file_record["file_id"] # 从记录中移除file_id，因为它是key
|
||||
self.files_meta[file_id] = file_record
|
||||
self._save_metadata()
|
||||
|
||||
@ -89,7 +89,7 @@
|
||||
:loading="state.chunkLoading"
|
||||
:disabled="(uploadMode === 'file' && fileList.length === 0) || (uploadMode === 'url' && !urlList.trim())"
|
||||
>
|
||||
生成分块
|
||||
添加到知识库
|
||||
</a-button>
|
||||
</template>
|
||||
<div class="add-files-content">
|
||||
@ -219,7 +219,7 @@
|
||||
<a-modal
|
||||
v-model:open="state.fileDetailModalVisible"
|
||||
:title="selectedFile?.filename || '文件详情'"
|
||||
width="800px"
|
||||
width="1200px"
|
||||
:footer="null"
|
||||
>
|
||||
<div class="file-detail-content" v-if="selectedFile">
|
||||
@ -1550,7 +1550,7 @@ const getKbTypeColor = (type) => {
|
||||
}
|
||||
|
||||
.header-container {
|
||||
padding: 8px;
|
||||
padding: 8px 16px;
|
||||
height: 54px;
|
||||
}
|
||||
|
||||
|
||||
@ -17,6 +17,9 @@
|
||||
<div class="status-wrapper">
|
||||
<div class="status-indicator" :class="graphStatusClass"></div>
|
||||
</div>
|
||||
<a-button type="default" @click="openLink('http://localhost:7474/')" :icon="h(GlobalOutlined)">
|
||||
Neo4j 浏览器
|
||||
</a-button>
|
||||
<a-button type="primary" @click="state.showModal = true" ><UploadOutlined/> 上传文件</a-button>
|
||||
<a-button v-if="unindexedCount > 0" type="primary" @click="indexNodes" :loading="state.indexing">
|
||||
<SyncOutlined/> 为{{ unindexedCount }}个节点添加索引
|
||||
@ -89,10 +92,10 @@
|
||||
|
||||
<script setup>
|
||||
import { Graph } from "@antv/g6";
|
||||
import { computed, onMounted, reactive, ref } from 'vue';
|
||||
import { computed, onMounted, reactive, ref, h } from 'vue';
|
||||
import { message, Button as AButton } from 'ant-design-vue';
|
||||
import { useConfigStore } from '@/stores/config';
|
||||
import { UploadOutlined, SyncOutlined } from '@ant-design/icons-vue';
|
||||
import { UploadOutlined, SyncOutlined, GlobalOutlined } from '@ant-design/icons-vue';
|
||||
import HeaderComponent from '@/components/HeaderComponent.vue';
|
||||
import { neo4jApi } from '@/apis/graph_api';
|
||||
import { useUserStore } from '@/stores/user';
|
||||
@ -388,6 +391,10 @@ const getAuthHeaders = () => {
|
||||
return userStore.getAuthHeaders();
|
||||
};
|
||||
|
||||
const openLink = (url) => {
|
||||
window.open(url, '_blank')
|
||||
}
|
||||
|
||||
</script>
|
||||
|
||||
<style lang="less" scoped>
|
||||
|
||||
Loading…
Reference in New Issue
Block a user