Merge pull request #262 from Pythagora-io/optimizations

optimize get_fe_states

core/cli/main.py

@@ -209,7 +209,7 @@ async def run_pythagora_session(sm: StateManager, ui: UIBase, args: Namespace):
             return False
 
         # SPECIFICATION
-        fe_states = await sm.get_fe_states()
+        fe_states = await sm.get_fe_states(limit=10)
         be_back_logs, last_task_in_db = await sm.get_be_back_logs()
 
         if sm.current_state.specification and sm.current_state.specification.original_description:

core/db/models/project_state.py

@@ -903,7 +903,9 @@ class ProjectState(Base):
             return results
 
     @staticmethod
-    async def get_fe_states(session: "AsyncSession", branch_id: UUID) -> Optional["ProjectState"]:
+    async def get_fe_states(
+        session: "AsyncSession", branch_id: UUID, limit: Optional[int] = None
+    ) -> Optional["ProjectState"]:
         query = select(ProjectState).where(
             and_(
                 ProjectState.branch_id == branch_id,
@@ -931,16 +933,26 @@ class ProjectState(Base):
         result = await session.execute(query)
         fe_end = result.scalars().one_or_none()
 
-        query = select(ProjectState).where(
-            and_(
-                ProjectState.branch_id == branch_id,
-                ProjectState.step_index >= fe_start.step_index,
-                ProjectState.step_index <= fe_end.step_index,
+        query = (
+            select(ProjectState)
+            .where(
+                and_(
+                    ProjectState.branch_id == branch_id,
+                    ProjectState.step_index >= fe_start.step_index,
+                    ProjectState.step_index <= fe_end.step_index,
+                )
             )
+            .order_by(ProjectState.step_index.desc())
         )
 
+        if limit:
+            query = query.limit(limit)
+
         results = await session.execute(query)
-        return results.scalars().all()
+        states = results.scalars().all()
+
+        # Since we ordered by step_index desc and limited, we need to reverse to get chronological order
+        return list(reversed(states))
 
     @staticmethod
     def get_epic_task_number(state, current_task) -> (int, int):

core/state/state_manager.py

@@ -159,8 +159,8 @@ class StateManager:
             self.current_session, self.current_state.branch_id, start_state_id, end_state_id, limit
         )
 
-    async def get_fe_states(self) -> Optional[ProjectState]:
-        return await ProjectState.get_fe_states(self.current_session, self.current_state.branch_id)
+    async def get_fe_states(self, limit: Optional[int] = None) -> Optional[ProjectState]:
+        return await ProjectState.get_fe_states(self.current_session, self.current_state.branch_id, limit)
 
     async def get_be_back_logs(self):
         """