73 lines
3.5 KiB
Python
73 lines
3.5 KiB
Python
"""
|
|
init_knowledge_base.py — скрипт инициализации базы знаний.
|
|
|
|
Загружает все текстовые документы из указанной директории
|
|
в векторное хранилище Qdrant.
|
|
|
|
Использование:
|
|
python init_knowledge_base.py [директория]
|
|
|
|
По умолчанию — папка ./docs
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from vector_store import init_collection, add_documents
|
|
|
|
# Поддерживаемые расширения
|
|
TEXT_EXTENSIONS = {".txt", ".md", ".rst", ".csv", ".json"}
|
|
|
|
|
|
def load_directory(docs_dir: str) -> None:
|
|
"""
|
|
Рекурсивно обходит директорию и загружает все текстовые файлы в Qdrant.
|
|
"""
|
|
base = Path(docs_dir)
|
|
if not base.exists():
|
|
print(f"[init] Директория «{docs_dir}» не найдена. Создаю пример.")
|
|
base.mkdir(parents=True, exist_ok=True)
|
|
# Создаём демонстрационный файл
|
|
sample = base / "sample.txt"
|
|
sample.write_text(
|
|
"LangChain — это фреймворк для построения приложений на основе больших языковых моделей. "
|
|
"Он предоставляет инструменты для цепочек, агентов, памяти и интеграции с внешними источниками данных.\n\n"
|
|
"Qdrant — это высокопроизводительная векторная база данных с открытым исходным кодом, "
|
|
"разработанная специально для хранения и поиска плотных векторов. "
|
|
"Поддерживает фильтрацию по метаданным и масштабируется горизонтально.\n\n"
|
|
"Ollama позволяет запускать большие языковые модели локально на вашем компьютере "
|
|
"без подключения к внешним API. Поддерживает модели llama3, mistral, gemma и другие.",
|
|
encoding="utf-8",
|
|
)
|
|
print(f"[init] Создан демонстрационный файл: {sample}")
|
|
|
|
files = [f for f in base.rglob("*") if f.suffix.lower() in TEXT_EXTENSIONS]
|
|
|
|
if not files:
|
|
print(f"[init] В директории «{docs_dir}» нет поддерживаемых файлов.")
|
|
return
|
|
|
|
print(f"[init] Найдено файлов: {len(files)}")
|
|
print("[init] Инициализация коллекции Qdrant...")
|
|
init_collection()
|
|
|
|
total_chunks = 0
|
|
for file in files:
|
|
try:
|
|
content = file.read_text(encoding="utf-8", errors="ignore")
|
|
if not content.strip():
|
|
print(f" ⚠ Пропущен пустой файл: {file.name}")
|
|
continue
|
|
n = add_documents(content, title=file.name)
|
|
total_chunks += n
|
|
print(f" ✓ {file.name} → {n} чанков")
|
|
except Exception as e:
|
|
print(f" ✗ Ошибка при обработке {file.name}: {e}")
|
|
|
|
print(f"\n[init] Загрузка завершена. Всего добавлено чанков: {total_chunks}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
directory = sys.argv[1] if len(sys.argv) > 1 else "./docs"
|
|
load_directory(directory) |