mirror of
https://github.com/imartinez/privateGPT.git
synced 2025-04-28 11:36:25 +00:00
022bd718e3
* Remove all global settings state * chore: remove autogenerated class * chore: cleanup * chore: merge conflicts
90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
import argparse
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from private_gpt.di import global_injector
|
|
from private_gpt.server.ingest.ingest_service import IngestService
|
|
from private_gpt.server.ingest.ingest_watcher import IngestWatcher
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ingest_service = global_injector.get(IngestService)
|
|
|
|
parser = argparse.ArgumentParser(prog="ingest_folder.py")
|
|
parser.add_argument("folder", help="Folder to ingest")
|
|
parser.add_argument(
|
|
"--watch",
|
|
help="Watch for changes",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
)
|
|
parser.add_argument(
|
|
"--log-file",
|
|
help="Optional path to a log file. If provided, logs will be written to this file.",
|
|
type=str,
|
|
default=None,
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Set up logging to a file if a path is provided
|
|
if args.log_file:
|
|
file_handler = logging.FileHandler(args.log_file, mode="a")
|
|
file_handler.setFormatter(
|
|
logging.Formatter(
|
|
"[%(asctime)s.%(msecs)03d] [%(levelname)s] %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
)
|
|
logger.addHandler(file_handler)
|
|
|
|
|
|
total_documents = 0
|
|
current_document_count = 0
|
|
|
|
|
|
def count_documents(folder_path: Path) -> None:
|
|
global total_documents
|
|
for file_path in folder_path.iterdir():
|
|
if file_path.is_file():
|
|
total_documents += 1
|
|
elif file_path.is_dir():
|
|
count_documents(file_path)
|
|
|
|
|
|
def _recursive_ingest_folder(folder_path: Path) -> None:
|
|
global current_document_count, total_documents
|
|
for file_path in folder_path.iterdir():
|
|
if file_path.is_file():
|
|
current_document_count += 1
|
|
progress_msg = f"Document {current_document_count} of {total_documents} ({(current_document_count / total_documents) * 100:.2f}%)"
|
|
logger.info(progress_msg)
|
|
_do_ingest(file_path)
|
|
elif file_path.is_dir():
|
|
_recursive_ingest_folder(file_path)
|
|
|
|
|
|
def _do_ingest(changed_path: Path) -> None:
|
|
try:
|
|
if changed_path.exists():
|
|
logger.info(f"Started ingesting {changed_path}")
|
|
ingest_service.ingest(changed_path.name, changed_path)
|
|
logger.info(f"Completed ingesting {changed_path}")
|
|
except Exception:
|
|
logger.exception(
|
|
f"Failed to ingest document: {changed_path}, find the exception attached"
|
|
)
|
|
|
|
|
|
path = Path(args.folder)
|
|
if not path.exists():
|
|
raise ValueError(f"Path {args.folder} does not exist")
|
|
|
|
# Count total documents before ingestion
|
|
count_documents(path)
|
|
|
|
_recursive_ingest_folder(path)
|
|
if args.watch:
|
|
logger.info(f"Watching {args.folder} for changes, press Ctrl+C to stop...")
|
|
watcher = IngestWatcher(args.folder, _do_ingest)
|
|
watcher.start()
|