Refactor microagent MarketFunctions to take market_id as input (#66)

prediction_market_agent/agents/known_outcome_agent/benchmark.py

@@ -39,7 +39,7 @@ def to_market(self) -> AgentMarket:
             ),
             volume=None,
             created_time=None,
-            close_time=None,
+            close_time=utcnow() + timedelta(days=1),
             resolution=None,
             outcomes=["YES", "NO"],
         )

prediction_market_agent/agents/known_outcome_agent/deploy.py

@@ -1,5 +1,6 @@
 import getpass
 import random
+import typing as t
 from decimal import Decimal
 
 from loguru import logger
@@ -33,7 +34,7 @@ class DeployableKnownOutcomeAgent(DeployableAgent):
     def load(self) -> None:
         self.markets_with_known_outcomes: dict[str, Result] = {}
 
-    def pick_markets(self, markets: list[AgentMarket]) -> list[AgentMarket]:
+    def pick_markets(self, markets: t.Sequence[AgentMarket]) -> list[AgentMarket]:
         picked_markets: list[AgentMarket] = []
         for market in markets:
             if not isinstance(market, OmenAgentMarket):

prediction_market_agent/agents/microchain_agent/functions.py

@@ -1,6 +1,5 @@
 import typing as t
 from decimal import Decimal
-from typing import cast
 
 from eth_utils import to_checksum_address
 from microchain import Function
@@ -19,8 +18,8 @@
 from prediction_market_agent.agents.microchain_agent.utils import (
     MicroMarket,
     get_balance,
-    get_binary_market_from_question,
     get_binary_markets,
+    get_example_market_id,
     get_market_token_balance,
     get_no_outcome,
     get_yes_outcome,
@@ -63,11 +62,18 @@ def __init__(self, market_type: MarketType) -> None:
         self.market_type = market_type
         super().__init__()
 
+    @property
+    def currency(self) -> Currency:
+        return self.market_type.market_class.currency
+
 
 class GetMarkets(MarketFunction):
     @property
     def description(self) -> str:
-        return "Use this function to get a list of predction markets and the current yes prices"
+        return (
+            "Use this function to get a list of predction market questions, "
+            "and the corresponding market IDs"
+        )
 
     @property
     def example_args(self) -> list[str]:
@@ -80,7 +86,27 @@ def __call__(self) -> list[str]:
         ]
 
 
-class GetPropabilityForQuestion(MarketFunction):
+class GetMarketProbability(MarketFunction):
+    @property
+    def description(self) -> str:
+        return (
+            f"Use this function to get the probability of a 'Yes' outcome for "
+            f"a binary prediction market. This is equivalent to the price of "
+            f"the 'Yes' token in {self.currency}. Pass in the market id as a "
+            f"string."
+        )
+
+    @property
+    def example_args(self) -> list[str]:
+        return [get_example_market_id(self.market_type)]
+
+    def __call__(self, market_id: str) -> list[str]:
+        return [
+            str(self.market_type.market_class.get_binary_market(id=market_id).p_yes)
+        ]
+
+
+class PredictPropabilityForQuestion(MarketFunction):
     @property
     def description(self) -> str:
         return "Use this function to research the probability of an event occuring"
@@ -106,21 +132,23 @@ def __init__(self, market_type: MarketType, outcome: str):
 
     @property
     def description(self) -> str:
-        return f"Use this function to buy {self.outcome} 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 "
+            f"prediction market. The first parameter is the market id. The "
+            f"second parameter specifies how much {self.currency} you spend."
+        )
 
     @property
     def example_args(self) -> list[t.Union[str, float]]:
-        return ["Will Joe Biden get reelected in 2024?", 2.3]
+        return [get_example_market_id(self.market_type), 2.3]
 
-    def __call__(self, market: str, amount: float) -> str:
+    def __call__(self, market_id: str, amount: float) -> str:
         outcome_bool = get_boolean_outcome(self.outcome)
 
-        market_obj: AgentMarket = get_binary_market_from_question(
-            market=market, market_type=self.market_type
+        market_obj: AgentMarket = self.market_type.market_class.get_binary_market(
+            market_id
         )
-        market_obj = cast(
-            OmenAgentMarket, market_obj
-        )  # TODO fix with 0.10.0 PMAT release
+        market_obj = t.cast(OmenAgentMarket, market_obj)
         outcome_index = market_obj.get_outcome_index(self.outcome)
         market_index_set = outcome_index + 1
 
@@ -140,7 +168,7 @@ def __call__(self, market: str, amount: float) -> str:
             )
             - before_balance
         )
