32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
import os
|
|
from dotenv import load_dotenv
|
|
from vectorstore import create_vectorstore, load_documents
|
|
|
|
def main():
|
|
load_dotenv() # загружаем TAVILY_API_KEY и другие переменные
|
|
persist_dir = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")
|
|
vectorstore = create_vectorstore(persist_directory=persist_dir)
|
|
|
|
# Проверяем, пуста ли база данных
|
|
try:
|
|
if hasattr(vectorstore, "_collection"):
|
|
doc_count = vectorstore._collection.count()
|
|
else:
|
|
# запасной способ: пробуем выполнить пустой поиск
|
|
doc_count = len(vectorstore.similarity_search("test", k=1))
|
|
except Exception:
|
|
doc_count = 0
|
|
|
|
if doc_count == 0:
|
|
docs_dir = os.path.join(os.path.dirname(__file__), "documents")
|
|
if not os.path.isdir(docs_dir):
|
|
print(f"Каталог {docs_dir} не найден.")
|
|
return
|
|
print(f"Загрузка документов из {docs_dir} в ChromaDB ({persist_dir})...")
|
|
load_documents(docs_dir, vectorstore)
|
|
print("Загрузка завершена.")
|
|
else:
|
|
print(f"База данных в {persist_dir} уже содержит {doc_count} документов. Загрузка пропущена.")
|
|
|
|
if __name__ == "__main__":
|
|
main() |