Add larger_than arg to AgentMarket.get_positions (#287)

prediction_market_agent_tooling/markets/agent_market.py

@@ -199,11 +199,16 @@ def get_token_balance(self, user_id: str, outcome: str) -> TokenAmount:
         raise NotImplementedError("Subclasses must implement this method")
 
     @classmethod
-    def get_positions(cls, user_id: str, liquid_only: bool = False) -> list[Position]:
+    def get_positions(
+        cls, user_id: str, liquid_only: bool = False, larger_than: float = 0
+    ) -> list[Position]:
         """
         Get all non-zero positions a user has in any market.
 
         If `liquid_only` is True, only return positions that can be sold.
+
+        If `larger_than` is not None, only return positions with a larger number
+        of tokens than this amount.
         """
         raise NotImplementedError("Subclasses must implement this method")
 

prediction_market_agent_tooling/markets/data_models.py

@@ -50,6 +50,13 @@ class Position(BaseModel):
     market_id: str
     amounts: dict[OutcomeStr, TokenAmount]
 
+    @property
+    def total_amount(self) -> TokenAmount:
+        return TokenAmount(
+            amount=sum(amount.amount for amount in self.amounts.values()),
+            currency=self.amounts[next(iter(self.amounts.keys()))].currency,
+        )
+
     def __str__(self) -> str:
         amounts_str = ", ".join(
             f"{amount.amount} '{outcome}' tokens"

prediction_market_agent_tooling/markets/omen/omen.py

@@ -367,11 +367,16 @@ def get_token_balance(
         )
 
     @classmethod
-    def get_positions(cls, user_id: str, liquid_only: bool = False) -> list[Position]:
+    def get_positions(
+        cls,
+        user_id: str,
+        liquid_only: bool = False,
+        larger_than: float = 0,
+    ) -> list[Position]:
         sgh = OmenSubgraphHandler()
         omen_positions = sgh.get_user_positions(
             better_address=Web3.to_checksum_address(user_id),
-            total_balance_bigger_than=wei_type(0),
+            total_balance_bigger_than=xdai_to_wei(xDai(larger_than)),
         )
 
         # Sort positions and corresponding markets by condition_id

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "prediction-market-agent-tooling"
-version = "0.39.1"
+version = "0.39.2"
 description = "Tools to benchmark, deploy and monitor prediction market agents."
 authors = ["Gnosis"]
 readme = "README.md"

tests/markets/omen/test_omen.py

@@ -123,6 +123,8 @@ def test_get_positions_0() -> None:
 def test_get_positions_1() -> None:
     """
     Check the user's positions against 'market.get_token_balance'
+
+    Also check that `larger_than` and `liquid_only` filters work
     """
     # Pick a user that has active positions
     user_address = Web3.to_checksum_address(
@@ -135,6 +137,23 @@ def test_get_positions_1() -> None:
     )
     assert len(positions) > len(liquid_positions)
 
+    # Get position id with smallest total amount
+    min_position_id = min(positions, key=lambda x: x.total_amount.amount).market_id
+    min_amount_position = next(
+        position for position in positions if position.market_id == min_position_id
+    )
+
+    large_positions = OmenAgentMarket.get_positions(
+        user_id=user_address, larger_than=min_amount_position.total_amount.amount
+    )
+    # Check that the smallest position has been filtered out
+    assert len(large_positions) == len(positions) - 1
+    assert all(position.market_id != min_position_id for position in large_positions)
+    assert all(
+        position.total_amount.amount > min_amount_position.total_amount.amount
+        for position in large_positions
+    )
+
     # Pick a single position to test, otherwise it can be very slow
     position = positions[0]