Skip to content

Commit

Permalink
For the users who have long conversation histories and a lot of items…
Browse files Browse the repository at this point in the history
… in sessions table, all the records were not displayed on the UI.

DynamoDB query operation was not returning all the records because the total size of records was more than 1MB and in such cases, DynamoDB query operation does not handle pagination automatically.

As a fix, added code to handle the pagination based on the 'LastEvaluatedKey' returned by DynamoDB.
  • Loading branch information
azaylamba authored and bigadsoleiman committed Jan 9, 2024
1 parent 96a56fe commit a31d7c9
Showing 1 changed file with 24 additions and 7 deletions.
31 changes: 24 additions & 7 deletions lib/shared/layers/python-sdk/python/genai_core/sessions.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,20 +25,37 @@ def get_session(session_id, user_id):


def list_sessions_by_user_id(user_id):
response = {}
items = []
try:
response = table.query(
KeyConditionExpression="UserId = :user_id",
ExpressionAttributeValues={":user_id": user_id},
IndexName=SESSIONS_BY_USER_ID_INDEX_NAME,
)
last_evaluated_key = None
while True:
if last_evaluated_key:
response = table.query(
KeyConditionExpression="UserId = :user_id",
ExpressionAttributeValues={":user_id": user_id},
IndexName=SESSIONS_BY_USER_ID_INDEX_NAME,
ExclusiveStartKey=last_evaluated_key,
)
else:
response = table.query(
KeyConditionExpression="UserId = :user_id",
ExpressionAttributeValues={":user_id": user_id},
IndexName=SESSIONS_BY_USER_ID_INDEX_NAME,
)

items.extend(response.get("Items", []))

last_evaluated_key = response.get("LastEvaluatedKey")
if not last_evaluated_key:
break

except ClientError as error:
if error.response["Error"]["Code"] == "ResourceNotFoundException":
print("No record found for user id: %s", user_id)
else:
print(error)

return response.get("Items", [])
return items


def delete_session(session_id, user_id):
Expand Down

0 comments on commit a31d7c9

Please sign in to comment.