-        return f"Bought {tokens} {self.outcome} outcome tokens of: {market}"
+        return f"Bought {tokens} {self.outcome} outcome tokens of market with id: {market_id}"
 
 
 class BuyYes(BuyTokens):
@@ -211,8 +239,9 @@ def __call__(self, summary: str) -> str:
 class GetBalance(MarketFunction):
     @property
     def description(self) -> str:
-        currency = self.market_type.market_class.currency
-        return f"Use this function to fetch your balance, given in {currency} units."
+        return (
+            f"Use this function to fetch your balance, given in {self.currency} units."
+        )
 
     @property
     def example_args(self) -> list[str]:
@@ -248,7 +277,8 @@ def __call__(self, user_address: str) -> list[OmenUserPosition]:
 # Functions that interact with the prediction markets
 MARKET_FUNCTIONS: list[type[MarketFunction]] = [
     GetMarkets,
-    GetPropabilityForQuestion,
+    GetMarketProbability,
+    PredictPropabilityForQuestion,
     GetBalance,
     BuyYes,
     BuyNo,

prediction_market_agent/agents/microchain_agent/utils.py

@@ -1,3 +1,4 @@
+import typing as t
 from decimal import Decimal
 
 from eth_typing import ChecksumAddress
@@ -28,23 +29,23 @@
 
 class MicroMarket(BaseModel):
     question: str
-    p_yes: float
+    id: str
 
     @staticmethod
     def from_agent_market(market: AgentMarket) -> "MicroMarket":
         return MicroMarket(
             question=market.question,
-            p_yes=float(market.p_yes),
+            id=market.id,
         )
 
     def __str__(self) -> str:
-        return f"'{self.question}' with probability of yes: {self.p_yes:.2%}"
+        return f"'{self.question}', id: {self.id}"
 
 
 def get_binary_markets(market_type: MarketType) -> list[AgentMarket]:
     # Get the 5 markets that are closing soonest
     cls = market_type.market_class
-    markets: list[AgentMarket] = cls.get_binary_markets(
+    markets: t.Sequence[AgentMarket] = cls.get_binary_markets(
         filter_by=FilterBy.OPEN,
         sort_by=(
             SortBy.NONE
@@ -53,7 +54,7 @@ def get_binary_markets(market_type: MarketType) -> list[AgentMarket]:
         ),
         limit=5,
     )
-    return markets
+    return list(markets)
 
 
 def get_balance(market_type: MarketType) -> BetAmount:
@@ -68,16 +69,6 @@ def get_balance(market_type: MarketType) -> BetAmount:
         raise ValueError(f"Market type '{market_type}' not supported")
 
 
-def get_binary_market_from_question(
-    market: str, market_type: MarketType
-) -> AgentMarket:
-    markets = get_binary_markets(market_type=market_type)
-    for m in markets:
-        if m.question == market:
-            return m
-    raise ValueError(f"Market '{market}' not found")
-
-
 def get_market_token_balance(
     user_address: ChecksumAddress, market_condition_id: HexBytes, market_index_set: int
 ) -> Wei:
@@ -104,3 +95,10 @@ def get_no_outcome(market_type: MarketType) -> str:
         return OMEN_FALSE_OUTCOME
     else:
         raise ValueError(f"Market type '{market_type}' not supported")
+
+
+def get_example_market_id(market_type: MarketType) -> str:
+    if market_type == MarketType.OMEN:
+        return "0x0020d13c89140b47e10db54cbd53852b90bc1391"
+    else:
+        raise ValueError(f"Market type '{market_type}' not supported")

pyproject.toml

@@ -35,7 +35,7 @@ poetry = "^1.7.1"
 poetry-plugin-export = "^1.6.0"
 functions-framework = "^3.5.0"
 cron-validator = "^1.0.8"
-prediction-market-agent-tooling = "^0.9.3"
+prediction-market-agent-tooling = "^0.11.0"
 pydantic-settings = "^2.1.0"
 autoflake = "^2.2.1"
 isort = "^5.13.2"

tests/agents/microchain/test_agent.py

@@ -0,0 +1,87 @@
+import json
+import typing as t
+
+import pytest
+from microchain import LLM, Agent, Engine, Function, OpenAIChatGenerator
+from microchain.functions import Reasoning, Stop
+from prediction_market_agent_tooling.markets.agent_market import AgentMarket
+from prediction_market_agent_tooling.markets.markets import MarketType
+from pydantic import BaseModel
+
+from prediction_market_agent.agents.microchain_agent.functions import (
+    GetMarketProbability,
+    GetMarkets,
+)
+from prediction_market_agent.utils import APIKeys
+from tests.utils import RUN_PAID_TESTS
+
+
+@pytest.fixture
+def generator() -> OpenAIChatGenerator:
+    return OpenAIChatGenerator(
+        model="gpt-4-turbo-preview",
+        api_key=APIKeys().openai_api_key.get_secret_value(),
+        api_base="https://api.openai.com/v1",
+        temperature=0.0,
+    )
+
+
+@pytest.mark.skipif(not RUN_PAID_TESTS, reason="This test costs money to run.")
+@pytest.mark.parametrize("market_type", [MarketType.OMEN])
+def test_get_probability(
+    market_type: MarketType, generator: OpenAIChatGenerator
+) -> None:
+    """
+    Test the the agent's ability to use the GetMarkets and GetMarketProbability
+    functions.
+
+    The agent should be able to find a market and its probability, and return
+    it in valid json format.
+    """
+
+    class MarketIDAndProbability(BaseModel):
+        market_id: str
+        probability: float
+
+    class Jsonify(Function):
+        @property
+        def description(self) -> str:
+            return "Use this function to jsonify the market id and probability"
+
+        @property
+        def example_args(self) -> list[t.Any]:
+            return ["0x1234", 0.5]
+
+        def __call__(self, market_id: str, p_yes: float) -> str:
+            return MarketIDAndProbability(
+                market_id=market_id, probability=p_yes
+            ).model_dump_json()
+
+    engine = Engine()
+    engine.register(Reasoning())
+    engine.register(Stop())
+    engine.register(GetMarkets(market_type=market_type))
+    engine.register(GetMarketProbability(market_type=market_type))
+    engine.register(Jsonify())
+    agent = Agent(llm=LLM(generator=generator), engine=engine)
+    agent.prompt = f"""Act as a agent to find any market and its probability, and return it in valid json format.
+    
+    You can use the following functions:
+
+    {engine.help}
+    
+    Only output valid Python function calls. When you have a market's id and
+    probability, return it in valid json format, with fields 'market_id' and 'probability'.
+
+    Once you have output the valid json, then stop.
+    """
+
+    agent.bootstrap = ['Reasoning("I need to reason step-by-step")']
+    agent.run(iterations=5)
+
+    # history[-1] is 'user' stop message
+    # history[-2] is 'assistant' stop function call
+    m_json = json.loads(agent.history[-3]["content"])
+    m = MarketIDAndProbability.model_validate(m_json)
+    market: AgentMarket = market_type.market_class.get_binary_market(m.market_id)
+    assert market.p_yes == m.probability

tests/agents/test_microchain.py → tests/agents/microchain/test_functions.py

@@ -2,6 +2,7 @@
 from eth_typing import HexAddress, HexStr
 from microchain import Engine
 from microchain.functions import Reasoning, Stop
+from prediction_market_agent_tooling.markets.agent_market import AgentMarket
 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 (
@@ -16,6 +17,7 @@
     BuyNo,
     BuyYes,
     GetBalance,
+    GetMarketProbability,
     GetMarkets,
     GetUserPositions,
 )
@@ -79,7 +81,7 @@ def test_balance_for_user_in_market() -> None:
     market_id = HexAddress(
         HexStr("0x59975b067b0716fef6f561e1e30e44f606b08803")
     )  # yes/no
-    market = subgraph_handler.get_omen_market(market_id)
+    market = subgraph_handler.get_omen_market_by_market_id(market_id)
     omen_agent_market = OmenAgentMarket.from_data_model(market)
     balance_yes = get_market_token_balance(
         user_address=Web3.to_checksum_address(user_address),
@@ -108,3 +110,14 @@ def test_engine_help(market_type: MarketType) -> None:
         engine.register(function(market_type=market_type))
 
     print(engine.help)
+
+
+@pytest.mark.parametrize("market_type", [MarketType.OMEN])
+def test_get_probability(market_type: MarketType) -> None:
+    market_id = "0x0020d13c89140b47e10db54cbd53852b90bc1391"
+    get_market_probability = GetMarketProbability(market_type=market_type)
+    # TODO this will need fixing once https://github.com/gnosis/prediction-market-agent-tooling/issues/181 is resolved
+    assert float(get_market_probability(market_id)[0]) == 0.5
+
+    market: AgentMarket = market_type.market_class.get_binary_market(market_id)
+    assert market.is_resolved()  # Probability wont change after resolution