Implemented PR comments

prediction_market_agent/agents/microchain_agent/functions.py

@@ -19,6 +19,7 @@
 
 from prediction_market_agent.agents.microchain_agent.utils import (
     MicroMarket,
+    fetch_public_key_from_env,
     get_omen_binary_market_from_question,
     get_omen_binary_markets,
     get_omen_market_token_balance,
@@ -106,6 +107,7 @@ def __call__(self) -> float:
 
 class BuyTokens(Function):
     def __init__(self, outcome: str):
+        self.user_address = fetch_public_key_from_env()
         self.outcome = outcome
         super().__init__()
 
@@ -121,14 +123,23 @@ def __call__(self, market: str, amount: float) -> str:
         outcome_bool = get_boolean_outcome(self.outcome)
 
         market_obj: OmenAgentMarket = get_omen_binary_market_from_question(market)
+        outcome_index = market_obj.get_outcome_index(self.outcome)
+        market_index_set = outcome_index + 1
+
         before_balance = get_omen_market_token_balance(
-            market=market_obj, outcome=outcome_bool
+            user_address=self.user_address,
+            market_condition_id=market_obj.condition.id,
+            market_index_set=market_index_set,
         )
         market_obj.place_bet(
             outcome_bool, BetAmount(amount=Decimal(amount), currency=Currency.xDai)
         )
         tokens = (
-            get_omen_market_token_balance(market=market_obj, outcome=outcome_bool)
+            get_omen_market_token_balance(
+                user_address=self.user_address,
+                market_condition_id=market_obj.condition.id,
+                market_index_set=market_index_set,
+            )
             - before_balance
         )
         return f"Bought {tokens} {self.outcome} outcome tokens of: {market}"
@@ -243,12 +254,9 @@ def example_args(self) -> list[str]:
         return ["0x2DD9f5678484C1F59F97eD334725858b938B4102"]
 
     def __call__(self, user_address: str) -> list[OmenUserPosition]:
-        user_address_checksummed = to_checksum_address(user_address)
-        omen_subgraph_handler = OmenSubgraphHandler()
-        user_positions = omen_subgraph_handler.get_user_positions(
-            better_address=user_address_checksummed
+        return OmenSubgraphHandler().get_user_positions(
+            better_address=to_checksum_address(user_address)
         )
-        return user_positions
 
 
 ALL_FUNCTIONS = [

prediction_market_agent/agents/microchain_agent/utils.py

@@ -1,12 +1,23 @@
+import os
 from typing import List, cast
 
+from eth_typing import ChecksumAddress
 from prediction_market_agent_tooling.markets.agent_market import (
     AgentMarket,
     FilterBy,
     SortBy,
 )
 from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
-from pydantic import BaseModel
+from prediction_market_agent_tooling.markets.omen.omen_contracts import (
+    OmenConditionalTokenContract,
+)
+from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import (
+    OmenSubgraphHandler,
+)
+from prediction_market_agent_tooling.tools.hexbytes_custom import HexBytes
+from prediction_market_agent_tooling.tools.web3_utils import private_key_to_public_key
+from pydantic import BaseModel, SecretStr
+from web3.types import Wei
 
 
 class MicroMarket(BaseModel):
@@ -42,6 +53,22 @@ def get_omen_binary_market_from_question(market: str) -> OmenAgentMarket:
     raise ValueError(f"Market '{market}' not found")
 
 
-def get_omen_market_token_balance(market: OmenAgentMarket, outcome: bool) -> float:
-    # TODO implement this
-    return 7.3
+def get_omen_market_token_balance(
+    user_address: ChecksumAddress, market_condition_id: HexBytes, market_index_set: int
+) -> Wei:
+    # We get the multiple positions for each market
+    positions = OmenSubgraphHandler().get_positions(market_condition_id)
+    # Find position matching market_outcome
+    position_for_index_set = next(
+        p for p in positions if market_index_set in p.indexSets
+    )
+    position_as_int = int(position_for_index_set.id.hex(), 16)
+    balance = OmenConditionalTokenContract().balanceOf(user_address, position_as_int)
+    return balance
+
+
+def fetch_public_key_from_env() -> ChecksumAddress:
+    private_key = os.environ.get("BET_FROM_PRIVATE_KEY")
+    if private_key is None:
+        raise EnvironmentError("Could not load private key using env var")
+    return private_key_to_public_key(SecretStr(private_key))

tests/agents/test_microchain.py

@@ -17,11 +17,9 @@
     GetUserPositions,
     GetWalletBalance,
 )
-from prediction_market_agent.agents.microchain_agent.tools import (
-    get_omen_market_token_balance,
-)
 from prediction_market_agent.agents.microchain_agent.utils import (
     get_omen_binary_markets,
+    get_omen_market_token_balance,
 )
 from tests.utils import RUN_PAID_TESTS
 
@@ -54,13 +52,13 @@ def test_buy_no() -> None:
     print(buy_yes(market.question, 0.0001))
 
 
-def test_replicator_has_balance_lt_0() -> None:
-    balance = GetWalletBalance().__call__(REPLICATOR_ADDRESS)
+def test_replicator_has_balance_gt_0() -> None:
+    balance = GetWalletBalance()(REPLICATOR_ADDRESS)
     assert balance > 0
 
 
 def test_agent_0_has_bet_on_market() -> None:
-    user_positions = GetUserPositions().__call__(AGENT_0_ADDRESS)
+    user_positions = GetUserPositions()(AGENT_0_ADDRESS)
     # Assert 3 conditionIds are included
     expected_condition_ids = [
         HexBytes("0x9c7711bee0902cc8e6838179058726a7ba769cc97d4d0ea47b31370d2d7a117b"),