huangfu/DronePlanning
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

AI重构

Browse Source
This commit is contained in:
huangfu
2026-02-20 23:04:09 +08:00
parent 43d69d5a99
commit 5d8412bcb6
244 changed files with 9304 additions and 1463 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
tools/rag/ingest.py
View File

@@ -5,7 +5,6 @@ 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模块
@@ -19,6 +18,11 @@ SCRIPT_DIR = Path(__file__).resolve().parent
KNOWLEDGE_BASE_DIR = SCRIPT_DIR / "knowledge_base"
VECTOR_STORE_DIR = SCRIPT_DIR / "vector_store"
COLLECTION_NAME = "drone_docs"
COLLECTIONS_BY_KB_TYPE = {
"location": "location_kb",
"pattern": "pattern_kb",
"rules": "rules_kb",
}
# EMBEDDING_MODEL_NAME = "bge-small-zh-v1.5" # 不再需要，模型名在函数内部处理
# --- 自定义远程嵌入函数 ---
@@ -66,6 +70,18 @@ class RemoteEmbeddingFunction(EmbeddingFunction[Embeddable]):
return []
def _infer_kb_type(file_path: Path, base_dir: Path) -> str:
try:
rel = file_path.relative_to(base_dir)
except ValueError:
return "location"
if len(rel.parts) > 1:
first = rel.parts[0].lower()
if first in COLLECTIONS_BY_KB_TYPE:
return first
return "location"
def get_documents(directory: Path):
"""从知识库目录加载所有文档并进行切分"""
documents = []
@@ -74,11 +90,12 @@ def get_documents(directory: Path):
if file_path.is_file() and not file_path.name.startswith('.'):
try:
# 对简单文本文件直接读取
kb_type = _infer_kb_type(file_path, directory)
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)}
"metadata": {"source": str(file_path.name), "kb_type": kb_type}
})
logging.info(f"成功处理文本文件: {file_path.name}")
# 特别处理常规的JSON文件
@@ -89,13 +106,13 @@ def get_documents(directory: Path):
for element in data['elements']:
documents.append({
"text": json.dumps(element, ensure_ascii=False),
"metadata": {"source": str(file_path.name)}
"metadata": {"source": str(file_path.name), "kb_type": kb_type}
})
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)}
"metadata": {"source": str(file_path.name), "kb_type": kb_type}
})
logging.info(f"成功处理JSON文件: {file_path.name} (作为单个文档)")
# 新增：专门处理我们生成的 NDJSON 文件
@@ -115,7 +132,7 @@ def get_documents(directory: Path):
text = json.dumps(record, ensure_ascii=False)
documents.append({
"text": text,
"metadata": {"source": str(file_path.name)}
"metadata": {"source": str(file_path.name), "kb_type": kb_type}
})
count += 1
except json.JSONDecodeError:
@@ -128,7 +145,7 @@ def get_documents(directory: Path):
for element in elements:
documents.append({
"text": element.text,
"metadata": {"source": str(file_path.name)}
"metadata": {"source": str(file_path.name), "kb_type": kb_type}
})
logging.info(f"成功处理文件: {file_path.name} (使用unstructured)")
except Exception as e:
@@ -159,38 +176,48 @@ def main():
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
)
# 3. 创建或获取集合（兼容旧集合 + 新多知识库集合）
collections = {
"fallback": client.get_or_create_collection(name=COLLECTION_NAME, embedding_function=embedding_func)
}
for kb_type, coll_name in COLLECTIONS_BY_KB_TYPE.items():
collections[kb_type] = client.get_or_create_collection(name=coll_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))]
grouped_docs = {}
for doc in docs_to_ingest:
kb_type = doc["metadata"].get("kb_type", "location")
grouped_docs.setdefault(kb_type, []).append(doc)
# 先入旧集合，保持兼容
all_texts = [doc['text'] for doc in docs_to_ingest]
all_metadatas = [doc['metadata'] for doc in docs_to_ingest]
all_ids = [f"doc_{KNOWLEDGE_BASE_DIR.name}_{i}" for i in range(len(all_texts))]
try:
# ChromaDB的add方法会自动处理嵌入
collection.add(
documents=doc_texts,
metadatas=metadatas,
ids=ids
)
logging.info("所有文档块已成功入库！")
collections["fallback"].add(documents=all_texts, metadatas=all_metadatas, ids=all_ids)
logging.info("兼容集合 drone_docs 入库完成。")
except Exception as e:
logging.error(f"向ChromaDB添加文档时出错: {e}")
logging.error(f"向兼容集合添加文档时出错: {e}")
for kb_type, docs in grouped_docs.items():
if kb_type not in collections:
continue
doc_texts = [doc["text"] for doc in docs]
metadatas = [doc["metadata"] for doc in docs]
ids = [f"doc_{kb_type}_{i}" for i in range(len(doc_texts))]
try:
collections[kb_type].add(documents=doc_texts, metadatas=metadatas, ids=ids)
logging.info(f"{kb_type} 集合入库完成，条目数: {len(doc_texts)}")
except Exception as e:
logging.error(f"向 {kb_type} 集合添加文档时出错: {e}")
# 验证一下
count = collection.count()
logging.info(f"数据库中现在有 {count} 个条目。")
for name, collection in collections.items():
try:
count = collection.count()
logging.info(f"集合 {name} 当前条目数: {count}")
except Exception:
pass
print("\n✅ 数据入库完成！")
print(f"知识库位于: {KNOWLEDGE_BASE_DIR}")