diff --git a/cognee/infrastructure/files/storage/S3FileStorage.py b/cognee/infrastructure/files/storage/S3FileStorage.py index 6218d6240..078d5fe2a 100644 --- a/cognee/infrastructure/files/storage/S3FileStorage.py +++ b/cognee/infrastructure/files/storage/S3FileStorage.py @@ -21,10 +21,15 @@ class S3FileStorage(Storage): def __init__(self, storage_path: str): self.storage_path = storage_path s3_config = get_s3_config() - if s3_config.aws_access_key_id is not None and s3_config.aws_secret_access_key is not None: + if ( + s3_config.aws_access_key_id is not None + and s3_config.aws_secret_access_key is not None + and s3_config.aws_session_token is not None + ): self.s3 = s3fs.S3FileSystem( key=s3_config.aws_access_key_id, secret=s3_config.aws_secret_access_key, + token=s3_config.aws_session_token, anon=False, endpoint_url=s3_config.aws_endpoint_url, client_kwargs={"region_name": s3_config.aws_region}, diff --git a/cognee/infrastructure/files/storage/s3_config.py b/cognee/infrastructure/files/storage/s3_config.py index 0b9372b7e..cefe5cd2f 100644 --- a/cognee/infrastructure/files/storage/s3_config.py +++ b/cognee/infrastructure/files/storage/s3_config.py @@ -8,6 +8,7 @@ class S3Config(BaseSettings): aws_endpoint_url: Optional[str] = None aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None + aws_session_token: Optional[str] = None model_config = SettingsConfigDict(env_file=".env", extra="allow") diff --git a/cognee/tasks/ingestion/resolve_data_directories.py b/cognee/tasks/ingestion/resolve_data_directories.py index 1d3124a0c..aa2f95303 100644 --- a/cognee/tasks/ingestion/resolve_data_directories.py +++ b/cognee/tasks/ingestion/resolve_data_directories.py @@ -32,7 +32,10 @@ async def resolve_data_directories( import s3fs fs = s3fs.S3FileSystem( - key=s3_config.aws_access_key_id, secret=s3_config.aws_secret_access_key, anon=False + key=s3_config.aws_access_key_id, + secret=s3_config.aws_secret_access_key, + token=s3_config.aws_session_token, + anon=False, ) for item in data: diff --git a/distributed/pyproject.toml b/distributed/pyproject.toml index fc04c6fcc..f61157236 100644 --- a/distributed/pyproject.toml +++ b/distributed/pyproject.toml @@ -40,7 +40,7 @@ dependencies = [ "aiohttp>=3.11.14,<4.0.0", "aiofiles>=23.2.1,<24.0.0", "rdflib>=7.1.4,<7.2.0", - "pypdf>=4.1.0,<6.0.0", + "pypdf>=4.1.0,<7.0.0", "jinja2>=3.1.3,<4", "matplotlib>=3.8.3,<4", "networkx>=3.4.2,<4", diff --git a/pyproject.toml b/pyproject.toml index ece238338..5e825713e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ dependencies = [ "aiohttp>=3.11.14,<4.0.0", "aiofiles>=23.2.1,<24.0.0", "rdflib>=7.1.4,<7.2.0", - "pypdf>=4.1.0,<6.0.0", + "pypdf>=4.1.0,<7.0.0", "jinja2>=3.1.3,<4", "matplotlib>=3.8.3,<4", "networkx>=3.4.2,<4",