diff --git a/init_db.py b/init_db.py index 1241777..2cc23ee 100644 --- a/init_db.py +++ b/init_db.py @@ -6,13 +6,27 @@ def main(): load_dotenv() # загружаем TAVILY_API_KEY и другие переменные persist_dir = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db") vectorstore = create_vectorstore(persist_directory=persist_dir) - 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("Загрузка завершена.") + + # Проверяем, пуста ли база данных + 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() \ No newline at end of file