fix: jupiter plugin (#74)

* Fix Jupiter plugin - Execute tx - Add convert to base units tool - Add ability to set custom SLP tokens * Add changesets
goat-sdk · Dec 16, 2024 · a66ceec · a66ceec
1 parent 786e966
commit a66ceec
Show file tree

Hide file tree

Showing 18 changed files with 249 additions and 88 deletions.
diff --git a/typescript/.changeset/clean-wombats-tell.md b/typescript/.changeset/clean-wombats-tell.md
@@ -0,0 +1,10 @@
+---
+"@goat-sdk/adapter-vercel-ai": patch
+"@goat-sdk/plugin-spl-token": patch
+"goat-examples-vercel-ai-solana": patch
+"@goat-sdk/plugin-jupiter": patch
+"@goat-sdk/wallet-solana": patch
+"@goat-sdk/core": patch
+---
+
+Fix jupiter swap tokens tx
diff --git a/typescript/examples/vercel-ai/solana/.env.template b/typescript/examples/vercel-ai/solana/.env.template
@@ -1,2 +1,3 @@
 OPENAI_API_KEY=
-WALLET_MNEMONIC=
+SOLANA_RPC_URL=
+SOLANA_PRIVATE_KEY=
diff --git a/typescript/examples/vercel-ai/solana/index.ts b/typescript/examples/vercel-ai/solana/index.ts
@@ -6,35 +6,30 @@ import { solana } from "@goat-sdk/wallet-solana";
 
 import { sendSOL } from "@goat-sdk/core";
 import { Connection, Keypair } from "@solana/web3.js";
-import * as bip39 from "bip39";
 
-require("dotenv").config();
-
-const connection = new Connection("https://api.mainnet-beta.solana.com", "confirmed");
+import { jupiter } from "@goat-sdk/plugin-jupiter";
+import { splToken } from "@goat-sdk/plugin-spl-token";
+import base58 from "bs58";
 
-const mnemonic = process.env.WALLET_MNEMONIC;
-
-if (!mnemonic) {
-    throw new Error("WALLET_MNEMONIC is not set in the environment");
-}
+require("dotenv").config();
 
-const seed = bip39.mnemonicToSeedSync(mnemonic);
-const keypair = Keypair.fromSeed(Uint8Array.from(seed).subarray(0, 32));
+const connection = new Connection(process.env.SOLANA_RPC_URL as string);
+const keypair = Keypair.fromSecretKey(base58.decode(process.env.SOLANA_PRIVATE_KEY as string));
 
 (async () => {
     const tools = await getOnChainTools({
         wallet: solana({
             keypair,
             connection,
         }),
-        plugins: [sendSOL()],
+        plugins: [sendSOL(), jupiter({ connection }), splToken({ connection, network: "mainnet" })],
     });
 
     const result = await generateText({
         model: openai("gpt-4o-mini"),
         tools: tools,
         maxSteps: 5,
-        prompt: "send 0.0001 SOL to <recipient_address>",
+        prompt: "Swap 0.05 USDC to GOAT, return the transaction hash, make sure you check i have enough USDC to cover the swap",
     });
 
     console.log(result.text);

diff --git a/typescript/examples/vercel-ai/solana/package.json b/typescript/examples/vercel-ai/solana/package.json
@@ -12,11 +12,13 @@
         "@ai-sdk/openai": "^1.0.4",
         "@goat-sdk/adapter-vercel-ai": "workspace:*",
         "@goat-sdk/core": "workspace:*",
-        "@goat-sdk/plugin-erc20": "workspace:*",
+        "@goat-sdk/plugin-spl-token": "workspace:*",
+        "@goat-sdk/plugin-jupiter": "workspace:*",
         "@goat-sdk/wallet-solana": "workspace:*",
         "@solana/web3.js": "catalog:",
         "ai": "catalog:",
         "bip39": "^3.1.0",
+        "bs58": "^6.0.0",
         "dotenv": "^16.4.5"
     }
 }
diff --git a/typescript/packages/adapters/vercel-ai/src/index.ts b/typescript/packages/adapters/vercel-ai/src/index.ts
@@ -1,10 +1,4 @@
-import {
-    type GetToolsParams,
-    type Tool,
-    type WalletClient,
-    addParametersToDescription,
-    getTools,
-} from "@goat-sdk/core";
+import { type GetToolsParams, type Tool, type WalletClient, getTools } from "@goat-sdk/core";
 
 import { type CoreTool, tool } from "ai";
 import type { z } from "zod";
@@ -24,7 +18,7 @@ export async function getOnChainTools<TWalletClient extends WalletClient>({
 
     for (const t of tools) {
         aiTools[t.name] = tool({
-            description: addParametersToDescription(t.description, t.parameters),
+            description: t.description,
             parameters: t.parameters,
             execute: async (arg: z.output<typeof t.parameters>) => {
                 return await t.method(arg);

diff --git a/typescript/packages/core/src/wallets/solana.ts b/typescript/packages/core/src/wallets/solana.ts
@@ -7,6 +7,7 @@ export function isSolanaWalletClient(wallet: WalletClient): wallet is SolanaWall
 
 export type SolanaTransaction = {
     instructions: TransactionInstruction[];
+    addressLookupTableAddresses?: string[];
 };
 
 export type SolanaReadRequest = {

diff --git a/typescript/packages/plugins/jupiter/src/parameters.ts b/typescript/packages/plugins/jupiter/src/parameters.ts
@@ -8,9 +8,9 @@ import {
 import { z } from "zod";
 
 export const getQuoteParametersSchema: z.ZodType<QuoteGetRequest> = z.object({
-    inputMint: z.string().describe("The token to swap from"),
-    outputMint: z.string().describe("The token to swap to"),
-    amount: z.number().describe("The amount of tokens to swap"),
+    inputMint: z.string().describe("The token address of the token to swap from"),
+    outputMint: z.string().describe("The token address of the token to swap to"),
+    amount: z.number().describe("The amount of tokens to swap in the tokens base unit"),
     slippageBps: z.number().optional().describe("The slippage in bps"),
     autoSlippage: z.boolean().optional().describe("Whether to use auto slippage"),
     autoSlippageCollisionUsdValue: z.number().optional().describe("The collision USD value for auto slippage"),
@@ -44,11 +44,11 @@ export const swapInfoSchema: z.ZodType<SwapInfo> = z.object({
 });
 
 export const quoteResponseSchema: z.ZodType<QuoteResponse> = z.object({
-    inputMint: z.string().describe("The token to swap from"),
-    inAmount: z.string().describe("The amount of tokens to swap"),
-    outputMint: z.string().describe("The token to swap to"),
-    outAmount: z.string().describe("The amount of tokens to swap"),
-    otherAmountThreshold: z.string().describe("The amount of tokens to swap"),
+    inputMint: z.string().describe("The token address of the token to swap from"),
+    inAmount: z.string().describe("The amount of tokens to swap in the tokens base unit"),
+    outputMint: z.string().describe("The token address of the token to swap to"),
+    outAmount: z.string().describe("The amount of tokens to swap in the tokens base unit"),
+    otherAmountThreshold: z.string().describe("The amount of tokens to swap in the tokens base unit"),
     swapMode: z.enum(["ExactIn", "ExactOut"]).describe("The swap mode"),
     slippageBps: z.number().describe("The slippage in bps"),
     computedAutoSlippage: z.number().optional().describe("The computed auto slippage"),

diff --git a/typescript/packages/plugins/jupiter/src/plugin.ts b/typescript/packages/plugins/jupiter/src/plugin.ts
@@ -1,10 +1,14 @@
 import type { Plugin, SolanaWalletClient } from "@goat-sdk/core";
 import { createJupiterApiClient } from "@jup-ag/api";
-import { TransactionMessage, VersionedTransaction } from "@solana/web3.js";
+import { ComputeBudgetProgram, type Connection, PublicKey, TransactionInstruction } from "@solana/web3.js";
 import type { z } from "zod";
-import { getQuoteParametersSchema, quoteResponseSchema } from "./parameters";
+import { getQuoteParametersSchema } from "./parameters";
 
-export function jupiter(): Plugin<SolanaWalletClient> {
+export function jupiter({
+    connection,
+}: {
+    connection: Connection;
+}): Plugin<SolanaWalletClient> {
     return {
         name: "jupiter",
         supportsSmartWallets: () => false,
@@ -15,33 +19,62 @@ export function jupiter(): Plugin<SolanaWalletClient> {
                     name: "get_quote",
                     description: "This {{tool}} gets a quote for a swap on the Jupiter DEX.",
                     parameters: getQuoteParametersSchema,
-                    method: (parameters: z.infer<typeof getQuoteParametersSchema>) =>
-                        createJupiterApiClient().quoteGet(parameters),
+                    method: async (parameters: z.infer<typeof getQuoteParametersSchema>) => {
+                        // get the token info
+                        try {
+                            const response = await createJupiterApiClient().quoteGet(parameters);
+
+                            console.log(JSON.stringify(response, null, 2));
+                            return response;
+                        } catch (error: unknown) {
+                            if (error && typeof error === "object" && "response" in error) {
+                                const response = error.response as Response;
+
+                                const result = await response.json();
+                                console.error(result);
+                                throw new Error(`Failed to get quote: ${result.error}`);
+                            }
+                            throw error;
+                        }
+                    },
                 },
                 {
-                    name: "get_swap_transaction",
-                    description: "This {{tool}} returns a transaction to swap tokens on the Jupiter DEX.",
-                    parameters: quoteResponseSchema,
-                    method: async (parameters: z.infer<typeof quoteResponseSchema>) => {
-                        const response = await createJupiterApiClient().swapPost({
-                            swapRequest: {
-                                userPublicKey: walletClient.getAddress(),
-                                quoteResponse: parameters,
-                            },
+                    name: "swap_tokens",
+                    description: "This {{tool}} swaps tokens on the Jupiter DEX",
+                    parameters: getQuoteParametersSchema,
+                    method: async (parameters: z.infer<typeof getQuoteParametersSchema>) => {
+                        const jupiterApiClient = createJupiterApiClient();
+                        const quoteResponse = await jupiterApiClient.quoteGet(parameters);
+                        const { swapInstruction, addressLookupTableAddresses } =
+                            await jupiterApiClient.swapInstructionsPost({
+                                swapRequest: {
+                                    userPublicKey: walletClient.getAddress(),
+                                    quoteResponse: quoteResponse,
+                                },
+                            });
+
+                        const deserializedInstruction = new TransactionInstruction({
+                            programId: new PublicKey(swapInstruction.programId),
+                            keys: swapInstruction.accounts.map((key) => ({
+                                pubkey: new PublicKey(key.pubkey),
+                                isSigner: key.isSigner,
+                                isWritable: key.isWritable,
+                            })),
+                            data: Buffer.from(swapInstruction.data, "base64"),
+                        });
+
+                        // TODO: Make this dynamic?
+                        const computeBudgetIx = ComputeBudgetProgram.setComputeUnitLimit({
+                            units: 400000, // Adjust this value as needed
                         });
 
-                        const serializedTransaction = response.swapTransaction;
-                        const deserializedTransaction = VersionedTransaction.deserialize(
-                            Buffer.from(serializedTransaction, "base64"),
-                        );
-                        const instructions = TransactionMessage.decompile(deserializedTransaction.message).instructions;
+                        const { hash } = await walletClient.sendTransaction({
+                            instructions: [computeBudgetIx, deserializedInstruction],
+                            addressLookupTableAddresses,
+                        });
 
                         return {
-                            serializedTransaction,
-                            instructions,
-                            lastValidBlockHeight: response.lastValidBlockHeight,
-                            prioritizationFeeLamports: response.prioritizationFeeLamports,
-                            dynamicSlippageReport: response.dynamicSlippageReport,
+                            hash,
                         };
                     },
                 },

diff --git a/typescript/packages/plugins/spl-token/README.md b/typescript/packages/plugins/spl-token/README.md
@@ -0,0 +1,50 @@
+# Goat SPL Token Plugin 🐐 - TypeScript
+
+SPL Token plugin for Goat. Allows you to create tools for transferring and getting the balance of SPL tokens.
+
+## Installation
+```
+npm install @goat-sdk/plugin-spl-token
+```
+
+## Usage
+
+```typescript
+import { splToken, USDC, GOAT } from "@goat-sdk/plugin-spl-token";
+
+const plugin = splToken({
+    connection,
+    network: "mainnet",
+    tokens: [USDC, GOAT],
+});
+```
+
+### Adding custom tokens
+```typescript
+import { splToken } from "@goat-sdk/plugin-spl-token";
+
+
+const plugin = splToken({
+    tokens: [
+        USDC,
+        {
+            decimals: 9,
+            symbol: "POPCAT",
+            name: "Popcat",
+            mintAddresses: {
+                "mainnet": "7GCihgDB8fe6KNjn2MYtkzZcRjQy3t9GHdC8uHYmW2hr",
+            },
+        },
+    ],
+});
+```
+
+## Goat
+
+<div align="center">
+Go out and eat some grass.
+
+[Docs](https://ohmygoat.dev) | [Examples](https://github.com/goat-sdk/goat/tree/main/typescript/examples) | [Discord](https://discord.gg/goat-sdk)</div>
+
+## Goat 🐐
+Goat 🐐 (Great Onchain Agent Toolkit) is an open-source library enabling AI agents to interact with blockchain protocols and smart contracts via their own wallets.
diff --git a/typescript/packages/plugins/spl-token/src/methods/convert-to-base-unit.ts b/typescript/packages/plugins/spl-token/src/methods/convert-to-base-unit.ts
@@ -0,0 +1,4 @@
+export async function convertToBaseUnit(amount: number, decimals: number) {
+    const baseUnit = amount * 10 ** decimals;
+    return baseUnit;
+}
diff --git a/typescript/packages/plugins/spl-token/src/parameters.ts b/typescript/packages/plugins/spl-token/src/parameters.ts
@@ -1,8 +1,7 @@
 import { z } from "zod";
-import { splTokenSymbolSchema } from "./tokens";
 
 export const getTokenMintAddressBySymbolParametersSchema = z.object({
-    symbol: splTokenSymbolSchema.describe("The symbol of the token to get the mint address of"),
+    symbol: z.string().describe("The symbol of the token to get the mint address of (e.g USDC, GOAT, SOL)"),
 });
 
 export const getTokenBalanceByMintAddressParametersSchema = z.object({
@@ -15,3 +14,8 @@ export const transferTokenByMintAddressParametersSchema = z.object({
     to: z.string().describe("The address to transfer the token to"),
     amount: z.string().describe("The amount of tokens to transfer"),
 });
+
+export const convertToBaseUnitParametersSchema = z.object({
+    amount: z.number().describe("The amount of tokens to convert to base unit"),
+    decimals: z.number().describe("The decimals of the token"),
+});
diff --git a/typescript/packages/plugins/spl-token/src/plugin.ts b/typescript/packages/plugins/spl-token/src/plugin.ts
@@ -1,16 +1,17 @@
 import type { Plugin, SolanaWalletClient } from "@goat-sdk/core";
 import type { Connection } from "@solana/web3.js";
-import type { SolanaNetwork } from "./tokens";
+import type { SolanaNetwork, Token } from "./tokens";
+import { SPL_TOKENS } from "./tokens";
 import { getTools } from "./utils/getTools";
-
 export function splToken({
     connection,
     network,
-}: { connection: Connection; network: SolanaNetwork }): Plugin<SolanaWalletClient> {
+    tokens = SPL_TOKENS,
+}: { connection: Connection; network: SolanaNetwork; tokens?: Token[] }): Plugin<SolanaWalletClient> {
     return {
         name: "splToken",
         supportsSmartWallets: () => false,
         supportsChain: (chain) => chain.type === "solana",
-        getTools: async (walletClient: SolanaWalletClient) => getTools(walletClient, connection, network),
+        getTools: async (walletClient: SolanaWalletClient) => getTools(walletClient, connection, network, tokens),
     };
 }
diff --git a/typescript/packages/plugins/spl-token/src/tokens.ts b/typescript/packages/plugins/spl-token/src/tokens.ts
@@ -1,13 +1,8 @@
-import { z } from "zod";
-
 export type SolanaNetwork = "devnet" | "mainnet";
 
-export const splTokenSymbolSchema = z.enum(["USDC"]);
-export type SplTokenSymbol = z.infer<typeof splTokenSymbolSchema>;
-
 export type Token = {
     decimals: number;
-    symbol: SplTokenSymbol;
+    symbol: string;
     name: string;
     mintAddresses: Record<SolanaNetwork, string | null>;
 };
@@ -19,12 +14,32 @@ export type NetworkSpecificToken = Omit<Token, "mintAddresses"> & {
 
 export const USDC: Token = {
     decimals: 6,
-    symbol: splTokenSymbolSchema.Enum.USDC,
+    symbol: "USDC",
     name: "USDC",
     mintAddresses: {
         devnet: "4zMMC9srt5Ri5X14GAgXhaHii3GnPAEERYPJgZJDncDU",
         mainnet: "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
     },
 };
 
-export const SPL_TOKENS: Token[] = [USDC];
+export const GOAT: Token = {
+    decimals: 6,
+    symbol: "GOAT",
+    name: "GOAT",
+    mintAddresses: {
+        mainnet: "CzLSujWBLFsSjncfkh59rUFqvafWcY5tzedWJSuypump",
+        devnet: null,
+    },
+};
+
+export const SOL: Token = {
+    decimals: 9,
+    symbol: "SOL",
+    name: "Wrapped SOL",
+    mintAddresses: {
+        mainnet: "So11111111111111111111111111111111111111112",
+        devnet: "So11111111111111111111111111111111111111112",
+    },
+};
+
+export const SPL_TOKENS: Token[] = [USDC, GOAT, SOL];
diff --git a/typescript/packages/plugins/spl-token/src/utils/getTokenInfoBySymbol.ts b/typescript/packages/plugins/spl-token/src/utils/getTokenInfoBySymbol.ts
@@ -0,0 +1,11 @@
+import type { SolanaNetwork, Token } from "../tokens";
+
+export function getTokenInfoBySymbol(symbol: string, tokens: Token[], network: SolanaNetwork) {
+    const token = tokens.find((token) => [token.symbol, token.symbol.toLowerCase()].includes(symbol));
+    return {
+        symbol: token?.symbol,
+        mintAddress: token?.mintAddresses[network],
+        decimals: token?.decimals,
+        name: token?.name,
+    };
+}