prediction-market-agent-tooling 0.52.2__py3-none-any.whl → 0.54.0__py3-none-any.whl

prediction_market_agent_tooling/deploy/agent.py

@@ -8,10 +8,8 @@ from datetime import timedelta
 from enum import Enum
 from functools import cached_property
 
-from pydantic import BaseModel, BeforeValidator, computed_field
+from pydantic import BeforeValidator, computed_field
 from typing_extensions import Annotated
-from web3 import Web3
-from web3.constants import HASH_ZERO
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.deploy.betting_strategy import (
@@ -32,11 +30,12 @@ from prediction_market_agent_tooling.deploy.gcp.utils import (
     gcp_function_is_active,
     gcp_resolve_api_keys_secrets,
 )
-from prediction_market_agent_tooling.gtypes import HexStr, xDai, xdai_type
+from prediction_market_agent_tooling.gtypes import xDai, xdai_type
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.markets.agent_market import (
     AgentMarket,
     FilterBy,
+    ProcessedMarket,
     SortBy,
 )
 from prediction_market_agent_tooling.markets.data_models import (
@@ -49,28 +48,16 @@ from prediction_market_agent_tooling.markets.markets import (
     MarketType,
     have_bet_on_market_since,
 )
-from prediction_market_agent_tooling.markets.omen.data_models import (
-    ContractPrediction,
-    IPFSAgentResult,
-)
 from prediction_market_agent_tooling.markets.omen.omen import (
-    is_minimum_required_balance,
-    redeem_from_all_user_positions,
     withdraw_wxdai_to_xdai_to_keep_balance,
 )
-from prediction_market_agent_tooling.markets.omen.omen_contracts import (
-    OmenAgentResultMappingContract,
-)
 from prediction_market_agent_tooling.monitor.monitor_app import (
     MARKET_TYPE_TO_DEPLOYED_AGENT,
 )
-from prediction_market_agent_tooling.tools.hexbytes_custom import HexBytes
-from prediction_market_agent_tooling.tools.ipfs.ipfs_handler import IPFSHandler
 from prediction_market_agent_tooling.tools.is_invalid import is_invalid
 from prediction_market_agent_tooling.tools.is_predictable import is_predictable_binary
 from prediction_market_agent_tooling.tools.langfuse_ import langfuse_context, observe
 from prediction_market_agent_tooling.tools.utils import DatetimeUTC, utcnow
-from prediction_market_agent_tooling.tools.web3_utils import ipfscidv0_to_byte32
 
 MAX_AVAILABLE_MARKETS = 20
 TRADER_TAG = "trader"
@@ -122,11 +109,6 @@ class OutOfFundsError(ValueError):
     pass
 
 
-class ProcessedMarket(BaseModel):
-    answer: ProbabilisticAnswer
-    trades: list[PlacedTrade]
-
-
 class AnsweredEnum(str, Enum):
     ANSWERED = "answered"
     NOT_ANSWERED = "not_answered"
@@ -294,7 +276,6 @@ def {entrypoint_function_name}(request) -> str:
 
 class DeployableTraderAgent(DeployableAgent):
     bet_on_n_markets_per_run: int = 1
-    min_required_balance_to_operate: xDai | None = xdai_type(1)
     min_balance_to_keep_in_native_currency: xDai | None = xdai_type(0.1)
     allow_invalid_questions: bool = False
     same_market_bet_interval: timedelta = timedelta(hours=24)
@@ -353,38 +334,25 @@ class DeployableTraderAgent(DeployableAgent):
             ]
         )
 
-    def check_min_required_balance_to_operate(
-        self,
-        market_type: MarketType,
-        check_for_gas: bool = True,
-        check_for_trades: bool = True,
-    ) -> None:
+    def check_min_required_balance_to_operate(self, market_type: MarketType) -> None:
         api_keys = APIKeys()
-        if (
-            market_type == MarketType.OMEN
-            and check_for_gas
-            and not is_minimum_required_balance(
-                api_keys.public_key,
-                min_required_balance=xdai_type(0.001),
-                sum_wxdai=False,
-            )
-        ):
+
+        if not market_type.market_class.verify_operational_balance(api_keys):
             raise CantPayForGasError(
-                f"{api_keys.public_key=} doesn't have enough xDai to pay for gas."
+                f"{api_keys=} doesn't have enough operational balance."
             )
-        if self.min_required_balance_to_operate is None:
-            return
-        if (
-            market_type == MarketType.OMEN
-            and check_for_trades
-            and not is_minimum_required_balance(
-                api_keys.bet_from_address,
-                min_required_balance=self.min_required_balance_to_operate,
-            )
-        ):
+
+    def check_min_required_balance_to_trade(self, market: AgentMarket) -> None:
+        api_keys = APIKeys()
+
+        # Get the strategy to know how much it will bet.
+        strategy = self.get_betting_strategy(market)
+        # Have a little bandwidth after the bet.
+        min_required_balance_to_trade = strategy.maximum_possible_bet_amount * 1.01
+
+        if market.get_trade_balance(api_keys) < min_required_balance_to_trade:
             raise OutOfFundsError(
-                f"Minimum required balance {self.min_required_balance_to_operate} "
-                f"for agent with address {api_keys.bet_from_address=} is not met."
+                f"Minimum required balance {min_required_balance_to_trade} for agent is not met."
             )
 
     def have_bet_on_market_since(self, market: AgentMarket, since: timedelta) -> bool:
