prediction-market-agent-tooling 0.49.0__py3-none-any.whl → 0.49.2__py3-none-any.whl

prediction_market_agent_tooling/markets/omen/omen_contracts.py

@@ -5,6 +5,7 @@ from datetime import datetime, timedelta
 from enum import Enum
 
 from web3 import Web3
+from web3.constants import HASH_ZERO
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.gtypes import (
@@ -20,10 +21,15 @@ from prediction_market_agent_tooling.gtypes import (
     Wei,
     int_to_hexbytes,
     wei_type,
-    xdai_type,
 )
 from prediction_market_agent_tooling.markets.omen.data_models import (
     INVALID_ANSWER_HEX_BYTES,
+    ConditionPreparationEvent,
+    ContractPrediction,
+    FPMMFundingAddedEvent,
+    OmenFixedProductMarketMakerCreationEvent,
+    RealitioLogNewQuestionEvent,
+    format_realitio_question,
 )
 from prediction_market_agent_tooling.tools.contract import (
     ContractDepositableWrapperERC20OnGnosisChain,
@@ -38,7 +44,6 @@ from prediction_market_agent_tooling.tools.web3_utils import (
     ZERO_BYTES,
     byte32_to_ipfscidv0,
     ipfscidv0_to_byte32,
-    xdai_to_wei,
 )
 
 
@@ -84,6 +89,10 @@ class OmenOracleContract(ContractOnGnosisChain):
         )
 
 
+def build_parent_collection_id() -> HexStr:
+    return HASH_ZERO  # Taken from Olas
+
+
 class OmenConditionalTokenContract(ContractOnGnosisChain):
     # Contract ABI taken from https://gnosisscan.io/address/0xCeAfDD6bc0bEF976fdCd1112955828E00543c0Ce#code.
     abi: ABI = abi_field_validator(
@@ -177,8 +186,8 @@ class OmenConditionalTokenContract(ContractOnGnosisChain):
         api_keys: APIKeys,
         collateral_token_address: HexAddress,
         condition_id: HexBytes,
-        parent_collection_id: HexStr,
         index_sets: t.List[int],
+        parent_collection_id: HexStr = build_parent_collection_id(),
         web3: Web3 | None = None,
     ) -> TxReceipt:
         return self.send(
@@ -247,8 +256,8 @@ class OmenConditionalTokenContract(ContractOnGnosisChain):
         outcomes_slot_count: int,
         tx_params: t.Optional[TxParams] = None,
         web3: Web3 | None = None,
-    ) -> TxReceipt:
-        return self.send(
+    ) -> ConditionPreparationEvent:
+        receipt_tx = self.send(
             api_keys=api_keys,
             function_name="prepareCondition",
             function_params=[
@@ -260,6 +269,15 @@ class OmenConditionalTokenContract(ContractOnGnosisChain):
             web3=web3,
         )
 
+        event_logs = (
+            self.get_web3_contract(web3=web3)
+            .events.ConditionPreparation()
+            .process_receipt(receipt_tx)
+        )
+        cond_event = ConditionPreparationEvent(**event_logs[0]["args"])
+
+        return cond_event
+
 
 class OmenFixedProductMarketMakerContract(ContractOnGnosisChain):
     # File content taken from https://github.com/protofire/omen-exchange/blob/master/app/src/abi/marketMaker.json.
@@ -412,7 +430,8 @@ class sDaiContract(ContractERC4626OnGnosisChain):
     )
 
 
-OMEN_DEFAULT_MARKET_FEE = 0.02  # 2% fee from the buying shares amount.
+OMEN_DEFAULT_MARKET_FEE_PERC = 0.02  # 2% fee from the buying shares amount.
+REALITY_DEFAULT_FINALIZATION_TIMEOUT = timedelta(days=3)
 
 
 class OmenFixedProductMarketMakerFactoryContract(ContractOnGnosisChain):
@@ -433,14 +452,15 @@ class OmenFixedProductMarketMakerFactoryContract(ContractOnGnosisChain):
         condition_id: HexBytes,
         initial_funds_wei: Wei,
         collateral_token_address: ChecksumAddress,
-        fee: float = OMEN_DEFAULT_MARKET_FEE,
+        fee: Wei,  # This is actually fee in %, 'where 100% == 1 xDai'.
+        distribution_hint: list[OmenOutcomeToken] | None = None,
         tx_params: t.Optional[TxParams] = None,
         web3: Web3 | None = None,
-    ) -> TxReceipt:
-        fee_wei = xdai_to_wei(
-            xdai_type(fee)
-        )  # We need to convert this to the wei units, but in reality it's % fee as stated in the `OMEN_DEFAULT_MARKET_FEE` variable.
-        return self.send(
+    ) -> tuple[
+        OmenFixedProductMarketMakerCreationEvent, FPMMFundingAddedEvent, TxReceipt
+    ]:
+        web3 = web3 or self.get_web3()
+        receipt_tx = self.send(
             api_keys=api_keys,
             function_name="create2FixedProductMarketMaker",
             function_params=dict(
@@ -450,19 +470,41 @@ class OmenFixedProductMarketMakerFactoryContract(ContractOnGnosisChain):
                 conditionalTokens=OmenConditionalTokenContract().address,
                 collateralToken=collateral_token_address,
                 conditionIds=[condition_id],
-                fee=fee_wei,
+                fee=fee,
                 initialFunds=initial_funds_wei,
-                distributionHint=[],
+                distributionHint=distribution_hint or [],
             ),
             tx_params=tx_params,
             web3=web3,
         )
 
