Add order by and order direction field to mongodb stores (#85)

Adds the order by and order direction fields to the messages and states
objects, along with some tests.

Author: MichaelSun48

agentex/src/adapters/crud_store/adapter_mongodb.py

@@ -510,6 +510,8 @@ async def list(
         filters: dict[str, Any] | None = None,
         limit: int | None = None,
         page_number: int | None = None,
+        order_by: str | None = None,
+        order_direction: str | None = None,
     ) -> list[T]:
         """
         List all documents in the collection.
@@ -525,8 +527,20 @@ async def list(
         else:
             cursor = self.collection.find()
         cursor = cursor.skip(skip).limit(limit)
-        # Consistent sorting so that pagination works correctly
-        cursor = cursor.sort([("_id", 1)])
+
+        sort_list = []
+        if order_by:
+            direction = (
+                pymongo.DESCENDING
+                if order_direction and order_direction.lower() == "desc"
+                else pymongo.ASCENDING
+            )
+            sort_list.append((order_by, direction))
+
+        # Always use _id as tiebreaker
+        sort_list.append(("_id", pymongo.ASCENDING))
+        cursor = cursor.sort(sort_list)
+
         try:
             return [self._deserialize(doc) for doc in cursor]
         except Exception as e:

agentex/src/api/routes/messages.py

@@ -115,9 +115,15 @@ async def list_messages(
     message_use_case: DMessageUseCase,
     limit: int = 50,
     page_number: int = 1,
+    order_by: str | None = None,
+    order_direction: str = "desc",
 ) -> list[TaskMessage]:
     task_message_entities = await message_use_case.list_messages(
-        task_id=task_id, limit=limit, page_number=page_number
+        task_id=task_id,
+        limit=limit,
+        page_number=page_number,
+        order_by=order_by,
+        order_direction=order_direction,
     )
 
     return [

agentex/src/api/routes/states.py

@@ -62,9 +62,16 @@ async def filter_states(
     agent_id: str | None = Query(None, description="Agent ID"),
     limit: int = Query(50, description="Limit", ge=1),
     page_number: int = Query(1, description="Page number", ge=1),
+    order_by: str | None = Query(None, description="Field to order by"),
+    order_direction: str = Query("desc", description="Order direction (asc or desc)"),
 ) -> list[State]:
     state_entities = await states_use_case.list(
-        task_id=task_id, agent_id=agent_id, limit=limit, page_number=page_number
+        task_id=task_id,
+        agent_id=agent_id,
+        limit=limit,
+        page_number=page_number,
+        order_by=order_by,
+        order_direction=order_direction,
     )
     logger.info(f"Listing states: {state_entities}")
     return [State.model_validate(state_entity) for state_entity in state_entities]

agentex/src/domain/services/task_message_service.py

@@ -43,26 +43,35 @@ async def get_message(self, message_id: str) -> TaskMessageEntity:
         return await self.repository.get(id=message_id)
 
     async def get_messages(
-        self, task_id: str, limit: int, page_number: int
+        self,
+        task_id: str,
+        limit: int,
+        page_number: int,
+        order_by: str | None = None,
+        order_direction: str = "desc",
     ) -> list[TaskMessageEntity]:
         """
         Get all messages for a specific task.
 
         Args:
             task_id: The task ID
             limit: Optional limit on the number of messages to return
+            order_by: Optional field name to order by (defaults to created_at)
+            order_direction: Optional direction to order by ("asc" or "desc", defaults to "desc")
 
         Returns:
             List of TaskMessageEntity objects for the task
         """
-        # Sort by created_at in ascending order (oldest first)
-        # This is typically what we want for conversation history
+        # Default to created_at descending (newest first)
+        sort_field = order_by or "created_at"
+        sort_direction = 1 if order_direction.lower() == "asc" else -1
+
         return await self.repository.find_by_field(
             "task_id",
             task_id,
             limit=limit,
             page_number=page_number,
-            sort_by={"created_at": 1},
+            sort_by={sort_field: sort_direction},
         )
 
     async def append_message(

agentex/src/domain/use_cases/messages_use_case.py

@@ -105,20 +105,31 @@ async def get_message(self, message_id: str) -> TaskMessageEntity | None:
         return await self.task_message_service.get_message(message_id=message_id)
 
     async def list_messages(
-        self, task_id: str, limit: int, page_number: int
+        self,
+        task_id: str,
+        limit: int,
+        page_number: int,
+        order_by: str | None = None,
+        order_direction: str = "desc",
     ) -> list[TaskMessageEntity]:
         """
         Get all messages for a task.
 
         Args:
             task_id: The task ID
             limit: Optional limit on the number of messages to return
+            order_by: Optional field name to order by (defaults to created_at)
+            order_direction: Optional direction to order by ("asc" or "desc", defaults to "desc")
 
         Returns:
             List of TaskMessageEntity objects for the task
         """
         return await self.task_message_service.get_messages(
-            task_id=task_id, limit=limit, page_number=page_number
+            task_id=task_id,
+            limit=limit,
+            page_number=page_number,
+            order_by=order_by,
+            order_direction=order_direction,
         )
 
 

agentex/src/domain/use_cases/states_use_case.py

@@ -35,25 +35,22 @@ async def list(
         page_number: int,
         task_id: str | None = None,
         agent_id: str | None = None,
+        order_by: str | None = None,
+        order_direction: str = "desc",
     ) -> list[StateEntity]:
-        if task_id and agent_id:
-            return await self.task_state_repository.list(
-                filters={"task_id": task_id, "agent_id": agent_id},
-                limit=limit,
-                page_number=page_number,
-            )
-        elif task_id:
-            return await self.task_state_repository.list(
-                filters={"task_id": task_id}, limit=limit, page_number=page_number
-            )
-        elif agent_id:
-            return await self.task_state_repository.list(
-                filters={"agent_id": agent_id}, limit=limit, page_number=page_number
-            )
-        else:
-            return await self.task_state_repository.list(
-                limit=limit, page_number=page_number
-            )
+        filters = {}
+        if task_id:
+            filters["task_id"] = task_id
+        if agent_id:
+            filters["agent_id"] = agent_id
+
+        return await self.task_state_repository.list(
+            filters=filters if filters else None,
+            limit=limit,
+            page_number=page_number,
+            order_by=order_by,
+            order_direction=order_direction,
+        )
 
     async def update(self, id: str, task_id: str, state: dict[str, Any]) -> StateEntity:
         task_state = await self.task_state_repository.get(id=id)

agentex/tests/integration/api/messages/test_messages_api.py

@@ -289,3 +289,138 @@ async def test_list_messages_pagination(
         assert {(d["id"], d["content"]["content"]) for d in paginated_messages} == {
             (d.id, d.content.content) for d in test_pagination_messages
         }
+
+    async def test_list_messages_with_order_by(
+        self, isolated_client, isolated_repositories
+    ):
+        """Test that list messages endpoint supports order_by parameter"""
+        # Given - Create an agent and task
+        agent_repo = isolated_repositories["agent_repository"]
+        agent = AgentEntity(
+            id=orm_id(),
+            name="order-by-message-agent",
+            description="Agent for order_by message testing",
+            acp_url="http://test-acp:8000",
+            acp_type=ACPType.SYNC,
+        )
+        await agent_repo.create(agent)
+
+        task_repo = isolated_repositories["task_repository"]
+        task = TaskEntity(
+            id=orm_id(),
+            name="order-by-message-task",
+            status=TaskStatus.RUNNING,
+            status_reason="Task for order_by message testing",
+        )
+        await task_repo.create(agent_id=agent.id, task=task)
+
+        # Create multiple messages
+        message_repo = isolated_repositories["task_message_repository"]
+        messages = []
+        for i in range(3):
+            message = TaskMessageEntity(
+                id=orm_id(),
+                task_id=task.id,
+                content=TextContentEntity(
+                    type="text", author="user", content=f"Order test message {i}"
+                ),
+                streaming_status="DONE",
+            )
+            messages.append(await message_repo.create(message))
+
+        # When - Request messages with order_by=created_at and order_direction=asc
+        response_asc = await isolated_client.get(
+            "/messages",
+            params={
+                "task_id": task.id,
+                "order_by": "created_at",
+                "order_direction": "asc",
+            },
+        )
+
+        # Then - Should return messages in ascending order
+        assert response_asc.status_code == 200
+        messages_asc = response_asc.json()
+        assert len(messages_asc) == 3
+
+        # Verify ascending order
+        for i in range(len(messages_asc) - 1):
+            assert messages_asc[i]["created_at"] <= messages_asc[i + 1]["created_at"]
+
+        # When - Request messages with order_by=created_at and order_direction=desc
+        response_desc = await isolated_client.get(
+            "/messages",
+            params={
+                "task_id": task.id,
+                "order_by": "created_at",
+                "order_direction": "desc",
+            },
+        )
+
+        # Then - Should return messages in descending order
+        assert response_desc.status_code == 200
+        messages_desc = response_desc.json()
+        assert len(messages_desc) == 3
+
+        # Verify descending order
+        for i in range(len(messages_desc) - 1):
+            assert messages_desc[i]["created_at"] >= messages_desc[i + 1]["created_at"]
+
+        # Verify asc and desc return different orderings (first element of asc should be in last position of desc)
+        # Note: We only check the first element reversal since items with identical timestamps
+        # may have unpredictable relative ordering
+        assert messages_asc[0]["id"] == messages_desc[-1]["id"]
+
+    async def test_list_messages_order_by_defaults_to_desc(
+        self, isolated_client, isolated_repositories
+    ):
+        """Test that order_direction defaults to desc for messages (newest first)"""
+        # Given - Create an agent and task
+        agent_repo = isolated_repositories["agent_repository"]
+        agent = AgentEntity(
+            id=orm_id(),
+            name="order-default-message-agent",
+            description="Agent for order default message testing",
+            acp_url="http://test-acp:8000",
+            acp_type=ACPType.SYNC,
+        )
+        await agent_repo.create(agent)
+
+        task_repo = isolated_repositories["task_repository"]
+        task = TaskEntity(
+            id=orm_id(),
+            name="order-default-message-task",
+            status=TaskStatus.RUNNING,
+            status_reason="Task for order default message testing",
+        )
+        await task_repo.create(agent_id=agent.id, task=task)
+
+        # Create multiple messages
+        message_repo = isolated_repositories["task_message_repository"]
+        for i in range(3):
+            message = TaskMessageEntity(
+                id=orm_id(),
+                task_id=task.id,
+                content=TextContentEntity(
+                    type="text", author="user", content=f"Default order message {i}"
+                ),
+                streaming_status="DONE",
+            )
+            await message_repo.create(message)
+
+        # When - Request messages without specifying order_direction
+        response = await isolated_client.get(
+            "/messages",
+            params={"task_id": task.id},
+        )
+
+        # Then - Should return messages successfully
+        assert response.status_code == 200
+        messages = response.json()
+        assert len(messages) == 3
+
+        # Verify descending order - items with same timestamp may have any relative order,
+        # so we only check that timestamps are non-increasing (allowing equal timestamps)
+        timestamps = [m["created_at"] for m in messages]
+        # Sort timestamps descending and verify the returned order matches a valid descending order
+        assert timestamps == sorted(timestamps, reverse=True)