@@ -447,6 +415,19 @@ class DeployableTraderAgent(DeployableAgent):
     ) -> None:
         self.update_langfuse_trace_by_market(market_type, market)
 
+        api_keys = APIKeys()
+
+        self.check_min_required_balance_to_trade(market)
+
+        if market_type.is_blockchain_market:
+            # Exchange wxdai back to xdai if the balance is getting low, so we can keep paying for fees.
+            if self.min_balance_to_keep_in_native_currency is not None:
+                withdraw_wxdai_to_xdai_to_keep_balance(
+                    api_keys,
+                    min_required_balance=self.min_balance_to_keep_in_native_currency,
+                    withdraw_multiplier=2,
+                )
+
     def process_market(
         self,
         market_type: MarketType,
@@ -510,72 +491,16 @@ class DeployableTraderAgent(DeployableAgent):
         market: AgentMarket,
         processed_market: ProcessedMarket,
     ) -> None:
-        if market_type != MarketType.OMEN:
-            logger.info(
-                f"Skipping after_process_market since market_type {market_type} != OMEN"
-            )
-            return
         keys = APIKeys()
-        self.store_prediction(
-            market_id=market.id, processed_market=processed_market, keys=keys
-        )
-
-    def store_prediction(
-        self, market_id: str, processed_market: ProcessedMarket, keys: APIKeys
-    ) -> None:
-        reasoning = (
-            processed_market.answer.reasoning
-            if processed_market.answer.reasoning
-            else ""
-        )
-
-        ipfs_hash_decoded = HexBytes(HASH_ZERO)
-        if keys.enable_ipfs_upload:
-            logger.info("Storing prediction on IPFS.")
-            ipfs_hash = IPFSHandler(keys).store_agent_result(
-                IPFSAgentResult(reasoning=reasoning)
-            )
-            ipfs_hash_decoded = ipfscidv0_to_byte32(ipfs_hash)
-
-        tx_hashes = [
-            HexBytes(HexStr(i.id)) for i in processed_market.trades if i.id is not None
-        ]
-        prediction = ContractPrediction(
-            publisher=keys.public_key,
-            ipfs_hash=ipfs_hash_decoded,
-            tx_hashes=tx_hashes,
-            estimated_probability_bps=int(processed_market.answer.p_yes * 10000),
-        )
-        tx_receipt = OmenAgentResultMappingContract().add_prediction(
-            api_keys=keys,
-            market_address=Web3.to_checksum_address(market_id),
-            prediction=prediction,
-        )
-        logger.info(
-            f"Added prediction to market {market_id}. - receipt {tx_receipt['transactionHash'].hex()}."
-        )
+        market.store_prediction(processed_market=processed_market, keys=keys)
 
     def before_process_markets(self, market_type: MarketType) -> None:
         """
         Executes actions that occur before bets are placed.
         """
         api_keys = APIKeys()
