Add subgraph for omen image mapping

prediction_market_agent_tooling/markets/omen/omen_contracts.py

@@ -630,6 +630,10 @@ class OmenThumbnailMapping(ContractOnGnosisChain):
         "0xe0cf08311F03850497B0ed6A2cf067f1750C3eFc"
     )
 
+    @staticmethod
+    def construct_ipfs_url(ipfs_hash: IPFSCIDVersion0) -> str:
+        return f"https://ipfs.io/ipfs/{ipfs_hash}"
+
     def get(
         self,
         market_address: ChecksumAddress,
@@ -646,7 +650,7 @@ def get_url(
         web3: Web3 | None = None,
     ) -> str | None:
         hash_ = self.get(market_address, web3)
-        return f"https://ipfs.io/ipfs/{hash_}" if hash_ is not None else None
+        return self.construct_ipfs_url(hash_) if hash_ is not None else None
 
     def set(
         self,

prediction_market_agent_tooling/markets/omen/omen_subgraph_handler.py

@@ -2,8 +2,11 @@
 import typing as t
 from datetime import datetime
 
+import requests
 import tenacity
 from eth_typing import ChecksumAddress
+from PIL import Image
+from PIL.Image import Image as ImageType
 from subgrounds import FieldPath, Subgrounds
 
 from prediction_market_agent_tooling.config import APIKeys
@@ -20,9 +23,15 @@
     RealityAnswer,
     RealityQuestion,
 )
+from prediction_market_agent_tooling.markets.omen.omen_contracts import (
+    OmenThumbnailMapping,
+)
 from prediction_market_agent_tooling.tools.singleton import SingletonMeta
 from prediction_market_agent_tooling.tools.utils import to_int_timestamp, utcnow
-from prediction_market_agent_tooling.tools.web3_utils import ZERO_BYTES
+from prediction_market_agent_tooling.tools.web3_utils import (
+    ZERO_BYTES,
+    byte32_to_ipfscidv0,
+)
 
 
 class OmenSubgraphHandler(metaclass=SingletonMeta):
@@ -36,6 +45,11 @@ class OmenSubgraphHandler(metaclass=SingletonMeta):
 
     REALITYETH_GRAPH_URL = "https://gateway-arbitrum.network.thegraph.com/api/{graph_api_key}/subgraphs/id/E7ymrCnNcQdAAgLbdFWzGE5mvr5Mb5T9VfT43FqA7bNh"
 
+    # TODO: Switch to arbitrum subgraph once it's published.
+    OMEN_IMAGE_MAPPING_GRAPH_URL = (
+        "https://api.studio.thegraph.com/query/63564/omen-thumbnailmapping/v0.0.3"
+    )
+
     INVALID_ANSWER = "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
 
     def __init__(self) -> None:
@@ -66,6 +80,11 @@ def __init__(self) -> None:
                 graph_api_key=keys.graph_api_key.get_secret_value()
             )
         )
+        self.omen_image_mapping_subgraph = self.sg.load_subgraph(
+            self.OMEN_IMAGE_MAPPING_GRAPH_URL.format(
+                graph_api_key=keys.graph_api_key.get_secret_value()
+            )
+        )
 
     def _get_fields_for_bets(self, bets_field: FieldPath) -> list[FieldPath]:
         markets = bets_field.fpmm
@@ -609,3 +628,23 @@ def get_market_from_user_position(
                 f"Incorrect number of markets fetched {len(markets)}, expected 1."
             )
         return markets[0]
+
+    def get_market_image_url(self, market_id: HexAddress) -> str | None:
+        image = self.omen_image_mapping_subgraph.Query.omenThumbnailMapping(
+            id=market_id.lower()
+        )
+        fields = [image.id, image.image_hash]
+        result = self.sg.query_json(fields)
+        items = self._parse_items_from_json(result)
+        if not items:
+            return None
+        parsed = byte32_to_ipfscidv0(HexBytes(items[0]["image_hash"]))
+        return OmenThumbnailMapping.construct_ipfs_url(parsed)
+
+    def get_market_image(self, market_id: HexAddress) -> ImageType | None:
+        image_url = self.get_market_image_url(market_id)
+        return (
+            Image.open(requests.get(image_url, stream=True).raw)
+            if image_url is not None
+            else None
+        )

tests/markets/omen/test_omen_subgraph_handler.py

@@ -285,3 +285,22 @@ def test_omen_get_non_existing_market_by_id() -> None:
     with pytest.raises(ValueError) as e:
         OmenSubgraphHandler().get_omen_market_by_market_id(HexAddress(HexStr("0x123")))
     assert "Fetched wrong number of markets. Expected 1 but got 0" in str(e)
+
+
+def test_get_existing_image() -> None:
+    market_id = HexAddress(HexStr("0x0539590c0cf0d929e3f40b290fda04b9b4a8cf68"))
+    image_url = OmenSubgraphHandler().get_market_image_url(market_id)
+    assert (
+        image_url
+        == "https://ipfs.io/ipfs/QmRiPQ4x7jAgKMyJMV8Uqw7vQir6vmgB851TXe7CgQx7cV"
+    )
+    image = OmenSubgraphHandler().get_market_image(market_id)
+    assert image is not None
+
+
+def test_get_non_existing_image() -> None:
+    market_id = HexAddress(HexStr("0xd37dfcee94666f83d7c334658719817bf8ee1bca"))
+    image_url = OmenSubgraphHandler().get_market_image_url(market_id)
+    assert image_url is None
+    image = OmenSubgraphHandler().get_market_image(market_id)
+    assert image is None