feat: Add ability to send directories with data to cognee

Add ability to send data directories to cognee

Feature COG-656
This commit is contained in:
Igor Ilic 2024-12-11 14:31:54 +01:00
parent 271834f62e
commit f3ce7be885
2 changed files with 62 additions and 0 deletions

View file

@ -3,3 +3,4 @@ from .save_data_to_storage import save_data_to_storage
from .save_data_item_to_storage import save_data_item_to_storage
from .save_data_item_with_metadata_to_storage import save_data_item_with_metadata_to_storage
from .ingest_data_with_metadata import ingest_data_with_metadata
from .resolve_data_directories import resolve_data_directories

View file

@ -0,0 +1,61 @@
# import os
# from typing import List, Union, BinaryIO
#
# def resolve_data_directories(data: Union[BinaryIO, List[BinaryIO], str, List[str]]):
# # We want to work with lists from now on
# if type(data) is not list:
# data = [data]
#
# # Check if data item in list is a directory
# for item in data:
# if type(item) is str:
# # If it's a directory add all files inside the directory to data list instead
# if os.path.isdir(item):
# pass
import os
from typing import List, Union, BinaryIO
def resolve_data_directories(data: Union[BinaryIO, List[BinaryIO], str, List[str]], include_subdirectories: bool = True):
"""
Resolves directories by replacing them with their contained files.
Args:
data: A single file, directory, or binary stream, or a list of such items.
include_subdirectories: Whether to include files in subdirectories recursively.
Returns:
A list of resolved files and binary streams.
"""
# Ensure `data` is a list
if not isinstance(data, list):
data = [data]
resolved_data = []
for item in data:
if isinstance(item, str): # Check if the item is a path
if os.path.isdir(item): # If it's a directory
if include_subdirectories:
# Recursively add all files in the directory and subdirectories
for root, _, files in os.walk(item):
resolved_data.extend([os.path.join(root, f) for f in files])
else:
# Add all files (not subdirectories) in the directory
resolved_data.extend(
[os.path.join(item, f) for f in os.listdir(item) if os.path.isfile(os.path.join(item, f))]
)
elif os.path.isfile(item): # If it's a file, add it to the resolved_data list
resolved_data.append(item)
else:
raise ValueError(f"Path '{item}' is neither a file nor a directory.")
elif isinstance(item, BinaryIO): # If it's a binary stream, add it directly
resolved_data.append(item)
else:
raise TypeError(f"Unsupported type: {type(item)}. Expected str or BinaryIO.")
return resolved_data
# Example usage:
# files = resolve_data_directories(["/path/to/dir", "/path/to/file.txt"], include_subdirectories=True)
# print(files)