+        market_event_logs = (
+            self.get_web3_contract(web3=web3)
+            .events.FixedProductMarketMakerCreation()
+            .process_receipt(receipt_tx)
+        )
+        market_event = OmenFixedProductMarketMakerCreationEvent(
+            **market_event_logs[0]["args"]
+        )
+        funding_event_logs = (
+            self.get_web3_contract(web3=web3)
+            .events.FPMMFundingAdded()
+            .process_receipt(receipt_tx)
+        )
+        funding_event = FPMMFundingAddedEvent(**funding_event_logs[0]["args"])
+
+        return market_event, funding_event, receipt_tx
+
 
 class Arbitrator(str, Enum):
-    KLEROS = "kleros"
+    KLEROS_511_JURORS_WITHOUT_APPEAL = "kleros_511_jurors_without_appeal"
+    KLEROS_31_JURORS_WITH_APPEAL = "kleros_31_jurors_with_appeal"
     DXDAO = "dxdao"
 
+    @property
+    def is_kleros(self) -> bool:
+        return self.value.startswith("kleros")
+
 
 class OmenDxDaoContract(ContractOnGnosisChain):
     # Contract ABI taken from https://gnosisscan.io/address/0xFe14059344b74043Af518d12931600C0f52dF7c5#code.
@@ -485,9 +527,22 @@ class OmenKlerosContract(ContractOnGnosisChain):
             "../../abis/omen_kleros.abi.json",
         )
     )
-    address: ChecksumAddress = Web3.to_checksum_address(
-        "0xe40DD83a262da3f56976038F1554Fe541Fa75ecd"
-    )
+
+    @staticmethod
+    def from_arbitrator(arbitrator: "Arbitrator") -> "OmenKlerosContract":
+        """
+        See https://docs.kleros.io/developer/deployment-addresses for all available addresses.
+        """
+        if arbitrator == Arbitrator.KLEROS_511_JURORS_WITHOUT_APPEAL:
+            address = "0xe40DD83a262da3f56976038F1554Fe541Fa75ecd"
+
+        elif arbitrator == Arbitrator.KLEROS_31_JURORS_WITH_APPEAL:
+            address = "0x29f39de98d750eb77b5fafb31b2837f079fce222"
+
+        else:
+            raise ValueError(f"Unsupported arbitrator: {arbitrator=}")
+
+        return OmenKlerosContract(address=Web3.to_checksum_address(address))
 
 
 class OmenRealitioContract(ContractOnGnosisChain):