-        if market_type == MarketType.OMEN:
-            # First, check if we have enough xDai to pay for gas, there is no way of doing anything without it.
-            self.check_min_required_balance_to_operate(
-                market_type, check_for_trades=False
-            )
-            # Omen is specific, because the user (agent) needs to manually withdraw winnings from the market.
-            redeem_from_all_user_positions(api_keys)
-            # After redeeming, check if we have enough xDai to pay for gas and place bets.
-            self.check_min_required_balance_to_operate(market_type)
-            # Exchange wxdai back to xdai if the balance is getting low, so we can keep paying for fees.
-            if self.min_balance_to_keep_in_native_currency is not None:
-                withdraw_wxdai_to_xdai_to_keep_balance(
-                    api_keys,
-                    min_required_balance=self.min_balance_to_keep_in_native_currency,
-                    withdraw_multiplier=2,
-                )
+        self.check_min_required_balance_to_operate(market_type)
+        market_type.market_class.redeem_winnings(api_keys)
 
     def process_markets(self, market_type: MarketType) -> None:
         """
@@ -589,9 +514,6 @@ class DeployableTraderAgent(DeployableAgent):
         processed = 0
 
         for market in available_markets:
-            # We need to check it again before each market bet, as the balance might have changed.
-            self.check_min_required_balance_to_operate(market_type)
-
             processed_market = self.process_market(market_type, market)
 
             if processed_market is not None:
@@ -603,7 +525,7 @@ class DeployableTraderAgent(DeployableAgent):
         logger.info("All markets processed.")
 
     def after_process_markets(self, market_type: MarketType) -> None:
-        pass
+        "Executes actions that occur after bets are placed."
 
     def run(self, market_type: MarketType) -> None:
         self.before_process_markets(market_type)

prediction_market_agent_tooling/deploy/betting_strategy.py

@@ -32,7 +32,12 @@ class BettingStrategy(ABC):
         answer: ProbabilisticAnswer,
         market: AgentMarket,
     ) -> list[Trade]:
-        pass
+        raise NotImplementedError("Subclass should implement this.")
+
+    @property
+    @abstractmethod
+    def maximum_possible_bet_amount(self) -> float:
+        raise NotImplementedError("Subclass should implement this.")
 
     def build_zero_token_amount(self, currency: Currency) -> TokenAmount:
         return TokenAmount(amount=0, currency=currency)
@@ -126,6 +131,10 @@ class MaxAccuracyBettingStrategy(BettingStrategy):
     def __init__(self, bet_amount: float):
         self.bet_amount = bet_amount
 
+    @property
+    def maximum_possible_bet_amount(self) -> float:
+        return self.bet_amount
+
     def calculate_trades(
         self,
         existing_position: Position | None,
@@ -168,6 +177,10 @@ class KellyBettingStrategy(BettingStrategy):
         self.max_bet_amount = max_bet_amount
         self.max_price_impact = max_price_impact
 
+    @property
+    def maximum_possible_bet_amount(self) -> float:
+        return self.max_bet_amount
+
     def calculate_trades(
         self,
         existing_position: Position | None,
@@ -282,6 +295,10 @@ class MaxAccuracyWithKellyScaledBetsStrategy(BettingStrategy):
     def __init__(self, max_bet_amount: float = 10):
         self.max_bet_amount = max_bet_amount
 
+    @property
+    def maximum_possible_bet_amount(self) -> float:
+        return self.max_bet_amount
+
     def adjust_bet_amount(
         self, existing_position: Position | None, market: AgentMarket
     ) -> float:

prediction_market_agent_tooling/markets/agent_market.py

@@ -11,7 +11,9 @@ from prediction_market_agent_tooling.markets.data_models import (
     Bet,
     BetAmount,
     Currency,
+    PlacedTrade,
     Position,
+    ProbabilisticAnswer,
     Resolution,
     ResolvedBet,
     TokenAmount,
@@ -25,6 +27,11 @@ from prediction_market_agent_tooling.tools.utils import (
 )
 
 
+class ProcessedMarket(BaseModel):
+    answer: ProbabilisticAnswer
+    trades: list[PlacedTrade]
+
+
 class SortBy(str, Enum):
     CLOSING_SOONEST = "closing-soonest"
     NEWEST = "newest"
@@ -198,6 +205,36 @@ class AgentMarket(BaseModel):
     def get_binary_market(id: str) -> "AgentMarket":
         raise NotImplementedError("Subclasses must implement this method")
 
+    @staticmethod
+    def redeem_winnings(api_keys: APIKeys) -> None:
+        """
+        On some markets (like Omen), it's needed to manually claim the winner bets. If it's not needed, just implement with `pass`.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @staticmethod
+    def get_trade_balance(api_keys: APIKeys) -> float:
+        """
+        Return balance that can be used to trade on the given market.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @staticmethod
+    def verify_operational_balance(api_keys: APIKeys) -> bool:
+        """
+        Return `True` if the user has enough of operational balance. If not needed, just return `True`.
+        For example: Omen needs at least some xDai in the wallet to execute transactions.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def store_prediction(
+        self, processed_market: ProcessedMarket, keys: APIKeys
+    ) -> None:
+        """
+        If market allows to upload predictions somewhere, implement it in this method.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
     @staticmethod
     def get_bets_made_since(
         better_address: ChecksumAddress, start_time: DatetimeUTC

prediction_market_agent_tooling/markets/manifold/manifold.py

@@ -125,6 +125,11 @@ class ManifoldAgentMarket(AgentMarket):
             )
         ]
 
+    @staticmethod
+    def redeem_winnings(api_keys: APIKeys) -> None:
+        # It's done automatically on Manifold.
+        pass
+
     @classmethod
     def get_user_url(cls, keys: APIKeys) -> str:
         return get_authenticated_user(keys.manifold_api_key.get_secret_value()).url

prediction_market_agent_tooling/markets/markets.py

@@ -46,6 +46,10 @@ class MarketType(str, Enum):
             raise ValueError(f"Unknown market type: {self}")
         return MARKET_TYPE_TO_AGENT_MARKET[self]
 
+    @property
+    def is_blockchain_market(self) -> bool:
+        return self in [MarketType.OMEN, MarketType.POLYMARKET]
+
 
 MARKET_TYPE_TO_AGENT_MARKET: dict[MarketType, type[AgentMarket]] = {
     MarketType.MANIFOLD: ManifoldAgentMarket,

prediction_market_agent_tooling/markets/omen/omen.py

@@ -4,6 +4,7 @@ from datetime import timedelta
 
 import tenacity
 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 (
@@ -23,6 +24,7 @@ from prediction_market_agent_tooling.markets.agent_market import (
     AgentMarket,
     FilterBy,
     MarketFees,
+    ProcessedMarket,
     SortBy,
 )
 from prediction_market_agent_tooling.markets.data_models import (
@@ -39,7 +41,9 @@ from prediction_market_agent_tooling.markets.omen.data_models import (
     PRESAGIO_BASE_URL,
     Condition,
     ConditionPreparationEvent,
+    ContractPrediction,
     CreatedMarket,
+    IPFSAgentResult,
     OmenBet,
     OmenMarket,
     OmenUserPosition,
@@ -50,6 +54,7 @@ from prediction_market_agent_tooling.markets.omen.omen_contracts import (
     OMEN_DEFAULT_MARKET_FEE_PERC,
     REALITY_DEFAULT_FINALIZATION_TIMEOUT,
     Arbitrator,
+    OmenAgentResultMappingContract,
     OmenConditionalTokenContract,
     OmenFixedProductMarketMakerContract,
     OmenFixedProductMarketMakerFactoryContract,
@@ -70,6 +75,7 @@ from prediction_market_agent_tooling.tools.contract import (
     to_gnosis_chain_contract,
 )
 from prediction_market_agent_tooling.tools.hexbytes_custom import HexBytes
+from prediction_market_agent_tooling.tools.ipfs.ipfs_handler import IPFSHandler
 from prediction_market_agent_tooling.tools.utils import (
     DatetimeUTC,
     calculate_sell_amount_in_collateral,
@@ -78,6 +84,7 @@ from prediction_market_agent_tooling.tools.utils import (
 from prediction_market_agent_tooling.tools.web3_utils import (
     add_fraction,
     get_receipt_block_timestamp,
+    ipfscidv0_to_byte32,
     remove_fraction,
     wei_to_xdai,
     xdai_to_wei,
@@ -392,6 +399,58 @@ class OmenAgentMarket(AgentMarket):
             )
         )
 
+    @staticmethod
+    def redeem_winnings(api_keys: APIKeys) -> None:
+        redeem_from_all_user_positions(api_keys)
+
+    @staticmethod
+    def get_trade_balance(api_keys: APIKeys, web3: Web3 | None = None) -> xDai:
+        return get_total_balance(
+            address=api_keys.bet_from_address, web3=web3, sum_xdai=True, sum_wxdai=True
+        )
+
+    @staticmethod
+    def verify_operational_balance(api_keys: APIKeys) -> bool:
+        return get_total_balance(
+            api_keys.public_key,  # Use `public_key`, not `bet_from_address` because transaction costs are paid from the EOA wallet.
+            sum_wxdai=False,
+        ) > xdai_type(0.001)
+
+    def store_prediction(
+        self, processed_market: ProcessedMarket, keys: APIKeys
+    ) -> None:
+        reasoning = (
+            processed_market.answer.reasoning
+            if processed_market.answer.reasoning
+            else ""
+        )
+
+        ipfs_hash_decoded = HexBytes(HASH_ZERO)
+        if keys.enable_ipfs_upload:
+            logger.info("Storing prediction on IPFS.")
+            ipfs_hash = IPFSHandler(keys).store_agent_result(
+                IPFSAgentResult(reasoning=reasoning)
+            )
+            ipfs_hash_decoded = ipfscidv0_to_byte32(ipfs_hash)
+
+        tx_hashes = [
+            HexBytes(HexStr(i.id)) for i in processed_market.trades if i.id is not None
+        ]
+        prediction = ContractPrediction(
+            publisher=keys.public_key,
+            ipfs_hash=ipfs_hash_decoded,
+            tx_hashes=tx_hashes,
+            estimated_probability_bps=int(processed_market.answer.p_yes * 10000),
+        )
+        tx_receipt = OmenAgentResultMappingContract().add_prediction(
+            api_keys=keys,
+            market_address=Web3.to_checksum_address(self.id),
+            prediction=prediction,
+        )
+        logger.info(
+            f"Added prediction to market {self.id}. - receipt {tx_receipt['transactionHash'].hex()}."
+        )
+
     @staticmethod
     def get_bets_made_since(
         better_address: ChecksumAddress, start_time: DatetimeUTC
@@ -1222,13 +1281,12 @@ def get_binary_market_p_yes_history(market: OmenAgentMarket) -> list[Probability
     return history
 
 
-def is_minimum_required_balance(
+def get_total_balance(
     address: ChecksumAddress,
-    min_required_balance: xDai,
     web3: Web3 | None = None,
     sum_xdai: bool = True,
     sum_wxdai: bool = True,
-) -> bool:
+) -> xDai:
     """
     Checks if the total balance of xDai and wxDai in the wallet is above the minimum required balance.
     """
@@ -1239,7 +1297,7 @@ def is_minimum_required_balance(
         total_balance += current_balances.xdai
     if sum_wxdai:
         total_balance += current_balances.wxdai
-    return total_balance >= min_required_balance
+    return xdai_type(total_balance)
 
 
 def withdraw_wxdai_to_xdai_to_keep_balance(

prediction_market_agent_tooling/tools/tavily/tavily_models.py

@@ -0,0 +1,84 @@
+import typing as t
+
+from pydantic import BaseModel
+from sqlalchemy import Column
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlmodel import ARRAY, Field, SQLModel, String
+
+from prediction_market_agent_tooling.tools.utils import DatetimeUTC, utcnow
+
+
+class TavilyResult(BaseModel):
+    title: str
+    url: str
+    content: str
+    score: float
+    raw_content: str | None
+
+
+class TavilyResponse(BaseModel):
+    query: str
+    follow_up_questions: str | None = None
+    answer: str
+    images: list[str]
+    results: list[TavilyResult]
+    response_time: float
+
+
+class TavilyResponseModel(SQLModel, table=True):
+    __tablename__ = "tavily_response"
+    __table_args__ = {"extend_existing": True}
+    id: int | None = Field(None, primary_key=True)
+    agent_id: str = Field(index=True, nullable=False)
+    # Parameters used to execute the search
+    query: str = Field(index=True, nullable=False)
+    search_depth: str
+    topic: str
+    days: int | None = Field(default=None, nullable=True)
+    max_results: int
+    include_domains: list[str] | None = Field(
+        None, sa_column=Column(ARRAY(String), nullable=True)
+    )
+    exclude_domains: list[str] | None = Field(
+        None, sa_column=Column(ARRAY(String), nullable=True)
+    )
+    include_answer: bool
+    include_raw_content: bool
+    include_images: bool
+    use_cache: bool
+    # Datetime at the time of search response and response from the search
+    datetime_: DatetimeUTC = Field(index=True, nullable=False)
+    response: dict[str, t.Any] = Field(sa_column=Column(JSONB, nullable=False))
+
+    @staticmethod
+    def from_model(
+        agent_id: str,
+        query: str,
+        search_depth: t.Literal["basic", "advanced"],
+        topic: t.Literal["general", "news"],
+        days: int | None,
+        max_results: int,
+        include_domains: t.Sequence[str] | None,
+        exclude_domains: t.Sequence[str] | None,
+        include_answer: bool,
+        include_raw_content: bool,
+        include_images: bool,
+        use_cache: bool,
+        response: TavilyResponse,
+    ) -> "TavilyResponseModel":
+        return TavilyResponseModel(
+            agent_id=agent_id,
+            query=query,
+            search_depth=search_depth,
+            topic=topic,
+            days=days,
+            max_results=max_results,
+            include_domains=sorted(include_domains) if include_domains else None,
+            exclude_domains=sorted(exclude_domains) if exclude_domains else None,
+            include_answer=include_answer,
+            include_raw_content=include_raw_content,
+            include_images=include_images,
+            use_cache=use_cache,
+            datetime_=utcnow(),
+            response=response.model_dump(),
+        )

prediction_market_agent_tooling/tools/{tavily_storage/tavily_storage.py → tavily/tavily_search.py}

@@ -4,16 +4,20 @@ import tenacity
 from tavily import TavilyClient
 
 from prediction_market_agent_tooling.config import APIKeys
-from prediction_market_agent_tooling.tools.tavily_storage.tavily_models import (
+from prediction_market_agent_tooling.tools.tavily.tavily_models import (
     TavilyResponse,
-    TavilyStorage,
+    TavilyResult,
 )
+from prediction_market_agent_tooling.tools.tavily.tavily_storage import TavilyStorage
+
+DEFAULT_SCORE_THRESHOLD = 0.75  # Based on some empirical testing, anything lower wasn't very relevant to the question being asked
 
 
 def tavily_search(
     query: str,
     search_depth: t.Literal["basic", "advanced"] = "advanced",
     topic: t.Literal["general", "news"] = "general",
+    days: int | None = None,
     max_results: int = 5,
     include_domains: t.Sequence[str] | None = None,
     exclude_domains: t.Sequence[str] | None = None,
@@ -35,6 +39,7 @@ def tavily_search(
             search_depth=search_depth,
             topic=topic,
             max_results=max_results,
+            days=days,
             include_domains=include_domains,
             exclude_domains=exclude_domains,
             include_answer=include_answer,
@@ -49,6 +54,7 @@ def tavily_search(
         search_depth=search_depth,
         topic=topic,
         max_results=max_results,
+        days=days,
         include_domains=include_domains,
         exclude_domains=exclude_domains,
         include_answer=include_answer,
@@ -63,6 +69,7 @@ def tavily_search(
             query=query,
             search_depth=search_depth,
             topic=topic,
+            days=days,
             max_results=max_results,
             include_domains=include_domains,
             exclude_domains=exclude_domains,
@@ -80,6 +87,7 @@ def _tavily_search(
     query: str,
     search_depth: t.Literal["basic", "advanced"],
     topic: t.Literal["general", "news"],
+    days: int | None,
     max_results: int,
     include_domains: t.Sequence[str] | None,
     exclude_domains: t.Sequence[str] | None,
@@ -99,6 +107,7 @@ def _tavily_search(
         query=query,
         search_depth=search_depth,
         topic=topic,
+        days=days,
         max_results=max_results,
         include_domains=include_domains,
         exclude_domains=exclude_domains,
@@ -108,3 +117,20 @@ def _tavily_search(
         use_cache=use_cache,
     )
     return response
+
+
+def get_related_news_since(
+    question: str,
+    days_ago: int,
+    score_threshold: float = DEFAULT_SCORE_THRESHOLD,
+    max_results: int = 3,
+    tavily_storage: TavilyStorage | None = None,
+) -> list[TavilyResult]:
+    news = tavily_search(
+        query=question,
+        days=days_ago,
+        max_results=max_results,
+        topic="news",
+        tavily_storage=tavily_storage,
+    )
+    return [r for r in news.results if r.score > score_threshold]

prediction_market_agent_tooling/tools/{tavily_storage/tavily_models.py → tavily/tavily_storage.py}

@@ -2,96 +2,15 @@ import typing as t
 from datetime import timedelta
 
 import tenacity
-from pydantic import BaseModel
-from sqlalchemy import Column
-from sqlalchemy.dialects.postgresql import JSONB
-from sqlmodel import (
-    ARRAY,
-    Field,
-    Session,
-    SQLModel,
-    String,
-    create_engine,
-    desc,
-    select,
-)
+from sqlmodel import Session, SQLModel, create_engine, desc, select
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.loggers import logger
-from prediction_market_agent_tooling.tools.utils import DatetimeUTC, utcnow
-
-
-class TavilyResult(BaseModel):
-    title: str
-    url: str
-    content: str
-    score: float
-    raw_content: str | None
-
-
-class TavilyResponse(BaseModel):
-    query: str
-    follow_up_questions: None = None
-    answer: str
-    images: list[str]
-    results: list[TavilyResult]
-    response_time: float
-
-
-class TavilyResponseModel(SQLModel, table=True):
-    __tablename__ = "tavily_response"
-    __table_args__ = {"extend_existing": True}
-    id: int | None = Field(None, primary_key=True)
-    agent_id: str = Field(index=True, nullable=False)
-    # Parameters used to execute the search
-    query: str = Field(index=True, nullable=False)
-    search_depth: str
-    topic: str
-    max_results: int
-    include_domains: list[str] | None = Field(
-        None, sa_column=Column(ARRAY(String), nullable=True)
-    )
-    exclude_domains: list[str] | None = Field(
-        None, sa_column=Column(ARRAY(String), nullable=True)
-    )
-    include_answer: bool
-    include_raw_content: bool
-    include_images: bool
-    use_cache: bool
-    # Datetime at the time of search response and response from the search
-    datetime_: DatetimeUTC = Field(index=True, nullable=False)
-    response: dict[str, t.Any] = Field(sa_column=Column(JSONB, nullable=False))
-
-    @staticmethod
-    def from_model(
-        agent_id: str,
-        query: str,
-        search_depth: t.Literal["basic", "advanced"],
-        topic: t.Literal["general", "news"],
-        max_results: int,
-        include_domains: t.Sequence[str] | None,
-        exclude_domains: t.Sequence[str] | None,
-        include_answer: bool,
-        include_raw_content: bool,
-        include_images: bool,
-        use_cache: bool,
-        response: TavilyResponse,
-    ) -> "TavilyResponseModel":
-        return TavilyResponseModel(
-            agent_id=agent_id,
-            query=query,
-            search_depth=search_depth,
-            topic=topic,
-            max_results=max_results,
-            include_domains=sorted(include_domains) if include_domains else None,
-            exclude_domains=sorted(exclude_domains) if exclude_domains else None,
-            include_answer=include_answer,
-            include_raw_content=include_raw_content,
-            include_images=include_images,
-            use_cache=use_cache,
-            datetime_=utcnow(),
-            response=response.model_dump(),
-        )
+from prediction_market_agent_tooling.tools.tavily.tavily_models import (
+    TavilyResponse,
+    TavilyResponseModel,
+)
+from prediction_market_agent_tooling.tools.utils import utcnow
 
 
 class TavilyStorage:
@@ -119,6 +38,7 @@ class TavilyStorage:
         query: str,
         search_depth: t.Literal["basic", "advanced"],
         topic: t.Literal["general", "news"],
+        days: int | None,
         max_results: int,
         include_domains: t.Sequence[str] | None,
         exclude_domains: t.Sequence[str] | None,
@@ -134,6 +54,7 @@ class TavilyStorage:
             search_depth=search_depth,
             topic=topic,
             max_results=max_results,
+            days=days,
             include_domains=include_domains,
             exclude_domains=exclude_domains,
             include_answer=include_answer,
@@ -152,6 +73,7 @@ class TavilyStorage:
         query: str,
         search_depth: t.Literal["basic", "advanced"],
         topic: t.Literal["general", "news"],
+        days: int | None,
         max_results: int,
         include_domains: t.Sequence[str] | None,
         exclude_domains: t.Sequence[str] | None,
@@ -167,6 +89,7 @@ class TavilyStorage:
                 .where(TavilyResponseModel.query == query)
                 .where(TavilyResponseModel.search_depth == search_depth)
                 .where(TavilyResponseModel.topic == topic)
+                .where(TavilyResponseModel.days == days)
                 .where(TavilyResponseModel.max_results == max_results)
                 .where(TavilyResponseModel.include_domains == include_domains)
                 .where(TavilyResponseModel.exclude_domains == exclude_domains)

{prediction_market_agent_tooling-0.52.2.dist-info → prediction_market_agent_tooling-0.54.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prediction-market-agent-tooling
-Version: 0.52.2
+Version: 0.54.0
 Summary: Tools to benchmark, deploy and monitor prediction market agents.
 Author: Gnosis
 Requires-Python: >=3.10,<3.12

{prediction_market_agent_tooling-0.52.2.dist-info → prediction_market_agent_tooling-0.54.0.dist-info}/RECORD

@@ -17,9 +17,9 @@ prediction_market_agent_tooling/benchmark/agents.py,sha256=B1-uWdyeN4GGKMWGK_-Cc
 prediction_market_agent_tooling/benchmark/benchmark.py,sha256=MqTiaaJ3cYiOLUVR7OyImLWxcEya3Rl5JyFYW-K0lwM,17097
 prediction_market_agent_tooling/benchmark/utils.py,sha256=D0MfUkVZllmvcU0VOurk9tcKT7JTtwwOp-63zuCBVuc,2880
 prediction_market_agent_tooling/config.py,sha256=WC30Nr16RGueTafA9i67OIB-6KDHZRryhiLPzebg9_I,6740
-prediction_market_agent_tooling/deploy/agent.py,sha256=3yookl4xiLLcIzvE5h-Rf6_NmBRZuKDsJu9vFJaSCtk,22294
+prediction_market_agent_tooling/deploy/agent.py,sha256=OAfD9B0rN6k3eISvRFJpHTS_dW4yMM2GlTvG8FMejd8,19232
 prediction_market_agent_tooling/deploy/agent_example.py,sha256=dIIdZashExWk9tOdyDjw87AuUcGyM7jYxNChYrVK2dM,1001
-prediction_market_agent_tooling/deploy/betting_strategy.py,sha256=6PceKCNa4CM8ws033JFoCDb2Xm9OmtMjCtdSWXD5BpQ,12495
+prediction_market_agent_tooling/deploy/betting_strategy.py,sha256=egarsbIiFDwCLQ0YtqzP3jdht2_PfhlT0yAmwYGJ10o,13009
 prediction_market_agent_tooling/deploy/constants.py,sha256=M5ty8URipYMGe_G-RzxRydK3AFL6CyvmqCraJUrLBnE,82
 prediction_market_agent_tooling/deploy/gcp/deploy.py,sha256=CYUgnfy-9XVk04kkxA_5yp0GE9Mw5caYqlFUZQ2j3ks,3739
 prediction_market_agent_tooling/deploy/gcp/kubernetes_models.py,sha256=OsPboCFGiZKsvGyntGZHwdqPlLTthITkNF5rJFvGgU8,2582
@@ -30,22 +30,22 @@ prediction_market_agent_tooling/jobs/jobs.py,sha256=I07yh0GJ-xhlvQaOUQB8xlSnihhc
 prediction_market_agent_tooling/jobs/jobs_models.py,sha256=I5uBTHJ2S1Wi3H4jDxxU7nsswSIP9r3BevHmljLh5Pg,1370
 prediction_market_agent_tooling/jobs/omen/omen_jobs.py,sha256=I2_vGrEJj1reSI8M377ab5QCsYNp_l4l4QeYEmDBkFM,3989
 prediction_market_agent_tooling/loggers.py,sha256=Am6HHXRNO545BO3l7Ue9Wb2TkYE1OK8KKhGbI3XypVU,3751
-prediction_market_agent_tooling/markets/agent_market.py,sha256=09Guz5y5Uq0K8nKFmy8SqxNfCYJX3auyxWLzIuOek2I,11119
+prediction_market_agent_tooling/markets/agent_market.py,sha256=A72Lf7bdFOTelLhUTSHmfLBOuTZ8b0Qd_IJIgdjzeB0,12470
 prediction_market_agent_tooling/markets/categorize.py,sha256=jsoHWvZk9pU6n17oWSCcCxNNYVwlb_NXsZxKRI7vmsk,1301
 prediction_market_agent_tooling/markets/data_models.py,sha256=jMqrSFO_w2z-5N3PFVgZqTHdVdkzSDhhzky2lHsGGKA,3621
 prediction_market_agent_tooling/markets/manifold/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 prediction_market_agent_tooling/markets/manifold/api.py,sha256=Fd0HYnstvvHO6AZkp1xiRlvCwQUc8kLR8DAj6PAZu0s,7297
 prediction_market_agent_tooling/markets/manifold/data_models.py,sha256=ylXIEHymx2RcCdOVUpKP4BcTqbzLu_Fd5gV1TGBntVk,6099
-prediction_market_agent_tooling/markets/manifold/manifold.py,sha256=IeVS3Suifryym-aLFDNEyyhdT3Gg03l4O39-PDpXmE8,4895
+prediction_market_agent_tooling/markets/manifold/manifold.py,sha256=qemQIwuFg4yf6egGWFp9lWpz1lXr02QiBeZ2akcT6II,5026
 prediction_market_agent_tooling/markets/manifold/utils.py,sha256=cPPFWXm3vCYH1jy7_ctJZuQH9ZDaPL4_AgAYzGWkoow,513
 prediction_market_agent_tooling/markets/market_fees.py,sha256=Q64T9uaJx0Vllt0BkrPmpMEz53ra-hMVY8Czi7CEP7s,1227
-prediction_market_agent_tooling/markets/markets.py,sha256=_3nV9QTT48G2oJ2egkuWA1UzrTOGY6x3mXqIRgDaVIo,3245
+prediction_market_agent_tooling/markets/markets.py,sha256=mwubc567OIlA32YKqlIdTloYV8FGJia9gPv0wE0xUEA,3368
 prediction_market_agent_tooling/markets/metaculus/api.py,sha256=4TRPGytQQbSdf42DCg2M_JWYPAuNjqZ3eBqaQBLkNks,2736
 prediction_market_agent_tooling/markets/metaculus/data_models.py,sha256=2wDZ0BmK9O5Lud-q-FCzgW0tsK9GxMU0rUMlcPxSS04,3184
 prediction_market_agent_tooling/markets/metaculus/metaculus.py,sha256=d5tHhekLfPGO6CfmWWJQWpJu5HyoAWBM2aGSR874Cms,3695
 prediction_market_agent_tooling/markets/omen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 prediction_market_agent_tooling/markets/omen/data_models.py,sha256=nCjsc-ylIzQOCK_1BW-5NoYrS-NIXz2Hg9N1-IqhhC8,27516
-prediction_market_agent_tooling/markets/omen/omen.py,sha256=3BF3oH78K_puwGktLkkOIxsa3G61Jb2zx6LZ_kjr3Hk,48638
+prediction_market_agent_tooling/markets/omen/omen.py,sha256=-9u8Tb3ADGy7vd0E0peTV8OqQi9zMaXBa4b0gJXpQpc,50773
 prediction_market_agent_tooling/markets/omen/omen_contracts.py,sha256=Zq7SncCq-hvpgXKsVruGBGCn1OhKZTe7r1qLdCTrT2w,28297
 prediction_market_agent_tooling/markets/omen/omen_resolving.py,sha256=iDWdjICGkt968exwCjY-6nsnQyrrNAg3YjnDdP430GQ,9415
 prediction_market_agent_tooling/markets/omen/omen_subgraph_handler.py,sha256=zQH3iu0SVH1RmE-W3NMEpcKVMILXJYxMhL6w1wh5RUo,37348
@@ -88,12 +88,13 @@ prediction_market_agent_tooling/tools/parallelism.py,sha256=6Gou0hbjtMZrYvxjTDFU
 prediction_market_agent_tooling/tools/safe.py,sha256=h0xOO0eNtitClf0fPkn-0oTc6A_bflDTee98V_aiV-A,5195
 prediction_market_agent_tooling/tools/singleton.py,sha256=CiIELUiI-OeS7U7eeHEt0rnVhtQGzwoUdAgn_7u_GBM,729
 prediction_market_agent_tooling/tools/streamlit_user_login.py,sha256=NXEqfjT9Lc9QtliwSGRASIz1opjQ7Btme43H4qJbzgE,3010
-prediction_market_agent_tooling/tools/tavily_storage/tavily_models.py,sha256=99S7w8BvnJRMOnUArGN0g4GVRoG8M0C-XyIFU8HnLn0,6374
-prediction_market_agent_tooling/tools/tavily_storage/tavily_storage.py,sha256=xrtQH9v5pXycBRyc5j45pWqkSffkoc9efNIU1_G633Q,3706
+prediction_market_agent_tooling/tools/tavily/tavily_models.py,sha256=Rz4tZzwCRzPaq49SFT33SCRQrqHXtqWdD9ajb2tGCWc,2723
+prediction_market_agent_tooling/tools/tavily/tavily_search.py,sha256=MK_ozeQbJ014HGiKFPDScjFYq0OGcjY1KPgc9A6qO0M,4511
+prediction_market_agent_tooling/tools/tavily/tavily_storage.py,sha256=t-tZzbCzBBdFedRZDuVBn3A3mIDX8Z5wza6SxWswu_E,4093
 prediction_market_agent_tooling/tools/utils.py,sha256=W-9SqeCKd51BYMRhDjYPQ7lfNO_zE9EvYpmu2r5WXGA,7163
 prediction_market_agent_tooling/tools/web3_utils.py,sha256=dkcjG-LtuaWRh7WEMzRGmZ5B5rsxZTlliFOI6fj-EJ8,11842
-prediction_market_agent_tooling-0.52.2.dist-info/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
-prediction_market_agent_tooling-0.52.2.dist-info/METADATA,sha256=OKILmVSPUD6h9Z8fyfp1ymB7FJF7WcLMQa5jkj9F-fk,8056
-prediction_market_agent_tooling-0.52.2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-prediction_market_agent_tooling-0.52.2.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
-prediction_market_agent_tooling-0.52.2.dist-info/RECORD,,
+prediction_market_agent_tooling-0.54.0.dist-info/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
+prediction_market_agent_tooling-0.54.0.dist-info/METADATA,sha256=iAGimZ2Kjtbx94FpzKsdYXKvhFO9e3Bd6xaAG5eQVg4,8056
+prediction_market_agent_tooling-0.54.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+prediction_market_agent_tooling-0.54.0.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
+prediction_market_agent_tooling-0.54.0.dist-info/RECORD,,