Implemented functions for microchain agent

prediction_market_agent/agents/microchain_agent/functions.py

@@ -1,13 +1,28 @@
+import os
 import pprint
 import typing as t
+from decimal import Decimal
+from typing import List
 
+from eth_typing import ChecksumAddress, HexAddress, HexStr
 from microchain import Function
-from prediction_market_agent_tooling.markets.agent_market import (
-    AgentMarket,
-    FilterBy,
-    SortBy,
-)
+from prediction_market_agent_tooling.gtypes import xDai
+from prediction_market_agent_tooling.markets.data_models import BetAmount, Currency
+from prediction_market_agent_tooling.markets.omen.data_models import OmenUserPosition
 from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
+from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import OmenSubgraphHandler
+from prediction_market_agent_tooling.tools.balances import get_balances
+from prediction_market_agent_tooling.tools.web3_utils import private_key_to_public_key
+from pydantic import SecretStr
+
+from prediction_market_agent.agents.microchain_agent.tools import (
+    get_omen_binary_market_from_question,
+    get_omen_binary_markets,
+    get_omen_market_token_balance, address_to_checksum_address,
+)
+
+#PRIVATE_KEY = os.getenv("BET_FROM_PRIVATE_KEY")
+#PUBLIC_KEY = private_key_to_public_key(SecretStr(PRIVATE_KEY))
 
 balance = 50
 outcomeTokens = {}
@@ -51,13 +66,7 @@ def example_args(self) -> list[str]:
         return []
 
     def __call__(self) -> list[str]:
-        # Get the 5 markets that are closing soonest
-        markets: list[AgentMarket] = OmenAgentMarket.get_binary_markets(
-            filter_by=FilterBy.OPEN,
-            sort_by=SortBy.CLOSING_SOONEST,
-            limit=5,
-        )
-
+        markets: list[OmenAgentMarket] = get_omen_binary_markets()
         market_questions_and_prices = []
         for market in markets:
             market_questions_and_prices.append(market.question)
@@ -98,44 +107,49 @@ def __call__(self) -> float:
         return balance
 
 
-class BuyYes(Function):
+class BuyTokens(Function):
+    def __init__(self, outcome: str):
+        self.outcome = outcome
+        super().__init__()
+
     @property
     def description(self) -> str:
-        return "Use this function to buy yes outcome tokens of a prediction market. The second parameter specifies how much $ you spend."
+        return f"Use this function to buy {self.outcome} outcome tokens of a prediction market. The second parameter specifies how much $ you spend."
 
     @property
     def example_args(self) -> list[t.Union[str, float]]:
-        return ["Will Joe Biden get reelected in 2024?", 2]
+        return ["Will Joe Biden get reelected in 2024?", 2.3]
+
+    def __call__(self, market: str, amount: float) -> str:
+        if self.outcome == "yes":
+            outcome_bool = True
+        elif self.outcome == "no":
+            outcome_bool = False
+        else:
+            raise ValueError(f"Invalid outcome: {self.outcome}")
+
+        market_obj: OmenAgentMarket = get_omen_binary_market_from_question(market)
+        before_balance = get_omen_market_token_balance(
+            market=market_obj, outcome=outcome_bool
+        )
+        market_obj.place_bet(
+            outcome_bool, BetAmount(amount=amount, currency=Currency.xDai)
+        )
+        tokens = (
+                get_omen_market_token_balance(market=market_obj, outcome=outcome_bool)
+                - before_balance
+        )
+        return f"Bought {tokens} {self.outcome} outcome tokens of: {market}"
 
-    def __call__(self, market: str, amount: int) -> str:
-        global balance
-        if amount > balance:
-            return (
-                f"Your balance of {balance} $ is not large enough to spend {amount} $."
-            )
 
-        balance -= amount
-        return "Bought " + str(amount * 2) + " yes outcome token of: " + market
+class BuyYes(BuyTokens):
+    def __init__(self):
+        super().__init__("yes")
 
 
-class BuyNo(Function):
-    @property
-    def description(self) -> str:
-        return "Use this function to buy no outcome tokens of a prdiction market. The second parameter specifies how much $ you spend."
-
-    @property
-    def example_args(self) -> list[t.Union[str, float]]:
-        return ["Will Joe Biden get reelected in 2024?", 4]
-
-    def __call__(self, market: str, amount: int) -> str:
-        global balance
-        if amount > balance:
-            return (
-                f"Your balance of {balance} $ is not large enough to spend {amount} $."
-            )
-
-        balance -= amount
-        return "Bought " + str(amount * 2) + " no outcome token of: " + market
+class BuyNo(BuyTokens):
+    def __init__(self):
+        super().__init__("no")
 
 
 class SellYes(Function):
