More filters for get_responses (#376)

prediction_market_agent_tooling/markets/omen/omen_subgraph_handler.py

@@ -608,21 +608,25 @@ def get_resolved_bets_with_valid_answer(
         )
         return [b for b in bets if b.fpmm.is_resolved_with_valid_answer]
 
-    def get_questions(
-        self,
+    @staticmethod
+    def get_reality_question_filters(
         user: HexAddress | None = None,
         claimed: bool | None = None,
         current_answer_before: datetime | None = None,
         finalized_before: datetime | None = None,
         finalized_after: datetime | None = None,
         id_in: list[str] | None = None,
+        question_id: HexBytes | None = None,
         question_id_in: list[HexBytes] | None = None,
-    ) -> list[RealityQuestion]:
+    ) -> dict[str, t.Any]:
         where_stms: dict[str, t.Any] = {}
 
         if user is not None:
             where_stms["user"] = user.lower()
 
+        if question_id is not None:
+            where_stms["questionId"] = question_id.hex()
+
         if claimed is not None:
             if claimed:
                 where_stms["historyHash"] = ZERO_BYTES.hex()
@@ -650,6 +654,27 @@ def get_questions(
         if question_id_in is not None:
             where_stms["questionId_in"] = [x.hex() for x in question_id_in]
 
+        return where_stms
+
+    def get_questions(
+        self,
+        user: HexAddress | None = None,
+        claimed: bool | None = None,
+        current_answer_before: datetime | None = None,
+        finalized_before: datetime | None = None,
+        finalized_after: datetime | None = None,
+        id_in: list[str] | None = None,
+        question_id_in: list[HexBytes] | None = None,
+    ) -> list[RealityQuestion]:
+        where_stms: dict[str, t.Any] = self.get_reality_question_filters(
+            user=user,
+            claimed=claimed,
+            finalized_before=finalized_before,
+            finalized_after=finalized_after,
+            current_answer_before=current_answer_before,
+            id_in=id_in,
+            question_id_in=question_id_in,
+        )
         questions = self.realityeth_subgraph.Query.questions(where=where_stms)
         fields = self._get_fields_for_reality_questions(questions)
         result = self.sg.query_json(fields)
@@ -669,11 +694,24 @@ def get_answers(self, question_id: HexBytes) -> list[RealityAnswer]:
         items = self._parse_items_from_json(result)
         return [RealityAnswer.model_validate(i) for i in items]
 
-    def get_responses(self, question_id: HexBytes) -> list[RealityResponse]:
-        response = self.realityeth_subgraph.Response
-        where_stms = [
-            response.question.questionId == question_id.hex(),
-        ]
+    def get_responses(
+        self,
+        user: HexAddress | None = None,
+        question_id: HexBytes | None = None,
+        question_claimed: bool | None = None,
+        question_finalized_before: datetime | None = None,
+    ) -> list[RealityResponse]:
+        where_stms: dict[str, t.Any] = {}
+
+        if user is not None:
+            where_stms["user"] = user.lower()
+
+        where_stms["question_"] = self.get_reality_question_filters(
+            question_id=question_id,
+            claimed=question_claimed,
+            finalized_before=question_finalized_before,
+        )
+
         responses = self.realityeth_subgraph.Query.responses(where=where_stms)
         fields = self._get_fields_for_responses(responses)
         result = self.sg.query_json(fields)

pyproject.toml

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

tests/markets/omen/test_omen_subgraph_handler.py

@@ -177,7 +177,7 @@ def test_get_responses(omen_subgraph_handler: OmenSubgraphHandler) -> None:
     question_id = HexBytes.fromhex(
         HexStr("0xdcb2691a9ec05e25a6e595a9972b482ea65b789d978b27c0c06ff97345fce919")
     )
-    responses = omen_subgraph_handler.get_responses(question_id)
+    responses = omen_subgraph_handler.get_responses(question_id=question_id)
     assert len(responses) == 1
     answer = responses[0]
     assert answer.question.user == HexAddress(

tests_integration_with_local_chain/markets/omen/test_omen.py

@@ -119,7 +119,7 @@ def test_create_bet_withdraw_resolve_market(
     # ToDo - Instead of subgraph, fetch data directly from contract.
     answers = omen_subgraph_handler.get_answers(market.question.id)
     assert len(answers) == 1, answers
-    responses = omen_subgraph_handler.get_responses(market.question.id)
+    responses = omen_subgraph_handler.get_responses(question_id=market.question.id)
     assert len(responses) == 1, responses
     # ToDo: Once this test is fixed, check how to assert this, currently `answer` is HexBytes and OMEN_FALSE_OUTCOME is string, so it will never be equal.
     # assert answers[0].answer == OMEN_FALSE_OUTCOME, answers[0]