close session after use

llm-server/models/repository/chat_history_repo.py

@@ -189,47 +189,46 @@ def get_chat_history_for_retrieval_chain(
 
 def get_unique_sessions_with_first_message_by_bot_id(
     bot_id: str, limit: int = 20, offset: int = 0
-) -> list[Dict[str, object]]:
-    # If a session is not provided, create a new one
-    session = Session()
-
-    # Use distinct to get unique session_ids
-    unique_session_ids = (
-        session.query(distinct(ChatHistory.session_id), ChatHistory.id)
-        .filter_by(chatbot_id=bot_id)
-        .order_by(ChatHistory.id.desc())
-        .limit(limit)
-        .offset(offset)
-        .all()
-    )
-
-    result_list = []
-
-    for session_id in unique_session_ids:
-        # Get the first message in each session
-        first_message = (
-            session.query(ChatHistory)
-            .filter_by(chatbot_id=bot_id, session_id=session_id[0])
-            .order_by(ChatHistory.id.asc())
-            .first()
+) -> List[Dict[str, object]]:
+    # Using a context manager to automatically close the session
+    with Session() as session:
+        # Use distinct to get unique session_ids
+        unique_session_ids = (
+            session.query(distinct(ChatHistory.session_id))
+            .filter_by(chatbot_id=bot_id)
+            # fix the next one
+            # .order_by(ChatHistory.id.desc())
+            .limit(limit)
+            .offset(offset)
+            .all()
         )
 
-        # Convert ChatHistory object to a dictionary
-        if first_message:
-            first_message_dict = {
-                column.key: getattr(first_message, column.key)
-                for column in class_mapper(ChatHistory).mapped_table.columns
-            }
-        else:
-            first_message_dict = None
+        result_list = []
+
+        for session_id in unique_session_ids:
+            # Get the first message in each session
+            first_message = (
+                session.query(ChatHistory)
+                .filter_by(chatbot_id=bot_id, session_id=session_id[0])
+                .order_by(ChatHistory.id.asc())
+                .first()
+            )
+
+            # Convert ChatHistory object to a dictionary
+            if first_message:
+                first_message_dict = {
+                    column.key: getattr(first_message, column.key)
+                    for column in class_mapper(ChatHistory).mapped_table.columns
+                }
+            else:
+                first_message_dict = None
 
-        # Create a dictionary with session_id and first_message
-        result_dict: Dict[str, object] = {
-            "session_id": session_id[0],
-            "first_message": first_message_dict,
-        }
+            # Create a dictionary with session_id and first_message
+            result_dict: Dict[str, object] = {
+                "session_id": session_id[0],
+                "first_message": first_message_dict,
+            }
 
-        result_list: List[Dict[str, object]] = []
-        result_list.append(result_dict)
+            result_list.append(result_dict)
 
-    return result_list
+    return result_list