@@ -209,6 +223,37 @@ def __call__(self, summary: str) -> str:
         return summary
 
 
+class GetWalletBalance(Function):
+    @property
+    def description(self) -> str:
+        return "Use this function to fetch the balance of a user, given in xDAI units."
+
+    @property
+    def example_args(self) -> list[str]:
+        return ["0x2DD9f5678484C1F59F97eD334725858b938B4102"]
+
+    def __call__(self, user_address: str) -> Decimal:
+        # We focus solely on xDAI balance for now to avoid the agent having to wrap/unwrap xDAI.
+        user_address_checksummed = address_to_checksum_address(user_address)
+        balance = get_balances(user_address_checksummed)
+        return balance.xdai
+
+
+class GetUserPositions(Function):
+    @property
+    def description(self) -> str:
+        return "Use this function to fetch the markets where the user has previously bet."
+
+    @property
+    def example_args(self) -> list[str]:
+        return ["0x2DD9f5678484C1F59F97eD334725858b938B4102"]
+
+    def __call__(self, user_address: str) -> list[OmenUserPosition]:
+        user_address_checksummed = address_to_checksum_address(user_address)
+        omen_subgraph_handler = OmenSubgraphHandler()
+        user_positions = omen_subgraph_handler.get_user_positions(better_address=user_address_checksummed)
+        return user_positions
+
 ALL_FUNCTIONS = [
     Sum,
     Product,
@@ -221,4 +266,6 @@ def __call__(self, summary: str) -> str:
     SellNo,
     # BalanceToOutcomes,
     SummarizeLearning,
+    GetWalletBalance,
+    GetUserPositions
 ]

prediction_market_agent/agents/microchain_agent/microchain_agent.py

@@ -1,34 +1,43 @@
 import os
 
 from dotenv import load_dotenv
+load_dotenv()
 from functions import ALL_FUNCTIONS
 from microchain import LLM, Agent, Engine, OpenAIChatGenerator
 from microchain.functions import Reasoning, Stop
 
-load_dotenv()
 
-engine = Engine()
-engine.register(Reasoning())
-engine.register(Stop())
-for function in ALL_FUNCTIONS:
-    engine.register(function())
-
-generator = OpenAIChatGenerator(
-    model="gpt-4-turbo-preview",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    api_base="https://api.openai.com/v1",
-    temperature=0.7,
-)
-agent = Agent(llm=LLM(generator=generator), engine=engine)
-agent.prompt = f"""Act as a agent. You can use the following functions:
-
-{engine.help}
-
-
-Only output valid Python function calls.
-
-"""
-
-agent.bootstrap = ['Reasoning("I need to reason step-by-step")']
-agent.run(iterations=3)
-generator.print_usage()
+
+
+def main() -> None:
+    engine = Engine()
+    engine.register(Reasoning())
+    engine.register(Stop())
+    for function in ALL_FUNCTIONS:
+        engine.register(function())
+
+    generator = OpenAIChatGenerator(
+        model="gpt-4-turbo-preview",
+        api_key=os.getenv("OPENAI_API_KEY"),
+        api_base="https://api.openai.com/v1",
+        temperature=0.7,
+    )
+    agent = Agent(llm=LLM(generator=generator), engine=engine)
+    agent.prompt = f"""Act as a agent.
+     Interact with any market available on Omen, a prediction market platform, in order to increase your balance.
+     You can use the following functions:
+
+    {engine.help}
+
+
+    Only output valid Python function calls.
+
+    """
+
+    agent.bootstrap = ['Reasoning("I need to reason step-by-step")']
+    agent.run(iterations=3)
+    generator.print_usage()
+
+
+if __name__ == "__main__":
+    main()

prediction_market_agent/agents/microchain_agent/tools.py