@@ -506,8 +561,8 @@ class OmenRealitioContract(ContractOnGnosisChain):
     def get_arbitrator_contract(
         arbitrator: Arbitrator,
     ) -> ContractOnGnosisChain:
-        if arbitrator == Arbitrator.KLEROS:
-            return OmenKlerosContract()
+        if arbitrator.is_kleros:
+            return OmenKlerosContract.from_arbitrator(arbitrator)
         if arbitrator == Arbitrator.DXDAO:
             return OmenDxDaoContract()
         raise ValueError(f"Unknown arbitrator: {arbitrator}")
@@ -521,25 +576,25 @@ class OmenRealitioContract(ContractOnGnosisChain):
         language: str,
         arbitrator: Arbitrator,
         opening: datetime,
-        timeout: timedelta = timedelta(days=1),
+        timeout: timedelta,
         nonce: int | None = None,
         tx_params: t.Optional[TxParams] = None,
         web3: Web3 | None = None,
-    ) -> HexBytes:
+    ) -> RealitioLogNewQuestionEvent:
         """
         After the question is created, you can find it at https://reality.eth.link/app/#!/creator/{from_address}.
         """
+        web3 = web3 or self.get_web3()
         arbitrator_contract_address = self.get_arbitrator_contract(arbitrator).address
         # See https://realitio.github.io/docs/html/contracts.html#templates
         # for possible template ids and how to format the question.
         template_id = 2
