Spaces:
Running
Running
File size: 4,257 Bytes
b655c88 0bcda63 b655c88 0bcda63 b655c88 0bcda63 b655c88 0bcda63 b655c88 0bcda63 b655c88 47738d8 b655c88 47738d8 b655c88 47738d8 b655c88 47738d8 b655c88 0bcda63 47738d8 b655c88 47738d8 b655c88 0bcda63 b655c88 0bcda63 b655c88 0bcda63 b655c88 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """
Module dynamic indexing: Thêm tài liệu PDF mới vào ChromaDB mà không cần rebuild toàn bộ.
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from loader import load_pdf_file
from chunking import chunk_documents
from indexing import get_collection, is_document_indexed, _make_chunk_id
def add_new_documents(documents: list) -> int:
"""
Thêm tài liệu mới vào ChromaDB.
Args:
documents: danh sách dict {"content": "...", "source": "..."}
Returns:
Số chunks đã thêm
"""
if not documents:
print("❌ Không có tài liệu mới để thêm!")
return 0
chunks = chunk_documents(documents)
if not chunks:
print("❌ Không tạo được chunks từ tài liệu!")
return 0
collection = get_collection()
from indexing import get_embedding_function
embedding_fn = get_embedding_function()
embedding_fn.set_mode("passage")
documents_list = []
metadatas_list = []
ids_list = []
per_source_counter: dict[str, int] = {}
for chunk in chunks:
content = chunk.get("content", "").strip()
if not content:
continue
metadata = chunk.get("metadata", {})
clean_metadata = {}
for k, v in metadata.items():
if isinstance(v, (str, int, float, bool)):
clean_metadata[k] = v
else:
clean_metadata[k] = str(v)
source = clean_metadata.get("source", "unknown")
idx = per_source_counter.get(source, 0)
per_source_counter[source] = idx + 1
documents_list.append(content)
metadatas_list.append(clean_metadata)
ids_list.append(_make_chunk_id(source, idx))
if not documents_list:
print("❌ Không có nội dung hợp lệ để thêm!")
return 0
batch_size = 500
total = len(documents_list)
skipped_existing = 0
inserted_new = 0
for start in range(0, total, batch_size):
end = min(start + batch_size, total)
batch_ids = ids_list[start:end]
existing = collection.get(ids=batch_ids, include=[])
existing_ids = set(existing.get("ids", []) if existing else [])
filtered_docs = []
filtered_metas = []
filtered_ids = []
for doc, meta, chunk_id in zip(
documents_list[start:end],
metadatas_list[start:end],
batch_ids,
):
if chunk_id in existing_ids:
skipped_existing += 1
continue
filtered_docs.append(doc)
filtered_metas.append(meta)
filtered_ids.append(chunk_id)
if not filtered_ids:
continue
collection.upsert(
documents=filtered_docs,
metadatas=filtered_metas,
ids=filtered_ids,
)
inserted_new += len(filtered_ids)
embedding_fn.set_mode("query")
print(f"✅ Đã thêm {inserted_new} chunks mới vào ChromaDB")
if skipped_existing:
print(f"⏭️ Bỏ qua {skipped_existing} chunks đã tồn tại")
print(f"📊 Tổng chunks hiện tại: {collection.count()}")
return inserted_new
def add_pdf_file(filepath: str) -> int:
"""
Thêm một file PDF mới vào ChromaDB.
Nếu file đã được index trước đó thì bỏ qua.
Args:
filepath: đường dẫn đến file PDF
Returns:
Số chunks đã thêm (0 nếu đã index hoặc lỗi)
"""
if not os.path.exists(filepath):
print(f"❌ File không tồn tại: {filepath}")
return 0
filename = os.path.basename(filepath)
if is_document_indexed(filename):
print(f"⏭️ Bỏ qua '{filename}' -- đã được index trước đó")
return 0
try:
content = load_pdf_file(filepath)
except Exception as e:
print(f"❌ Lỗi đọc PDF: {e}")
return 0
if not content or not content.strip():
print(f"⚠️ File PDF rỗng hoặc không trích xuất được text: {filepath}")
return 0
documents = [{"content": content.strip(), "source": filename}]
return add_new_documents(documents) |