@@ -0,0 +1,53 @@
+from eth_typing import HexStr, HexAddress, ChecksumAddress
+from prediction_market_agent_tooling.markets.agent_market import (
+    AgentMarket,
+    FilterBy,
+    SortBy,
+)
+from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
+from prediction_market_agent_tooling.markets.omen.omen_contracts import OmenConditionalTokenContract, \
+    OmenFixedProductMarketMakerContract
+from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import OmenSubgraphHandler
+from web3 import Web3
+from web3.types import Wei
+
+
+def address_to_checksum_address(address: str) -> ChecksumAddress:
+    return ChecksumAddress(HexAddress(HexStr(address)))
+
+def get_omen_binary_markets() -> list[OmenAgentMarket]:
+    # Get the 5 markets that are closing soonest
+    markets: list[AgentMarket] = OmenAgentMarket.get_binary_markets(
+        filter_by=FilterBy.OPEN,
+        sort_by=SortBy.CLOSING_SOONEST,
+        limit=5,
+    )
+
+
+def get_omen_binary_market_from_question(market: str) -> OmenAgentMarket:
+    markets = get_omen_binary_markets()
+    for m in markets:
+        if m.question == market:
+            return m
+    raise ValueError(f"Market '{market}' not found")
+
+
+def find_index_set_for_market_outcome(market: OmenAgentMarket, market_outcome: str):
+    try:
+        market_outcome_match = market.outcomes.index(market_outcome)
+        # Index_sets start at 1
+        return market_outcome_match + 1
+    except ValueError as e:
+        print (f"Market outcome {market_outcome} not present in market. Available outcomes: {market.outcomes}")
+        raise e
+
+
+def get_omen_market_token_balance(user_address: ChecksumAddress, market: OmenAgentMarket, market_outcome: str) -> Wei:
+    # We get the multiple positions for each market
+    positions = OmenSubgraphHandler().get_positions(market.condition.id)
+    # Find position matching market_outcome
+    index_set = find_index_set_for_market_outcome(market, market_outcome)
+    position_for_index_set = next(p for p in positions if p.indexSets.__contains__(index_set))
+    position_as_int = int(position_for_index_set.id.hex(), 16)
+    balance = OmenConditionalTokenContract().balanceOf(user_address, position_as_int)
+    return balance

tests/tools/microchain_agent/test_functions.py

@@ -0,0 +1,58 @@
+import pytest
+from dotenv import load_dotenv
+from eth_typing import HexStr, HexAddress
+from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
+from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import OmenSubgraphHandler
+from prediction_market_agent_tooling.tools.hexbytes_custom import HexBytes
+from web3 import Web3
+
+from prediction_market_agent.agents.microchain_agent.functions import GetWalletBalance, GetUserPositions
+from prediction_market_agent.agents.microchain_agent.tools import get_omen_market_token_balance
+
+REPLICATOR_ADDRESS = "0x993DFcE14768e4dE4c366654bE57C21D9ba54748"
+AGENT_0_ADDRESS = "0x2DD9f5678484C1F59F97eD334725858b938B4102"
+
+
+@pytest.fixture(scope="session", autouse=True)
+def do_something(request):
+    load_dotenv()
+    yield
+
+
+@pytest.fixture()
+def get_wallet_balance():
+    return GetWalletBalance()
+
+
+def test_replicator_has_balance_lt_0():
+    balance = GetWalletBalance().__call__(REPLICATOR_ADDRESS)
+    assert balance > 0
+
+
+def test_agent_0_has_bet_on_market():
+    user_positions = GetUserPositions().__call__(AGENT_0_ADDRESS)
+    # Assert 3 conditionIds are included
+    expected_condition_ids = [
+        HexBytes("0x9c7711bee0902cc8e6838179058726a7ba769cc97d4d0ea47b31370d2d7a117b"),
+        HexBytes("0xe2bf80af2a936cdabeef4f511620a2eec46f1caf8e75eb5dc189372367a9154c"),
+        HexBytes("0x3f8153364001b26b983dd92191a084de8230f199b5ad0b045e9e1df61089b30d"),
+    ]
+    unique_condition_ids = sum([u.position.conditionIds for u in user_positions], [])
+    assert set(expected_condition_ids).issubset(unique_condition_ids)
+
+def test_balance_for_user_in_market() -> None:
+    user_address = '0x2DD9f5678484C1F59F97eD334725858b938B4102'
+    subgraph_handler = OmenSubgraphHandler()
+    market_id = HexAddress(HexStr('0x59975b067b0716fef6f561e1e30e44f606b08803')) # yes/no
+    market = subgraph_handler.get_omen_market(market_id)
+    omen_agent_market = OmenAgentMarket.from_data_model(market)
+    outcomes = omen_agent_market.outcomes
+    balance_yes = get_omen_market_token_balance(user_address=Web3.to_checksum_address(user_address),
+                                            market=omen_agent_market,
+                                            market_outcome=outcomes[0])
+    assert balance_yes == 1959903969410997
+
+    balance_no = get_omen_market_token_balance(user_address=Web3.to_checksum_address(user_address),
+                                                market=omen_agent_market,
+                                                market_outcome=outcomes[1])
+    assert balance_no == 0