From a31d7c9b52c06cd55ab387f358ffb45c7db2076e Mon Sep 17 00:00:00 2001
From: Ajay Lamba <azaystudy@gmail.com>
Date: Wed, 27 Dec 2023 16:22:04 +0530
Subject: [PATCH] For the users who have long conversation histories and a lot
 of items 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.
---
 .../python-sdk/python/genai_core/sessions.py  | 31 ++++++++++++++-----
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/lib/shared/layers/python-sdk/python/genai_core/sessions.py b/lib/shared/layers/python-sdk/python/genai_core/sessions.py
index f5accdffd..8933bdd33 100644
--- a/lib/shared/layers/python-sdk/python/genai_core/sessions.py
+++ b/lib/shared/layers/python-sdk/python/genai_core/sessions.py
@@ -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):