refactor(database): 文件结构优化,顺便修复了 chroma 的 embedding 错误
This commit is contained in:
parent
f51e7e0557
commit
770d098f05
@ -11,7 +11,7 @@ from sqlalchemy.orm import Session
|
|||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from src import executor, config
|
from src import executor, config
|
||||||
from src.core import HistoryManager
|
from src.knowledge import HistoryManager
|
||||||
from src.agents import agent_manager
|
from src.agents import agent_manager
|
||||||
from src.models import select_model
|
from src.models import select_model
|
||||||
from src.utils.logging_config import logger
|
from src.utils.logging_config import logger
|
||||||
|
|||||||
@ -10,11 +10,11 @@ from src.config import Config # noqa: E402
|
|||||||
config = Config()
|
config = Config()
|
||||||
|
|
||||||
# 导入知识库相关模块
|
# 导入知识库相关模块
|
||||||
from src.core.kb_factory import KnowledgeBaseFactory # noqa: E402
|
from src.knowledge.kb_factory import KnowledgeBaseFactory # noqa: E402
|
||||||
from src.core.kb_manager import KnowledgeBaseManager # noqa: E402
|
from src.knowledge.kb_manager import KnowledgeBaseManager # noqa: E402
|
||||||
from src.core.lightrag_kb import LightRagKB # noqa: E402
|
from src.knowledge.lightrag_kb import LightRagKB # noqa: E402
|
||||||
from src.core.chroma_kb import ChromaKB # noqa: E402
|
from src.knowledge.chroma_kb import ChromaKB # noqa: E402
|
||||||
from src.core.milvus_kb import MilvusKB # noqa: E402
|
from src.knowledge.milvus_kb import MilvusKB # noqa: E402
|
||||||
|
|
||||||
# 注册知识库类型
|
# 注册知识库类型
|
||||||
KnowledgeBaseFactory.register("lightrag", LightRagKB, {
|
KnowledgeBaseFactory.register("lightrag", LightRagKB, {
|
||||||
@ -37,5 +37,5 @@ KnowledgeBaseFactory.register("milvus", MilvusKB, {
|
|||||||
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
||||||
knowledge_base = KnowledgeBaseManager(work_dir)
|
knowledge_base = KnowledgeBaseManager(work_dir)
|
||||||
|
|
||||||
from src.core import GraphDatabase # noqa: E402
|
from src.knowledge import GraphDatabase # noqa: E402
|
||||||
graph_base = GraphDatabase()
|
graph_base = GraphDatabase()
|
||||||
|
|||||||
@ -6,81 +6,19 @@ from pathlib import Path
|
|||||||
from typing import Optional, Dict, List, Any
|
from typing import Optional, Dict, List, Any
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
try:
|
import chromadb
|
||||||
import chromadb
|
from chromadb.config import Settings
|
||||||
from chromadb.config import Settings
|
from chromadb.api.types import EmbeddingFunction, Documents, Embeddings
|
||||||
from chromadb.api.types import EmbeddingFunction, Documents, Embeddings
|
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
|
||||||
except ImportError:
|
|
||||||
chromadb = None
|
|
||||||
EmbeddingFunction = None
|
|
||||||
Documents = None
|
|
||||||
Embeddings = 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.utils import logger, hashstr
|
||||||
from src import config
|
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):
|
class ChromaKB(KnowledgeBase):
|
||||||
"""基于 ChromaDB 的向量知识库实现"""
|
"""基于 ChromaDB 的向量知识库实现"""
|
||||||
|
|
||||||
@ -136,7 +74,10 @@ class ChromaKB(KnowledgeBase):
|
|||||||
|
|
||||||
try:
|
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}")
|
logger.info(f"Retrieved existing collection: {collection_name}")
|
||||||
|
|
||||||
# 检查现有集合的配置是否匹配当前的 embed_info
|
# 检查现有集合的配置是否匹配当前的 embed_info
|
||||||
@ -144,6 +85,7 @@ class ChromaKB(KnowledgeBase):
|
|||||||
collection_metadata = collection.metadata or {}
|
collection_metadata = collection.metadata or {}
|
||||||
current_model = collection_metadata.get("embedding_model", "unknown")
|
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:
|
if current_model != expected_model:
|
||||||
logger.warning(f"Collection {collection_name} uses model '{current_model}', but expected '{expected_model}'. Recreating collection.")
|
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):
|
def _get_embedding_function(self, embed_info: Dict):
|
||||||
"""获取 embedding 函数"""
|
"""获取 embedding 函数"""
|
||||||
if embed_info:
|
config_dict = get_embedding_config(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
|
|
||||||
|
|
||||||
# 返回符合 ChromaDB 0.4.16+ 接口的 EmbeddingFunction 实例
|
|
||||||
return OpenAIEmbeddingFunction(
|
return OpenAIEmbeddingFunction(
|
||||||
model=model,
|
model_name=config_dict["model"],
|
||||||
api_key=api_key,
|
api_key=config_dict["api_key"],
|
||||||
base_url=base_url
|
api_base=config_dict["base_url"].replace('/embeddings', '')
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _get_chroma_collection(self, db_id: str):
|
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]:
|
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)
|
||||||
|
|
||||||
# 简单的分块策略:按段落和长度分割
|
# 为 ChromaDB 添加特定的 metadata 格式
|
||||||
paragraphs = text.split('\n\n')
|
for chunk in chunks:
|
||||||
|
chunk["metadata"] = {
|
||||||
current_chunk = ""
|
"source": chunk["source"],
|
||||||
chunk_index = 0
|
"chunk_id": chunk["chunk_id"],
|
||||||
|
"full_doc_id": file_id
|
||||||
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
|
|
||||||
}
|
|
||||||
})
|
|
||||||
|
|
||||||
return chunks
|
return chunks
|
||||||
|
|
||||||
@ -285,35 +172,17 @@ class ChromaKB(KnowledgeBase):
|
|||||||
processed_items_info = []
|
processed_items_info = []
|
||||||
|
|
||||||
for item in items:
|
for item in items:
|
||||||
# 根据内容类型生成不同的ID和文件名
|
# 准备文件元数据
|
||||||
if content_type == "file":
|
metadata = prepare_item_metadata(item, content_type, db_id)
|
||||||
file_path = Path(item)
|
file_id = metadata["file_id"]
|
||||||
file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}"
|
filename = metadata["filename"]
|
||||||
file_type = file_path.suffix.lower().replace(".", "")
|
item_path = metadata["path"]
|
||||||
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
|
|
||||||
|
|
||||||
# 添加文件记录
|
# 添加文件记录
|
||||||
file_record = {
|
file_record = metadata.copy()
|
||||||
"database_id": db_id,
|
|
||||||
"filename": filename,
|
|
||||||
"path": item_path,
|
|
||||||
"file_type": file_type,
|
|
||||||
"status": "processing",
|
|
||||||
"created_at": time.time()
|
|
||||||
}
|
|
||||||
self.files_meta[file_id] = file_record
|
self.files_meta[file_id] = file_record
|
||||||
self._save_metadata()
|
self._save_metadata()
|
||||||
|
|
||||||
# 添加 file_id 到返回数据
|
|
||||||
file_record = file_record.copy()
|
|
||||||
file_record["file_id"] = file_id
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# 根据内容类型处理内容
|
# 根据内容类型处理内容
|
||||||
if content_type == "file":
|
if content_type == "file":
|
||||||
@ -1,5 +1,5 @@
|
|||||||
from typing import Dict, Type, Any
|
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
|
from src.utils import logger
|
||||||
|
|
||||||
|
|
||||||
@ -4,8 +4,8 @@ import time
|
|||||||
from typing import Dict, Optional, List, Any
|
from typing import Dict, Optional, List, Any
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
from src.core.knowledge_base import KnowledgeBase, KBNotFoundError, KBOperationError
|
from src.knowledge.knowledge_base import KnowledgeBase, KBNotFoundError, KBOperationError
|
||||||
from src.core.kb_factory import KnowledgeBaseFactory
|
from src.knowledge.kb_factory import KnowledgeBaseFactory
|
||||||
from src.utils import logger
|
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':
|
if file_ext == '.pdf':
|
||||||
# 使用 OCR 处理 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)
|
text = await parse_pdf_async(str(file_path_obj), params=params)
|
||||||
return f"# {file_path_obj.name}\n\n{text}"
|
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.utils import EmbeddingFunc, setup_logger
|
||||||
from lightrag.kg.shared_storage import initialize_pipeline_status
|
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 import config
|
||||||
from src.utils import logger, hashstr, get_docker_safe_url
|
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):
|
def _get_embedding_func(self, embed_info: Dict):
|
||||||
"""获取 embedding 函数"""
|
"""获取 embedding 函数"""
|
||||||
if embed_info:
|
config_dict = get_embedding_config(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
|
|
||||||
|
|
||||||
return EmbeddingFunc(
|
return EmbeddingFunc(
|
||||||
embedding_dim=dimension,
|
embedding_dim=config_dict["dimension"],
|
||||||
max_token_size=4096,
|
max_token_size=4096,
|
||||||
func=lambda texts: openai_embed(
|
func=lambda texts: openai_embed(
|
||||||
texts=texts,
|
texts=texts,
|
||||||
model=model,
|
model=config_dict["model"],
|
||||||
api_key=api_key,
|
api_key=config_dict["api_key"],
|
||||||
base_url=base_url.replace("/embeddings", ""),
|
base_url=config_dict["base_url"].replace("/embeddings", ""),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -164,35 +154,17 @@ class LightRagKB(KnowledgeBase):
|
|||||||
processed_items_info = []
|
processed_items_info = []
|
||||||
|
|
||||||
for item in items:
|
for item in items:
|
||||||
# 根据内容类型生成不同的ID和文件名
|
# 准备文件元数据
|
||||||
if content_type == "file":
|
metadata = prepare_item_metadata(item, content_type, db_id)
|
||||||
file_path = Path(item)
|
file_id = metadata["file_id"]
|
||||||
file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}"
|
filename = metadata["filename"]
|
||||||
file_type = file_path.suffix.lower().replace(".", "")
|
item_path = metadata["path"]
|
||||||
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
|
|
||||||
|
|
||||||
# 添加文件记录
|
# 添加文件记录
|
||||||
file_record = {
|
file_record = metadata.copy()
|
||||||
"database_id": db_id,
|
|
||||||
"filename": filename,
|
|
||||||
"path": item_path,
|
|
||||||
"file_type": file_type,
|
|
||||||
"status": "processing",
|
|
||||||
"created_at": time.time()
|
|
||||||
}
|
|
||||||
self.files_meta[file_id] = file_record
|
self.files_meta[file_id] = file_record
|
||||||
self._save_metadata()
|
self._save_metadata()
|
||||||
|
|
||||||
# 添加 file_id 到返回数据
|
|
||||||
file_record = file_record.copy()
|
|
||||||
file_record["file_id"] = file_id
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# 根据内容类型处理内容
|
# 根据内容类型处理内容
|
||||||
if content_type == "file":
|
if content_type == "file":
|
||||||
@ -18,7 +18,8 @@ except ImportError:
|
|||||||
utility = None
|
utility = None
|
||||||
Collection = 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.utils import logger, hashstr
|
||||||
from src import config
|
from src import config
|
||||||
|
|
||||||
@ -171,16 +172,10 @@ class MilvusKB(KnowledgeBase):
|
|||||||
|
|
||||||
def _get_embedding_function(self, embed_info: Dict):
|
def _get_embedding_function(self, embed_info: Dict):
|
||||||
"""获取 embedding 函数"""
|
"""获取 embedding 函数"""
|
||||||
if embed_info:
|
config_dict = get_embedding_config(embed_info)
|
||||||
model = embed_info["name"]
|
model = config_dict["model"]
|
||||||
api_key = os.getenv(embed_info["api_key"], embed_info["api_key"])
|
api_key = config_dict["api_key"]
|
||||||
base_url = embed_info["base_url"]
|
base_url = config_dict["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
|
|
||||||
|
|
||||||
# 返回同步的嵌入函数
|
# 返回同步的嵌入函数
|
||||||
def embedding_function(texts):
|
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]:
|
def _split_text_into_chunks(self, text: str, file_id: str, filename: str) -> List[Dict]:
|
||||||
"""将文本分割成块"""
|
"""将文本分割成块"""
|
||||||
chunks = []
|
return 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,
|
|
||||||
"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
|
|
||||||
|
|
||||||
async def add_content(self, db_id: str, items: List[str],
|
async def add_content(self, db_id: str, items: List[str],
|
||||||
params: Optional[Dict] = None) -> List[Dict]:
|
params: Optional[Dict] = None) -> List[Dict]:
|
||||||
@ -300,28 +246,15 @@ class MilvusKB(KnowledgeBase):
|
|||||||
processed_items_info = []
|
processed_items_info = []
|
||||||
|
|
||||||
for item in items:
|
for item in items:
|
||||||
# 根据内容类型生成不同的ID和文件名
|
# 准备文件元数据
|
||||||
if content_type == "file":
|
metadata = prepare_item_metadata(item, content_type, db_id)
|
||||||
file_path = Path(item)
|
file_id = metadata["file_id"]
|
||||||
file_id = f"file_{hashstr(str(file_path) + str(time.time()), 6)}"
|
filename = metadata["filename"]
|
||||||
file_type = file_path.suffix.lower().replace(".", "")
|
item_path = metadata["path"]
|
||||||
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
|
|
||||||
|
|
||||||
# 添加文件记录
|
# 添加文件记录
|
||||||
file_record = {
|
file_record = metadata.copy()
|
||||||
"database_id": db_id,
|
del file_record["file_id"] # 从记录中移除file_id,因为它是key
|
||||||
"filename": filename,
|
|
||||||
"path": item_path,
|
|
||||||
"file_type": file_type,
|
|
||||||
"status": "processing",
|
|
||||||
"created_at": time.time()
|
|
||||||
}
|
|
||||||
self.files_meta[file_id] = file_record
|
self.files_meta[file_id] = file_record
|
||||||
self._save_metadata()
|
self._save_metadata()
|
||||||
|
|
||||||
@ -89,7 +89,7 @@
|
|||||||
:loading="state.chunkLoading"
|
:loading="state.chunkLoading"
|
||||||
:disabled="(uploadMode === 'file' && fileList.length === 0) || (uploadMode === 'url' && !urlList.trim())"
|
:disabled="(uploadMode === 'file' && fileList.length === 0) || (uploadMode === 'url' && !urlList.trim())"
|
||||||
>
|
>
|
||||||
生成分块
|
添加到知识库
|
||||||
</a-button>
|
</a-button>
|
||||||
</template>
|
</template>
|
||||||
<div class="add-files-content">
|
<div class="add-files-content">
|
||||||
@ -219,7 +219,7 @@
|
|||||||
<a-modal
|
<a-modal
|
||||||
v-model:open="state.fileDetailModalVisible"
|
v-model:open="state.fileDetailModalVisible"
|
||||||
:title="selectedFile?.filename || '文件详情'"
|
:title="selectedFile?.filename || '文件详情'"
|
||||||
width="800px"
|
width="1200px"
|
||||||
:footer="null"
|
:footer="null"
|
||||||
>
|
>
|
||||||
<div class="file-detail-content" v-if="selectedFile">
|
<div class="file-detail-content" v-if="selectedFile">
|
||||||
@ -1550,7 +1550,7 @@ const getKbTypeColor = (type) => {
|
|||||||
}
|
}
|
||||||
|
|
||||||
.header-container {
|
.header-container {
|
||||||
padding: 8px;
|
padding: 8px 16px;
|
||||||
height: 54px;
|
height: 54px;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -17,6 +17,9 @@
|
|||||||
<div class="status-wrapper">
|
<div class="status-wrapper">
|
||||||
<div class="status-indicator" :class="graphStatusClass"></div>
|
<div class="status-indicator" :class="graphStatusClass"></div>
|
||||||
</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 type="primary" @click="state.showModal = true" ><UploadOutlined/> 上传文件</a-button>
|
||||||
<a-button v-if="unindexedCount > 0" type="primary" @click="indexNodes" :loading="state.indexing">
|
<a-button v-if="unindexedCount > 0" type="primary" @click="indexNodes" :loading="state.indexing">
|
||||||
<SyncOutlined/> 为{{ unindexedCount }}个节点添加索引
|
<SyncOutlined/> 为{{ unindexedCount }}个节点添加索引
|
||||||
@ -89,10 +92,10 @@
|
|||||||
|
|
||||||
<script setup>
|
<script setup>
|
||||||
import { Graph } from "@antv/g6";
|
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 { message, Button as AButton } from 'ant-design-vue';
|
||||||
import { useConfigStore } from '@/stores/config';
|
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 HeaderComponent from '@/components/HeaderComponent.vue';
|
||||||
import { neo4jApi } from '@/apis/graph_api';
|
import { neo4jApi } from '@/apis/graph_api';
|
||||||
import { useUserStore } from '@/stores/user';
|
import { useUserStore } from '@/stores/user';
|
||||||
@ -388,6 +391,10 @@ const getAuthHeaders = () => {
|
|||||||
return userStore.getAuthHeaders();
|
return userStore.getAuthHeaders();
|
||||||
};
|
};
|
||||||
|
|
||||||
|
const openLink = (url) => {
|
||||||
|
window.open(url, '_blank')
|
||||||
|
}
|
||||||
|
|
||||||
</script>
|
</script>
|
||||||
|
|
||||||
<style lang="less" scoped>
|
<style lang="less" scoped>
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user