|
| 1 | +use anda_core::{BoxError, FunctionDefinition, Tool}; |
| 2 | +use anda_engine::context::BaseCtx; |
| 3 | +use candid::Nat; |
| 4 | +use schemars::{schema_for, JsonSchema}; |
| 5 | +use serde::{Deserialize, Serialize}; |
| 6 | +use serde_json::{json, Value}; |
| 7 | +use std::sync::Arc; |
| 8 | + |
| 9 | +use super::ICPLedgers; |
| 10 | + |
| 11 | +/// Arguments for the balance of an account for a token |
| 12 | +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema)] |
| 13 | +pub struct BalanceOfArgs { |
| 14 | + /// ICP account address (principal) to query, e.g. "77ibd-jp5kr-moeco-kgoar-rro5v-5tng4-krif5-5h2i6-osf2f-2sjtv-kqe" |
| 15 | + pub account: String, |
| 16 | + /// Token symbol, e.g. "ICP" |
| 17 | + pub symbol: String, |
| 18 | +} |
| 19 | + |
| 20 | +/// ICP Ledger BalanceOf tool implementation |
| 21 | +#[derive(Debug, Clone)] |
| 22 | +pub struct BalanceOfTool { |
| 23 | + ledgers: Arc<ICPLedgers>, |
| 24 | + schema: Value, |
| 25 | +} |
| 26 | + |
| 27 | +impl BalanceOfTool { |
| 28 | + /// Creates a new BalanceOfTool instance |
| 29 | + pub fn new(ledgers: Arc<ICPLedgers>) -> Self { |
| 30 | + let mut schema = schema_for!(BalanceOfArgs); |
| 31 | + schema.meta_schema = None; // Remove the $schema field |
| 32 | + |
| 33 | + BalanceOfTool { |
| 34 | + ledgers, |
| 35 | + schema: json!(schema), |
| 36 | + } |
| 37 | + } |
| 38 | +} |
| 39 | + |
| 40 | +impl Tool<BaseCtx> for BalanceOfTool { |
| 41 | + const CONTINUE: bool = true; |
| 42 | + type Args = BalanceOfArgs; |
| 43 | + type Output = Nat; |
| 44 | + |
| 45 | + fn name(&self) -> String { |
| 46 | + format!("icp_ledger_balance_of") |
| 47 | + } |
| 48 | + |
| 49 | + fn description(&self) -> String { |
| 50 | + let tokens = self |
| 51 | + .ledgers |
| 52 | + .ledgers |
| 53 | + .keys() |
| 54 | + .map(|k| k.as_str()) |
| 55 | + .collect::<Vec<_>>(); |
| 56 | + format!( |
| 57 | + "Query the balance of the specified account on ICP network for the following tokens: {}", |
| 58 | + tokens.join(", ") |
| 59 | + ) |
| 60 | + } |
| 61 | + |
| 62 | + fn definition(&self) -> FunctionDefinition { |
| 63 | + FunctionDefinition { |
| 64 | + name: self.name(), |
| 65 | + description: self.description(), |
| 66 | + parameters: self.schema.clone(), |
| 67 | + strict: Some(true), |
| 68 | + } |
| 69 | + } |
| 70 | + |
| 71 | + async fn call(&self, ctx: BaseCtx, data: Self::Args) -> Result<Self::Output, BoxError> { |
| 72 | + self.ledgers.balance_of(&ctx, data).await |
| 73 | + } |
| 74 | +} |
| 75 | + |
| 76 | +#[cfg(test)] |
| 77 | +mod tests { |
| 78 | + use super::*; |
| 79 | + use candid::Principal; |
| 80 | + use std::collections::BTreeMap; |
| 81 | + |
| 82 | + #[tokio::test(flavor = "current_thread")] |
| 83 | + async fn test_icp_ledger_transfer() { |
| 84 | + let panda_ledger = Principal::from_text("druyg-tyaaa-aaaaq-aactq-cai").unwrap(); |
| 85 | + let ledgers = ICPLedgers { |
| 86 | + ledgers: BTreeMap::from([ |
| 87 | + ( |
| 88 | + String::from("ICP"), |
| 89 | + ( |
| 90 | + Principal::from_text("ryjl3-tyaaa-aaaaa-aaaba-cai").unwrap(), |
| 91 | + 8, |
| 92 | + ), |
| 93 | + ), |
| 94 | + (String::from("PANDA"), (panda_ledger, 8)), |
| 95 | + ]), |
| 96 | + from_user_subaccount: true, |
| 97 | + }; |
| 98 | + let ledgers = Arc::new(ledgers); |
| 99 | + let tool = BalanceOfTool::new(ledgers.clone()); |
| 100 | + let definition = tool.definition(); |
| 101 | + assert_eq!(definition.name, "icp_ledger_balance_of"); |
| 102 | + let s = serde_json::to_string_pretty(&definition).unwrap(); |
| 103 | + println!("{}", s); |
| 104 | + // { |
| 105 | + // "name": "icp_ledger_balance_of", |
| 106 | + // "description": "Query the balance of the specified account on ICP network for the following tokens: ICP, PANDA", |
| 107 | + // "parameters": { |
| 108 | + // "description": "Arguments for the balance of an account for a token", |
| 109 | + // "properties": { |
| 110 | + // "account": { |
| 111 | + // "description": "ICP account address (principal) to query, e.g. \"77ibd-jp5kr-moeco-kgoar-rro5v-5tng4-krif5-5h2i6-osf2f-2sjtv-kqe\"", |
| 112 | + // "type": "string" |
| 113 | + // }, |
| 114 | + // "symbol": { |
| 115 | + // "description": "Token symbol, e.g. \"ICP\"", |
| 116 | + // "type": "string" |
| 117 | + // } |
| 118 | + // }, |
| 119 | + // "required": [ |
| 120 | + // "account", |
| 121 | + // "symbol" |
| 122 | + // ], |
| 123 | + // "title": "BalanceOfArgs", |
| 124 | + // "type": "object" |
| 125 | + // }, |
| 126 | + // "strict": true |
| 127 | + // } |
| 128 | + } |
| 129 | +} |
0 commit comments