-        realitio_question = "␟".join(
-            [
-                question,
-                ",".join(f'"{o}"' for o in outcomes),
-                category,
-                language,
-            ]
+        realitio_question = format_realitio_question(
+            question=question,
+            outcomes=outcomes,
+            category=category,
+            language=language,
+            template_id=template_id,
         )
         receipt_tx = self.send(
             api_keys=api_keys,
@@ -557,10 +612,15 @@ class OmenRealitioContract(ContractOnGnosisChain):
             tx_params=tx_params,
             web3=web3,
         )
-        question_id = HexBytes(
-            receipt_tx["logs"][0]["topics"][1]
-        )  # The question id is available in the first emitted log, in the second topic.
-        return question_id
+
+        event_logs = (
+            self.get_web3_contract(web3=web3)
+            .events.LogNewQuestion()
+            .process_receipt(receipt_tx)
+        )
+        question_event = RealitioLogNewQuestionEvent(**event_logs[0]["args"])
+
+        return question_event
 
     def submitAnswer(
         self,
@@ -670,6 +730,93 @@ class OmenRealitioContract(ContractOnGnosisChain):
     ) -> TxReceipt:
         return self.send(api_keys=api_keys, function_name="withdraw", web3=web3)
 
+    def getOpeningTS(
+        self,
+        question_id: HexBytes,
+        web3: Web3 | None = None,
+    ) -> int:
+        ts: int = self.call(
+            function_name="getOpeningTS",
+            function_params=[question_id],
+            web3=web3,
+        )
+        return ts
+
+    def getFinalizeTS(
+        self,
+        question_id: HexBytes,
+        web3: Web3 | None = None,
+    ) -> int:
+        ts: int = self.call(
+            function_name="getFinalizeTS",
+            function_params=[question_id],
+            web3=web3,
+        )
+        return ts
+
+    def isFinalized(
+        self,
+        question_id: HexBytes,
+        web3: Web3 | None = None,
+    ) -> bool:
+        is_finalized: bool = self.call(
+            function_name="isFinalized",
+            function_params=[question_id],
+            web3=web3,
+        )
+        return is_finalized
+
+    def isPendingArbitration(
+        self,
+        question_id: HexBytes,
+        web3: Web3 | None = None,
+    ) -> bool:
+        is_pending_arbitration: bool = self.call(
+            function_name="isPendingArbitration",
+            function_params=[question_id],
+            web3=web3,
+        )
+        return is_pending_arbitration
+
+
+class OmenAgentResultMappingContract(ContractOnGnosisChain):
+    # Contract ABI taken from built https://github.com/gnosis/labs-contracts.
+
+    abi: ABI = abi_field_validator(
+        os.path.join(
+            os.path.dirname(os.path.realpath(__file__)),
+            "../../abis/omen_agentresultmapping.abi.json",
+        )
+    )
+
+    address: ChecksumAddress = Web3.to_checksum_address(
+        "0x260E1077dEA98e738324A6cEfB0EE9A272eD471a"
+    )
+
+    def get_predictions(
+        self,
+        market_address: ChecksumAddress,
+        web3: Web3 | None = None,
+    ) -> list[ContractPrediction]:
+        prediction_tuples = self.call(
+            "getPredictions", function_params=[market_address], web3=web3
+        )
+        return [ContractPrediction.from_tuple(p) for p in prediction_tuples]
+
+    def add_prediction(
+        self,
+        api_keys: APIKeys,
+        market_address: ChecksumAddress,
+        prediction: ContractPrediction,
+        web3: Web3 | None = None,
+    ) -> TxReceipt:
+        return self.send(
+            api_keys=api_keys,
+            function_name="addPrediction",
+            function_params=[market_address, prediction.model_dump(by_alias=True)],
+            web3=web3,
+        )
+
 
 class OmenThumbnailMapping(ContractOnGnosisChain):
     # Contract ABI taken from built https://github.com/gnosis/labs-contracts.

prediction_market_agent_tooling/markets/omen/omen_resolving.py

@@ -240,7 +240,7 @@ def omen_resolve_market_tx(
     web3: Web3 | None = None,
 ) -> None:
     """
-    Market can be resolved 24h after last answer was submitted via `omen_submit_answer_market_tx`.
+    Market can be resolved after the answer if finalized on Reality.
     """
     oracle_contract = OmenOracleContract()
     oracle_contract.resolve(

prediction_market_agent_tooling/markets/omen/omen_subgraph_handler.py

@@ -20,6 +20,7 @@ from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.agent_market import FilterBy, SortBy
 from prediction_market_agent_tooling.markets.omen.data_models import (
     OMEN_BINARY_MARKET_OUTCOMES,
+    ContractPrediction,
     OmenBet,
     OmenMarket,
     OmenPosition,
@@ -60,6 +61,8 @@ class OmenSubgraphHandler(metaclass=SingletonMeta):
 
     OMEN_IMAGE_MAPPING_GRAPH_URL = "https://gateway-arbitrum.network.thegraph.com/api/{graph_api_key}/subgraphs/id/EWN14ciGK53PpUiKSm7kMWQ6G4iz3tDrRLyZ1iXMQEdu"
 
+    OMEN_AGENT_RESULT_MAPPING_GRAPH_URL = "https://gateway-arbitrum.network.thegraph.com/api/{graph_api_key}/subgraphs/id/GoE3UFyc8Gg9xzv92oinonyhRCphpGu62qB2Eh2XvJ8F"
+
     INVALID_ANSWER = "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
 
     def __init__(self) -> None:
@@ -96,6 +99,12 @@ class OmenSubgraphHandler(metaclass=SingletonMeta):
             )
         )
 
+        self.omen_agent_result_mapping_subgraph = self.sg.load_subgraph(
+            self.OMEN_AGENT_RESULT_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
         fields_for_markets = self._get_fields_for_markets(markets)
@@ -310,6 +319,16 @@ class OmenSubgraphHandler(metaclass=SingletonMeta):
                 sort_by_field = (
                     self.trades_subgraph.FixedProductMarketMaker.openingTimestamp
                 )
+            case SortBy.HIGHEST_LIQUIDITY:
+                sort_direction = "desc"
+                sort_by_field = (
+                    self.trades_subgraph.FixedProductMarketMaker.liquidityMeasure
+                )
+            case SortBy.LOWEST_LIQUIDITY:
+                sort_direction = "asc"
+                sort_by_field = (
+                    self.trades_subgraph.FixedProductMarketMaker.liquidityMeasure
+                )
             case SortBy.NONE:
                 sort_direction = None
                 sort_by_field = None
@@ -797,3 +816,29 @@ class OmenSubgraphHandler(metaclass=SingletonMeta):
             if image_url is not None
             else None
         )
+
+    def get_agent_results_for_market(
+        self, market_id: HexAddress | None = None
+    ) -> list[ContractPrediction]:
+        where_stms = {}
+        if market_id:
+            where_stms["marketAddress"] = market_id.lower()
+
+        prediction_added = (
+            self.omen_agent_result_mapping_subgraph.Query.predictionAddeds(
+                where=where_stms,
+                orderBy="blockNumber",
+                orderDirection="asc",
+            )
+        )
+        fields = [
+            prediction_added.publisherAddress,
+            prediction_added.ipfsHash,
+            prediction_added.txHashes,
+            prediction_added.estimatedProbabilityBps,
+        ]
+        result = self.sg.query_json(fields)
+        items = self._parse_items_from_json(result)
+        if not items:
+            return []
+        return [ContractPrediction.model_validate(i) for i in items]

prediction_market_agent_tooling/markets/polymarket/api.py

@@ -2,8 +2,8 @@ import typing as t
 
 import requests
 import tenacity
-from loguru import logger
 
+from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.polymarket.data_models import (
     POLYMARKET_FALSE_OUTCOME,
     POLYMARKET_TRUE_OUTCOME,

prediction_market_agent_tooling/monitor/monitor.py

@@ -27,8 +27,8 @@ from prediction_market_agent_tooling.markets.data_models import Resolution, Reso
 from prediction_market_agent_tooling.tools.parallelism import par_map
 from prediction_market_agent_tooling.tools.utils import (
     DatetimeWithTimezone,
-    add_utc_timezone_validator,
     check_not_none,
+    convert_to_utc_datetime,
     should_not_happen,
 )
 
@@ -49,10 +49,10 @@ class DeployedAgent(BaseModel):
     raw_env_vars: dict[str, str] | None = None
 
     _add_timezone_validator_start_time = field_validator("start_time")(
-        add_utc_timezone_validator
+        convert_to_utc_datetime
     )
     _add_timezone_validator_end_time = field_validator("end_time")(
-        add_utc_timezone_validator
+        convert_to_utc_datetime
     )
 
     def model_dump_prefixed(self) -> dict[str, t.Any]:

prediction_market_agent_tooling/monitor/monitor_app.py

@@ -28,8 +28,8 @@ from prediction_market_agent_tooling.monitor.monitor import (
 from prediction_market_agent_tooling.monitor.monitor_settings import MonitorSettings
 from prediction_market_agent_tooling.tools.utils import (
     DatetimeWithTimezone,
-    add_utc_timezone_validator,
     check_not_none,
+    convert_to_utc_datetime,
     utcnow,
 )
 
@@ -104,7 +104,7 @@ def monitor_app(
         st.selectbox(label="Market type", options=enabled_market_types, index=0)
     )
     start_time: DatetimeWithTimezone | None = (
-        add_utc_timezone_validator(
+        convert_to_utc_datetime(
             datetime.combine(
                 t.cast(
                     # This will be always a date for us, so casting.

prediction_market_agent_tooling/tools/contract.py

@@ -6,6 +6,7 @@ from contextlib import contextmanager
 
 from pydantic import BaseModel, field_validator
 from web3 import Web3
+from web3.contract.contract import Contract as Web3Contract
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.gtypes import (
@@ -21,7 +22,11 @@ from prediction_market_agent_tooling.tools.gnosis_rpc import (
     GNOSIS_NETWORK_ID,
     GNOSIS_RPC_URL,
 )
-from prediction_market_agent_tooling.tools.utils import should_not_happen
+from prediction_market_agent_tooling.tools.utils import (
+    DatetimeWithTimezone,
+    should_not_happen,
+    utc_timestamp_to_utc_datetime,
+)
 from prediction_market_agent_tooling.tools.web3_utils import (
     call_function_on_contract,
     send_function_on_contract_tx,
@@ -73,6 +78,10 @@ class ContractBaseClass(BaseModel):
         {}
     )  # Can be used to hold values that aren't going to change after getting them for the first time, as for example `symbol` of an ERC-20 token.
 
+    def get_web3_contract(self, web3: Web3 | None = None) -> Web3Contract:
+        web3 = web3 or self.get_web3()
+        return web3.eth.contract(address=self.address, abi=self.abi)
+
     def call(
         self,
         function_name: str,
@@ -423,6 +432,46 @@ class ContractERC4626OnGnosisChain(
         return to_gnosis_chain_contract(super().get_asset_token_contract(web3=web3))
 
 
+class DebuggingContract(ContractOnGnosisChain):
+    # Contract ABI taken from https://gnosisscan.io/address/0x5Aa82E068aE6a6a1C26c42E5a59520a74Cdb8998#code.
+    abi: ABI = abi_field_validator(
+        os.path.join(
+            os.path.dirname(os.path.realpath(__file__)),
+            "../abis/debuggingcontract.abi.json",
+        )
+    )
+    address: ChecksumAddress = Web3.to_checksum_address(
+        "0x5Aa82E068aE6a6a1C26c42E5a59520a74Cdb8998"
+    )
+
+    def getNow(
+        self,
+        web3: Web3 | None = None,
+    ) -> int:
+        now: int = self.call(
+            function_name="getNow",
+            web3=web3,
+        )
+        return now
+
+    def get_now(
+        self,
+        web3: Web3 | None = None,
+    ) -> DatetimeWithTimezone:
+        return utc_timestamp_to_utc_datetime(self.getNow(web3))
+
+    def inc(
+        self,
+        api_keys: APIKeys,
+        web3: Web3 | None = None,
+    ) -> TxReceipt:
+        return self.send(
+            api_keys=api_keys,
+            function_name="inc",
+            web3=web3,
+        )
+
+
 def contract_implements_function(
     contract_address: ChecksumAddress,
     function_name: str,

prediction_market_agent_tooling/tools/httpx_cached_client.py

@@ -0,0 +1,11 @@
+import hishel
+
+
+class HttpxCachedClient:
+    def __init__(self) -> None:
+        storage = hishel.FileStorage(ttl=3600, check_ttl_every=600)
+        controller = hishel.Controller(force_cache=True)
+        self.client = hishel.CacheClient(storage=storage, controller=controller)
+
+    def get_client(self) -> hishel.CacheClient:
+        return self.client

prediction_market_agent_tooling/tools/ipfs/ipfs_handler.py

@@ -0,0 +1,33 @@
+import json
+import tempfile
+
+from pinatapy import PinataPy
+
+from prediction_market_agent_tooling.config import APIKeys
+from prediction_market_agent_tooling.gtypes import IPFSCIDVersion0
+from prediction_market_agent_tooling.markets.omen.data_models import IPFSAgentResult
+
+
+class IPFSHandler:
+    def __init__(self, api_keys: APIKeys):
+        self.pinata = PinataPy(
+            api_keys.pinata_api_key.get_secret_value(),
+            api_keys.pinata_api_secret.get_secret_value(),
+        )
+
+    def upload_file(self, file_path: str) -> IPFSCIDVersion0:
+        return IPFSCIDVersion0(
+            self.pinata.pin_file_to_ipfs(file_path, save_absolute_paths=False)[
+                "IpfsHash"
+            ]
+        )
+
+    def store_agent_result(self, agent_result: IPFSAgentResult) -> IPFSCIDVersion0:
+        with tempfile.NamedTemporaryFile(mode="r+", encoding="utf-8") as json_file:
+            json.dump(agent_result.model_dump(), json_file, indent=4)
+            json_file.flush()
+            ipfs_hash = self.upload_file(json_file.name)
+            return ipfs_hash
+
+    def unpin_file(self, hash_to_remove: str) -> None:
+        self.pinata.remove_pin_from_ipfs(hash_to_remove=hash_to_remove)

prediction_market_agent_tooling/tools/is_predictable.py

@@ -1,7 +1,7 @@
 import tenacity
-from loguru import logger
 
 from prediction_market_agent_tooling.config import APIKeys
+from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.cache import persistent_inmemory_cache
 from prediction_market_agent_tooling.tools.langfuse_ import (
     get_langfuse_langchain_config,

prediction_market_agent_tooling/tools/langfuse_client_utils.py

@@ -6,6 +6,7 @@ from langfuse import Langfuse
 from langfuse.client import TraceWithDetails
 from pydantic import BaseModel
 
+from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.data_models import (
     PlacedTrade,
     ProbabilisticAnswer,
@@ -13,7 +14,7 @@ from prediction_market_agent_tooling.markets.data_models import (
     TradeType,
 )
 from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
-from prediction_market_agent_tooling.tools.utils import add_utc_timezone_validator
+from prediction_market_agent_tooling.tools.utils import convert_to_utc_datetime
 
 
 class ProcessMarketTrace(BaseModel):
@@ -23,11 +24,11 @@ class ProcessMarketTrace(BaseModel):
     trades: list[PlacedTrade]
 
     @property
-    def buy_trade(self) -> PlacedTrade:
+    def buy_trade(self) -> PlacedTrade | None:
         buy_trades = [t for t in self.trades if t.trade_type == TradeType.BUY]
-        if len(buy_trades) == 1:
-            return buy_trades[0]
-        raise ValueError("No buy trade found")
+        if len(buy_trades) > 1:
+            raise ValueError("Unhandled logic, check it outm please!")
+        return buy_trades[0] if buy_trades else None
 
     @staticmethod
     def from_langfuse_trace(
@@ -134,8 +135,10 @@ def get_trace_for_bet(
     traces_for_bet: list[ProcessMarketTrace] = []
     for t in traces:
         # Cannot use exact comparison due to gas fees
-        if t.buy_trade.outcome == bet.outcome and np.isclose(
-            t.buy_trade.amount.amount, bet.amount.amount
+        if (
+            t.buy_trade
+            and t.buy_trade.outcome == bet.outcome
+            and np.isclose(t.buy_trade.amount.amount, bet.amount.amount)
         ):
             traces_for_bet.append(t)
 
@@ -146,9 +149,21 @@ def get_trace_for_bet(
     else:
         # In-case there are multiple traces for the same market, get the closest
         # trace to the bet
+        bet_timestamp = convert_to_utc_datetime(bet.created_time)
         closest_trace_index = get_closest_datetime_from_list(
-            add_utc_timezone_validator(bet.created_time),
+            bet_timestamp,
             [t.timestamp for t in traces_for_bet],
         )
 
+        # Sanity check: Let's say the upper bound for time between
+        # `agent.process_market` being called and the bet being placed is 20
+        # minutes
+        candidate_trace = traces_for_bet[closest_trace_index]
+        if abs(candidate_trace.timestamp - bet_timestamp).total_seconds() > 1200:
+            logger.info(
+                f"Closest trace to bet has timestamp {candidate_trace.timestamp}, "
+                f"but bet was created at {bet_timestamp}. Not matching"
+            )
+            return None
+
         return traces_for_bet[closest_trace_index]