11from contextlib import asynccontextmanager
22from decimal import Decimal
3+ from typing import Annotated
34
45import click
56from ape .cli import ConnectedProviderCommand , account_option , network_option , verbosity_option
7+ from ape .types import AddressType
68from ape_tokens import Token , tokens
9+ from pydantic import Field
710
811from uniswap_sdk import Uniswap
912
@@ -168,8 +171,10 @@ async def lifespan(server):
168171
169172 # TODO: Move this to ape-tokens?
170173 @server .tool ()
171- async def get_token_balance (token : str ) -> Decimal :
172- """Get the token balance of the user's account."""
174+ async def get_token_balance (
175+ token : Annotated [str | AddressType , Field (description = "The token symbol or address" )],
176+ ) -> Decimal :
177+ """Get the balance of `token` in the user's account."""
173178
174179 if token in ("ether" , "ETH" ):
175180 return account .balance * Decimal ("1e-18" )
@@ -183,24 +188,90 @@ async def get_token_balance(token: str) -> Decimal:
183188 )
184189
185190 @server .tool ()
186- async def get_price (ctx : Context , base : str , quote : str ) -> Decimal :
187- """Get the current price of BASE in terms of QUOTE."""
188- uni = ctx .request_context .lifespan_context
191+ async def get_price (
192+ ctx : Context ,
193+ base : Annotated [
194+ str | AddressType ,
195+ Field (description = "The token symbol or address you want to know the price of" ),
196+ ],
197+ quote : Annotated [
198+ str | AddressType ,
199+ Field (description = "The token symbol or address which the price will be expressed" ),
200+ ],
201+ ) -> Decimal :
202+ """
203+ Returns the current exchange rate between two tokens `base` and `quote`, at the current
204+ market equilibrium observed across all relevant markets in the Uniswap protocol indexed by
205+ this tool. This is not the actual price at which a swap will occur, but rather the starting
206+ price at which a trade on Uniswap will begin with, getting worse as larger amounts of `base`
207+ are swapped through to obtain `quote`, due to the mechanics of the Uniswap AMM model.
208+
209+ **Important Note**: It is only intended to use this price as a reference.
210+ """
189211
212+ uni = ctx .request_context .lifespan_context
190213 return uni .price (base , quote )
191214
192215 @server .tool ()
193216 async def swap (
194217 ctx : Context ,
195- have : str ,
196- want : str ,
197- amount_in : Decimal | None = None ,
198- max_amount_in : Decimal | None = None ,
199- amount_out : Decimal | None = None ,
200- min_amount_out : Decimal | None = None ,
201- slippage : Decimal | None = None ,
218+ have : Annotated [
219+ str | AddressType ,
220+ Field (description = "The token symbol or address you want to sell" ),
221+ ],
222+ want : Annotated [
223+ str | AddressType ,
224+ Field (description = "The token symbol or address you want to buy" ),
225+ ],
226+ amount_in : Annotated [
227+ Decimal | None ,
228+ Field (
229+ description = "The amount of `have` tokens you want to sell."
230+ " Leave empty if using `amount_out`."
231+ ),
232+ ] = None ,
233+ max_amount_in : Annotated [
234+ Decimal | None ,
235+ Field (
236+ description = "The maximum amount of `have` tokens you are willing to sell."
237+ " Leave empty to auto-compute this amount when `amount_out` is provided."
238+ ),
239+ ] = None ,
240+ amount_out : Annotated [
241+ Decimal | None ,
242+ Field (
243+ description = "The amount of `want` tokens you want to buy."
244+ " Leave empty if using `amount_in`."
245+ ),
246+ ] = None ,
247+ min_amount_out : Annotated [
248+ Decimal | None ,
249+ Field (
250+ description = "The minimum amount of `want` tokens you want to buy."
251+ " Leave empty to auto-compute this amount when `amount_in` is provided."
252+ ),
253+ ] = None ,
254+ slippage : Annotated [
255+ Decimal | None ,
256+ Field (
257+ description = """
258+ The maximum change in equilibrium price you are willing to accept.
259+ Quantity is a value-less ratio, convert to a ratio if user specifies a percent.
260+ Leave empty to use the default of `0.005` (0.5%),
261+ or when `max_amount_in`/`min_amount_out` are provided.
262+ """
263+ ),
264+ ] = None ,
202265 ) -> str :
203- """Swap HAVE for WANT using the configured swap options."""
266+ """
267+ Performs a token swap, converting an amount of have tokens into want tokens. This function
268+ is designed to execute trades on-chain and accounts for real-world dynamics such as
269+ slippage and market impact.
270+
271+ **Important Note**: This function will account for market shifts, which can be set by the
272+ `slippage`, `max_amount_in`, or `min_amount_out` parameters. Use these to protect against
273+ adverse market changes while executing the user's order.
274+ """
204275
205276 # NOTE: FastMCP doesn't actually support `Decimal` auto-casting yet
206277 if isinstance (amount_in , str ):
0 commit comments