Ubuntu
169 lines
6.0 KiB
Python
169 lines
6.0 KiB
Python
"""
|
||
知识库管理模块 - 使用 BAAI/bge-small-zh-v1.5 中文嵌入模型（本地离线版）
|
||
"""
|
||
import os
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from langchain_chroma import Chroma
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain_community.document_loaders import TextLoader
|
||
from langchain_core.documents import Document
|
||
|
||
from config import (
|
||
KNOWLEDGE_DIR, CHROMA_DB_DIR,
|
||
CHUNK_SIZE, CHUNK_OVERLAP,
|
||
RETRIEVAL_TOP_K
|
||
)
|
||
|
||
|
||
class KnowledgeBase:
|
||
"""知识库管理类 - 基于本地 BGE 中文嵌入模型"""
|
||
|
||
def __init__(self):
|
||
self.db_dir = str(CHROMA_DB_DIR)
|
||
self.collection_name = "customer_service_kb"
|
||
|
||
# 从本地加载 BGE 模型
|
||
print("🔧 正在加载本地 BGE 中文嵌入模型...")
|
||
|
||
# 直接使用 sentence_transformers 加载本地模型
|
||
from sentence_transformers import SentenceTransformer
|
||
|
||
# 本地模型路径
|
||
local_model_path = '/home/ubuntu/.cache/huggingface/models--BAAI--bge-small-zh-v1.5/snapshots/'
|
||
# 找到具体的 snapshot 目录
|
||
if os.path.exists(local_model_path):
|
||
snapshots = os.listdir(local_model_path)
|
||
if snapshots:
|
||
model_path = os.path.join(local_model_path, snapshots[0])
|
||
else:
|
||
model_path = 'BAAI/bge-small-zh-v1.5' # fallback
|
||
else:
|
||
model_path = 'BAAI/bge-small-zh-v1.5'
|
||
|
||
print(f" 模型路径: {model_path}")
|
||
|
||
# 加载模型
|
||
self.model = SentenceTransformer(model_path, device='cpu')
|
||
|
||
# 创建自定义 Embeddings 类
|
||
from langchain_core.embeddings import Embeddings
|
||
|
||
class LocalBGEEmbeddings(Embeddings):
|
||
def __init__(self, model):
|
||
self.model = model
|
||
|
||
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||
"""嵌入文档"""
|
||
embeddings = self.model.encode(texts, normalize_embeddings=True)
|
||
return embeddings.tolist()
|
||
|
||
def embed_query(self, text: str) -> List[float]:
|
||
"""嵌入查询"""
|
||
embedding = self.model.encode(text, normalize_embeddings=True)
|
||
return embedding.tolist()
|
||
|
||
self.embeddings = LocalBGEEmbeddings(self.model)
|
||
print("✅ BGE 模型加载完成")
|
||
|
||
# 初始化或加载向量数据库
|
||
self.vectorstore = self._load_or_create_db()
|
||
|
||
def _load_or_create_db(self) -> Chroma:
|
||
"""加载或创建向量数据库"""
|
||
if os.path.exists(self.db_dir) and os.listdir(self.db_dir):
|
||
print(f"📚 加载已有知识库: {self.db_dir}")
|
||
return Chroma(
|
||
collection_name=self.collection_name,
|
||
embedding_function=self.embeddings,
|
||
persist_directory=self.db_dir
|
||
)
|
||
else:
|
||
print(f"🆕 创建新知识库: {self.db_dir}")
|
||
os.makedirs(self.db_dir, exist_ok=True)
|
||
return Chroma(
|
||
collection_name=self.collection_name,
|
||
embedding_function=self.embeddings,
|
||
persist_directory=self.db_dir
|
||
)
|
||
|
||
def load_documents(self, directory: str = None) -> List[Document]:
|
||
"""从目录加载所有文档"""
|
||
directory = directory or str(KNOWLEDGE_DIR)
|
||
documents = []
|
||
|
||
if not os.path.exists(directory):
|
||
print(f"⚠️ 目录不存在: {directory}")
|
||
return documents
|
||
|
||
for file_path in Path(directory).glob("*"):
|
||
if file_path.suffix.lower() == ".txt":
|
||
loader = TextLoader(str(file_path), encoding="utf-8")
|
||
documents.extend(loader.load())
|
||
print(f"✅ 加载 TXT: {file_path.name}")
|
||
|
||
elif file_path.suffix.lower() == ".md":
|
||
loader = TextLoader(str(file_path), encoding="utf-8")
|
||
documents.extend(loader.load())
|
||
print(f"✅ 加载 MD: {file_path.name}")
|
||
|
||
return documents
|
||
|
||
def add_documents(self, documents: List[Document]):
|
||
"""添加文档到知识库"""
|
||
if not documents:
|
||
print("⚠️ 没有文档需要添加")
|
||
return
|
||
|
||
# 文本分块
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=CHUNK_SIZE,
|
||
chunk_overlap=CHUNK_OVERLAP,
|
||
separators=["\n\n", "\n", "。", "，", " ", ""]
|
||
)
|
||
chunks = text_splitter.split_documents(documents)
|
||
print(f"📄 文档分块: {len(documents)} 个文档 → {len(chunks)} 个片段")
|
||
|
||
# 添加到向量数据库
|
||
self.vectorstore.add_documents(chunks)
|
||
print(f"✅ 已添加到知识库")
|
||
|
||
def search(self, query: str, k: int = None) -> List[Document]:
|
||
"""语义搜索相关知识"""
|
||
k = k or RETRIEVAL_TOP_K
|
||
results = self.vectorstore.similarity_search(query, k=k)
|
||
return results
|
||
|
||
def get_retriever(self):
|
||
"""获取检索器（用于 LangChain）"""
|
||
return self.vectorstore.as_retriever(search_kwargs={"k": RETRIEVAL_TOP_K})
|
||
|
||
|
||
def init_knowledge_base():
|
||
"""初始化知识库（首次运行时使用）"""
|
||
print("🔧 初始化知识库...")
|
||
kb = KnowledgeBase()
|
||
|
||
# 加载知识库目录中的文档
|
||
documents = kb.load_documents()
|
||
|
||
if documents:
|
||
kb.add_documents(documents)
|
||
print(f"✅ 知识库初始化完成，共 {len(documents)} 个文档")
|
||
else:
|
||
print("⚠️ 没有找到文档，请在 knowledge_base/ 目录添加文档后重新运行")
|
||
|
||
return kb
|
||
|
||
|
||
# 全局知识库实例（单例模式）
|
||
_knowledge_base = None
|
||
|
||
def get_knowledge_base() -> KnowledgeBase:
|
||
"""获取知识库单例"""
|
||
global _knowledge_base
|
||
if _knowledge_base is None:
|
||
_knowledge_base = KnowledgeBase()
|
||
return _knowledge_base
|