Enable AWS SSO credentials use instead of just IAM user

[prparmar77]

Problem Description

Provide a clear and concise explanation of the issue you're facing.
Use of AWS SSO credentials to create storage in DynamoDB table may fail as code is not able to validate temporary credentials AWS_SESSION_TOKEN issued by AWS SSO.

Proposed Solution

Explain your ideal solution. How should this feature work?

To create DynamoDb table to store user sessions it should be able to use AWS_SESSION_TOKEN which is gathered by user from console or so after being authenticated and granted with permissions

Alternatives Considered

Have you found any workarounds or alternative solutions?
Share other approaches you've tried or thought about.

Yes, the file $HOME/phidata/phienv/lib/python3.13/site-packages/phi/storage/agent/dynamodb.py should be modified to include 'aws_session_token' as optional str

Existing --
def init(
self,
table_name: str,
region_name: Optional[str] = None,
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
endpoint_url: Optional[str] = None,
create_table_if_not_exists: bool = True,
):
"""
Initialize the DynamoDbAgentStorage.

    Args:
        table_name (str): The name of the DynamoDB table.
        region_name (Optional[str]): AWS region name.
        aws_access_key_id (Optional[str]): AWS access key ID.
        aws_secret_access_key (Optional[str]): AWS secret access key.
        endpoint_url (Optional[str]): The complete URL to use for the constructed client.
        create_table_if_not_exists (bool): Whether to create the table if it does not exist.
    """
    self.table_name = table_name
    self.region_name = region_name
    self.endpoint_url = endpoint_url
    self.aws_access_key_id = aws_access_key_id
    self.aws_secret_access_key = aws_secret_access_key
    self.create_table_if_not_exists = create_table_if_not_exists

    # Initialize DynamoDB resource
    self.dynamodb = boto3.resource(
        "dynamodb",
        region_name=self.region_name,
        aws_access_key_id=self.aws_access_key_id,
        aws_secret_access_key=self.aws_secret_access_key,
        endpoint_url=self.endpoint_url,
    )

Updated ---

def init(
self,
table_name: str,
region_name: Optional[str] = None,
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
aws_session_token: Optional[str] = None,
endpoint_url: Optional[str] = None,
create_table_if_not_exists: bool = True,
):
"""
Initialize the DynamoDbAgentStorage.

    Args:
        table_name (str): The name of the DynamoDB table.
        region_name (Optional[str]): AWS region name.
        aws_access_key_id (Optional[str]): AWS access key ID.
        aws_secret_access_key (Optional[str]): AWS secret access key.
        aws_session_token (Optional[str]): AWS Access Token
        endpoint_url (Optional[str]): The complete URL to use for the constructed client.
        create_table_if_not_exists (bool): Whether to create the table if it does not exist.
    """
    self.table_name = table_name
    self.region_name = region_name
    self.endpoint_url = endpoint_url
    self.aws_access_key_id = aws_access_key_id
    self.aws_secret_access_key = aws_secret_access_key
    self.aws_session_token = aws_session_token
    self.create_table_if_not_exists = create_table_if_not_exists

    # Initialize DynamoDB resource
    self.dynamodb = boto3.resource(
        "dynamodb",
        region_name=self.region_name,
        aws_access_key_id=self.aws_access_key_id,
        aws_secret_access_key=self.aws_secret_access_key,
        aws_session_token=self.aws_session_token,
        endpoint_url=self.endpoint_url,
    )

Since the strings are optional it will work for both IAM user and AWS SSO authenticated user token

Additional context

Include any extra information that might be helpful, such as:

• Screenshots
• Examples of similar features elsewhere
• Any relevant links or references
• https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html

Would you like to work on this?

We welcome contributions! Let us know if you’d like to help implement this feature.
[ ] I’m open to collaborating but need guidance.

[Ansub]

Thanks for the suggestion, @prparmar77! 🙌

Supporting AWS SSO credentials alongside IAM users makes sense. The team will review this and assess the best approach. We’ll keep you posted! 🚀