MojoSamurai/AutoGPT

1import argparse2import logging3 4from autogpt.commands.file_operations import ingest_file, search_files5from autogpt.config import Config6from autogpt.memory import get_memory7 8cfg = Config()9 10 11def configure_logging():12    logging.basicConfig(13        filename="log-ingestion.txt",14        filemode="a",15        format="%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s",16        datefmt="%H:%M:%S",17        level=logging.DEBUG,18    )19    return logging.getLogger("AutoGPT-Ingestion")20 21 22def ingest_directory(directory, memory, args):23    """24    Ingest all files in a directory by calling the ingest_file function for each file.25 26    :param directory: The directory containing the files to ingest27    :param memory: An object with an add() method to store the chunks in memory28    """29    try:30        files = search_files(directory)31        for file in files:32            ingest_file(file, memory, args.max_length, args.overlap)33    except Exception as e:34        print(f"Error while ingesting directory '{directory}': {str(e)}")35 36 37def main() -> None:38    logger = configure_logging()39 40    parser = argparse.ArgumentParser(41        description="Ingest a file or a directory with multiple files into memory. "42        "Make sure to set your .env before running this script."43    )44    group = parser.add_mutually_exclusive_group(required=True)45    group.add_argument("--file", type=str, help="The file to ingest.")46    group.add_argument(47        "--dir", type=str, help="The directory containing the files to ingest."48    )49    parser.add_argument(50        "--init",51        action="store_true",52        help="Init the memory and wipe its content (default: False)",53        default=False,54    )55    parser.add_argument(56        "--overlap",57        type=int,58        help="The overlap size between chunks when ingesting files (default: 200)",59        default=200,60    )61    parser.add_argument(62        "--max_length",63        type=int,64        help="The max_length of each chunk when ingesting files (default: 4000)",65        default=4000,66    )67 68    args = parser.parse_args()69 70    # Initialize memory71    memory = get_memory(cfg, init=args.init)72    print("Using memory of type: " + memory.__class__.__name__)73 74    if args.file:75        try:76            ingest_file(args.file, memory, args.max_length, args.overlap)77            print(f"File '{args.file}' ingested successfully.")78        except Exception as e:79            logger.error(f"Error while ingesting file '{args.file}': {str(e)}")80            print(f"Error while ingesting file '{args.file}': {str(e)}")81    elif args.dir:82        try:83            ingest_directory(args.dir, memory, args)84            print(f"Directory '{args.dir}' ingested successfully.")85        except Exception as e:86            logger.error(f"Error while ingesting directory '{args.dir}': {str(e)}")87            print(f"Error while ingesting directory '{args.dir}': {str(e)}")88    else:89        print(90            "Please provide either a file path (--file) or a directory name (--dir)"91            " inside the auto_gpt_workspace directory as input."92        )93 94 95if __name__ == "__main__":96    main()97