|
import argparse |
|
import logging |
|
|
|
from autogpt.commands.file_operations import ingest_file, search_files |
|
from autogpt.config import Config |
|
from autogpt.memory import get_memory |
|
|
|
cfg = Config() |
|
|
|
|
|
def configure_logging(): |
|
logging.basicConfig( |
|
filename="log-ingestion.txt", |
|
filemode="a", |
|
format="%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s", |
|
datefmt="%H:%M:%S", |
|
level=logging.DEBUG, |
|
) |
|
return logging.getLogger("AutoGPT-Ingestion") |
|
|
|
|
|
def ingest_directory(directory, memory, args): |
|
""" |
|
Ingest all files in a directory by calling the ingest_file function for each file. |
|
|
|
:param directory: The directory containing the files to ingest |
|
:param memory: An object with an add() method to store the chunks in memory |
|
""" |
|
try: |
|
files = search_files(directory) |
|
for file in files: |
|
ingest_file(file, memory, args.max_length, args.overlap) |
|
except Exception as e: |
|
print(f"Error while ingesting directory '{directory}': {str(e)}") |
|
|
|
|
|
def main() -> None: |
|
logger = configure_logging() |
|
|
|
parser = argparse.ArgumentParser( |
|
description="Ingest a file or a directory with multiple files into memory. " |
|
"Make sure to set your .env before running this script." |
|
) |
|
group = parser.add_mutually_exclusive_group(required=True) |
|
group.add_argument("--file", type=str, help="The file to ingest.") |
|
group.add_argument( |
|
"--dir", type=str, help="The directory containing the files to ingest." |
|
) |
|
parser.add_argument( |
|
"--init", |
|
action="store_true", |
|
help="Init the memory and wipe its content (default: False)", |
|
default=False, |
|
) |
|
parser.add_argument( |
|
"--overlap", |
|
type=int, |
|
help="The overlap size between chunks when ingesting files (default: 200)", |
|
default=200, |
|
) |
|
parser.add_argument( |
|
"--max_length", |
|
type=int, |
|
help="The max_length of each chunk when ingesting files (default: 4000)", |
|
default=4000, |
|
) |
|
|
|
args = parser.parse_args() |
|
|
|
|
|
memory = get_memory(cfg, init=args.init) |
|
print("Using memory of type: " + memory.__class__.__name__) |
|
|
|
if args.file: |
|
try: |
|
ingest_file(args.file, memory, args.max_length, args.overlap) |
|
print(f"File '{args.file}' ingested successfully.") |
|
except Exception as e: |
|
logger.error(f"Error while ingesting file '{args.file}': {str(e)}") |
|
print(f"Error while ingesting file '{args.file}': {str(e)}") |
|
elif args.dir: |
|
try: |
|
ingest_directory(args.dir, memory, args) |
|
print(f"Directory '{args.dir}' ingested successfully.") |
|
except Exception as e: |
|
logger.error(f"Error while ingesting directory '{args.dir}': {str(e)}") |
|
print(f"Error while ingesting directory '{args.dir}': {str(e)}") |
|
else: |
|
print( |
|
"Please provide either a file path (--file) or a directory name (--dir)" |
|
" inside the auto_gpt_workspace directory as input." |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|