Fix GetWalletBalance.__call__ args (#63)

* Fix GetWalletBalance.__call__ args
* Use PMA APIKeys instead of importing from PMAT

.github/workflows/python_ci.yaml

@@ -10,7 +10,6 @@ env:
   MANIFOLD_API_KEY: ${{ secrets.MANIFOLD_API_KEY }}
   SERP_API_KEY: ${{ secrets.SERP_API_KEY }}
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-  BET_FROM_ADDRESS: ${{ secrets.BET_FROM_ADDRESS }}
   BET_FROM_PRIVATE_KEY: ${{ secrets.BET_FROM_PRIVATE_KEY }}
 
 jobs:

prediction_market_agent/agents/microchain_agent/functions.py

@@ -5,7 +5,6 @@
 
 from eth_utils import to_checksum_address
 from microchain import Function
-from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.markets.agent_market import AgentMarket
 from prediction_market_agent_tooling.markets.data_models import BetAmount, Currency
 from prediction_market_agent_tooling.markets.markets import MarketType
@@ -27,6 +26,7 @@
     get_no_outcome,
     get_yes_outcome,
 )
+from prediction_market_agent.utils import APIKeys
 
 balance = 50
 outcomeTokens = {}
@@ -236,9 +236,9 @@ def description(self) -> str:
     def example_args(self) -> list[str]:
         return []
 
-    def __call__(self, user_address: str) -> Decimal:
+    def __call__(self) -> Decimal:
         # We focus solely on xDAI balance for now to avoid the agent having to wrap/unwrap xDAI.
-        user_address_checksummed = to_checksum_address(user_address)
+        user_address_checksummed = APIKeys().bet_from_address
         balance = get_balances(user_address_checksummed)
         return balance.xdai
 

prediction_market_agent/agents/microchain_agent/microchain_agent.py

@@ -1,12 +1,9 @@
-import os
-
-from dotenv import load_dotenv
 from functions import MARKET_FUNCTIONS, MISC_FUNCTIONS
 from microchain import LLM, Agent, Engine, OpenAIChatGenerator
 from microchain.functions import Reasoning, Stop
 from prediction_market_agent_tooling.markets.markets import MarketType
 
-load_dotenv()
+from prediction_market_agent.utils import APIKeys
 
 engine = Engine()
 engine.register(Reasoning())
@@ -18,7 +15,7 @@
 
 generator = OpenAIChatGenerator(
     model="gpt-4-turbo-preview",
-    api_key=os.getenv("OPENAI_API_KEY"),
+    api_key=APIKeys().openai_api_key.get_secret_value(),
     api_base="https://api.openai.com/v1",
     temperature=0.7,
 )

tests/agents/test_microchain.py

@@ -1,5 +1,7 @@
 import pytest
 from eth_typing import HexAddress, HexStr
+from microchain import Engine
+from microchain.functions import Reasoning, Stop
 from prediction_market_agent_tooling.markets.markets import MarketType
 from prediction_market_agent_tooling.markets.omen.omen import OmenAgentMarket
 from prediction_market_agent_tooling.markets.omen.omen_subgraph_handler import (
@@ -9,6 +11,8 @@
 from web3 import Web3
 
 from prediction_market_agent.agents.microchain_agent.functions import (
+    MARKET_FUNCTIONS,
+    MISC_FUNCTIONS,
     BuyNo,
     BuyYes,
     GetMarkets,
@@ -50,7 +54,7 @@ def test_buy_no(market_type: MarketType) -> None:
 
 @pytest.mark.parametrize("market_type", [MarketType.OMEN])
 def test_replicator_has_balance_gt_0(market_type: MarketType) -> None:
-    balance = GetWalletBalance(market_type=market_type)(REPLICATOR_ADDRESS)
+    balance = GetWalletBalance(market_type=market_type)()
     assert balance > 0
 
 
@@ -91,3 +95,16 @@ def test_balance_for_user_in_market() -> None:
         market_index_set=market.condition.index_sets[1],
     )
     assert balance_no == 0
+
+
+@pytest.mark.parametrize("market_type", [MarketType.OMEN])
+def test_engine_help(market_type: MarketType) -> None:
+    engine = Engine()
+    engine.register(Reasoning())
+    engine.register(Stop())
+    for function in MISC_FUNCTIONS:
+        engine.register(function())
+    for function in MARKET_FUNCTIONS:
+        engine.register(function(market_type=market_type))
+
+    print(engine.help)