194 lines
8.4 KiB
Python
194 lines
8.4 KiB
Python
# 该代码用于将本地知识库中的文档导入到ChromaDB中，并使用远程嵌入模型进行向量化
|
||
import os
|
||
from pathlib import Path
|
||
import chromadb
|
||
# from chromadb.utils import embedding_functions - 不再需要
|
||
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings, Embeddable
|
||
from unstructured.partition.auto import partition
|
||
from rich.progress import track
|
||
import logging
|
||
import requests # 导入requests
|
||
import json # 导入json模块
|
||
|
||
# 配置日志
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
|
||
# --- 配置 ---
|
||
# 获取脚本所在目录，确保路径的正确性
|
||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||
KNOWLEDGE_BASE_DIR = SCRIPT_DIR / "knowledge_base"
|
||
VECTOR_STORE_DIR = SCRIPT_DIR / "vector_store"
|
||
COLLECTION_NAME = "drone_docs"
|
||
# EMBEDDING_MODEL_NAME = "bge-small-zh-v1.5" # 不再需要，模型名在函数内部处理
|
||
|
||
# --- 自定义远程嵌入函数 ---
|
||
class RemoteEmbeddingFunction(EmbeddingFunction[Embeddable]):
|
||
"""
|
||
一个使用远程、兼容OpenAI API的嵌入服务的嵌入函数。
|
||
"""
|
||
def __init__(self, api_url: str):
|
||
self._api_url = api_url
|
||
logging.info(f"自定义嵌入函数已初始化，将连接到: {self._api_url}")
|
||
|
||
def __call__(self, input: Embeddable) -> Embeddings:
|
||
"""
|
||
对输入的文档进行嵌入。
|
||
"""
|
||
# 我们的服务只能处理文本，所以检查输入是否为字符串列表
|
||
if not isinstance(input, list) or not all(isinstance(doc, str) for doc in input):
|
||
logging.error("此嵌入函数仅支持字符串列表（文档）作为输入。")
|
||
return []
|
||
|
||
try:
|
||
# 移除 "model" 参数，因为embedding服务可能不需要它
|
||
response = requests.post(
|
||
self._api_url,
|
||
json={"input": input},
|
||
headers={"Content-Type": "application/json"}
|
||
)
|
||
response.raise_for_status() # 如果请求失败则抛出HTTPError
|
||
|
||
# 按照OpenAI API的格式解析返回的嵌入向量
|
||
data = response.json().get("data", [])
|
||
if not data:
|
||
raise ValueError("API响应中没有找到'data'字段或'data'为空")
|
||
|
||
embeddings = [item['embedding'] for item in data]
|
||
return embeddings
|
||
|
||
except requests.RequestException as e:
|
||
logging.error(f"调用嵌入API失败: {e}")
|
||
# 返回一个空列表或根据需要处理错误
|
||
return []
|
||
except (ValueError, KeyError) as e:
|
||
logging.error(f"解析API响应失败: {e}")
|
||
logging.error(f"收到的响应内容: {response.text}")
|
||
return []
|
||
|
||
|
||
def get_documents(directory: Path):
|
||
"""从知识库目录加载所有文档并进行切分"""
|
||
documents = []
|
||
logging.info(f"从 '{directory}' 加载文档...")
|
||
for file_path in directory.rglob("*"):
|
||
if file_path.is_file() and not file_path.name.startswith('.'):
|
||
try:
|
||
# 对简单文本文件直接读取
|
||
if file_path.suffix in ['.txt', '.md']:
|
||
text = file_path.read_text(encoding='utf-8')
|
||
documents.append({
|
||
"text": text,
|
||
"metadata": {"source": str(file_path.name)}
|
||
})
|
||
logging.info(f"成功处理文本文件: {file_path.name}")
|
||
# 特别处理常规的JSON文件
|
||
elif file_path.suffix == '.json':
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
if 'elements' in data and isinstance(data['elements'], list):
|
||
for element in data['elements']:
|
||
documents.append({
|
||
"text": json.dumps(element, ensure_ascii=False),
|
||
"metadata": {"source": str(file_path.name)}
|
||
})
|
||
logging.info(f"成功处理JSON文件: {file_path.name}, 提取了 {len(data['elements'])} 个元素。")
|
||
else:
|
||
documents.append({
|
||
"text": json.dumps(data, ensure_ascii=False),
|
||
"metadata": {"source": str(file_path.name)}
|
||
})
|
||
logging.info(f"成功处理JSON文件: {file_path.name} (作为单个文档)")
|
||
# 新增：专门处理我们生成的 NDJSON 文件
|
||
elif file_path.suffix == '.ndjson':
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
count = 0
|
||
for line in f:
|
||
try:
|
||
record = json.loads(line)
|
||
if 'text' in record and isinstance(record['text'], str):
|
||
documents.append({
|
||
"text": record['text'],
|
||
"metadata": {"source": str(file_path.name)}
|
||
})
|
||
count += 1
|
||
except json.JSONDecodeError:
|
||
logging.warning(f"跳过无效的JSON行: {line.strip()}")
|
||
if count > 0:
|
||
logging.info(f"成功处理NDJSON文件: {file_path.name}, 提取了 {count} 个文档。")
|
||
# 对其他所有文件类型，使用unstructured
|
||
else:
|
||
elements = partition(filename=str(file_path))
|
||
for element in elements:
|
||
documents.append({
|
||
"text": element.text,
|
||
"metadata": {"source": str(file_path.name)}
|
||
})
|
||
logging.info(f"成功处理文件: {file_path.name} (使用unstructured)")
|
||
except Exception as e:
|
||
logging.error(f"处理文件 {file_path.name} 失败: {e}")
|
||
return documents
|
||
|
||
|
||
def main():
|
||
"""主函数，执行文档入库流程"""
|
||
if not KNOWLEDGE_BASE_DIR.exists():
|
||
KNOWLEDGE_BASE_DIR.mkdir(parents=True)
|
||
logging.warning(f"知识库目录不存在，已自动创建: {KNOWLEDGE_BASE_DIR}")
|
||
logging.warning("请向该目录中添加您的知识文件(如 .txt, .pdf, .md)。")
|
||
return
|
||
|
||
# 1. 加载并切分文档
|
||
docs_to_ingest = get_documents(KNOWLEDGE_BASE_DIR)
|
||
if not docs_to_ingest:
|
||
logging.warning("在知识库中未找到可处理的文档。")
|
||
return
|
||
|
||
# 2. 初始化ChromaDB客户端和远程嵌入函数
|
||
orin_ip = os.getenv("ORIN_IP", "localhost")
|
||
embedding_api_url = f"http://{orin_ip}:8090/v1/embeddings"
|
||
|
||
logging.info(f"正在初始化远程嵌入函数，目标服务地址: {embedding_api_url}")
|
||
embedding_func = RemoteEmbeddingFunction(api_url=embedding_api_url)
|
||
|
||
client = chromadb.PersistentClient(path=str(VECTOR_STORE_DIR))
|
||
|
||
# 3. 创建或获取集合
|
||
logging.info(f"正在访问ChromaDB集合: {COLLECTION_NAME}")
|
||
collection = client.get_or_create_collection(
|
||
name=COLLECTION_NAME,
|
||
embedding_function=embedding_func
|
||
)
|
||
|
||
# 4. 将文档向量化并存入数据库
|
||
logging.info(f"开始将 {len(docs_to_ingest)} 个文档块入库...")
|
||
|
||
# 为了避免重复添加，可以先检查
|
||
# (这里为了简单，我们每次都重新添加，生产环境需要更复杂的逻辑)
|
||
|
||
doc_texts = [doc['text'] for doc in docs_to_ingest]
|
||
metadatas = [doc['metadata'] for doc in docs_to_ingest]
|
||
ids = [f"doc_{KNOWLEDGE_BASE_DIR.name}_{i}" for i in range(len(doc_texts))]
|
||
|
||
try:
|
||
# ChromaDB的add方法会自动处理嵌入
|
||
collection.add(
|
||
documents=doc_texts,
|
||
metadatas=metadatas,
|
||
ids=ids
|
||
)
|
||
logging.info("所有文档块已成功入库！")
|
||
except Exception as e:
|
||
logging.error(f"向ChromaDB添加文档时出错: {e}")
|
||
|
||
|
||
# 验证一下
|
||
count = collection.count()
|
||
logging.info(f"数据库中现在有 {count} 个条目。")
|
||
|
||
print("\n✅ 数据入库完成！")
|
||
print(f"知识库位于: {KNOWLEDGE_BASE_DIR}")
|
||
print(f"向量数据库位于: {VECTOR_STORE_DIR}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |