chore(llm): optimize code

Author: CocDap

packages/llm/src/agent/api.ts

@@ -59,11 +59,11 @@ export class PolkadotAgentApi implements IPolkadotAgentApi {
   }
 
   getNativeBalanceTool(address: string): BalanceTool {
-    return checkBalanceTool(this.api.getAllApis(), address) as unknown as BalanceTool
+    return checkBalanceTool(this.api, address) as unknown as BalanceTool
   }
 
   transferNativeTool(signer: PolkadotSigner): TransferTool {
-    return transferNativeTool(this.api.getAllApis(), signer) as unknown as TransferTool
+    return transferNativeTool(this.api, signer) as unknown as TransferTool
   }
 
   xcmTransferNativeTool(signer: PolkadotSigner, sender: string): XcmTransferNativeAssetTool {

packages/llm/src/langchain/balance.ts

@@ -1,26 +1,27 @@
 import { tool } from "@langchain/core/tools"
-import type { Api, KnownChainId } from "@polkadot-agent-kit/common"
+import type { KnownChainId } from "@polkadot-agent-kit/common"
 import { formatBalance } from "@polkadot-agent-kit/common"
+import type { PolkadotApi } from "@polkadot-agent-kit/core"
 import { getNativeBalance } from "@polkadot-agent-kit/core"
 import type { z } from "zod"
 
 import type { BalanceToolResult, balanceToolSchema } from "../types"
 import { toolConfigBalance } from "../types/balance"
 import { ToolNames } from "../types/common"
-import { executeTool, getApiForChain, validateAndFormatAddress } from "../utils"
+import { executeTool, validateAndFormatAddress } from "../utils"
 
 /**
  * Returns a tool that checks the balance of a specific address
  * @param apis - Map of chain IDs to API instances
  * @param address - The address to check the balance for
  * @returns A dynamic structured tool that checks the balance of the specified address
  */
-export const checkBalanceTool = (apis: Map<KnownChainId, Api<KnownChainId>>, address: string) => {
+export const checkBalanceTool = (polkadotApi: PolkadotApi, address: string) => {
   return tool(async ({ chain }: z.infer<typeof balanceToolSchema>) => {
     return executeTool<BalanceToolResult>(
       ToolNames.CHECK_BALANCE,
       async () => {
-        const api = getApiForChain(apis, chain)
+        const api = polkadotApi.getApi(chain as KnownChainId)
         const formattedAddress = validateAndFormatAddress(address, chain as KnownChainId)
         const balanceInfo = await getNativeBalance(api, formattedAddress)
         const formattedBalance = formatBalance(balanceInfo.balance, balanceInfo.decimals)

packages/llm/src/langchain/transfer.ts

@@ -1,29 +1,27 @@
 import { tool } from "@langchain/core/tools"
-import type { Api, KnownChainId } from "@polkadot-agent-kit/common"
+import type { KnownChainId } from "@polkadot-agent-kit/common"
 import { getDecimalsByChainId, parseUnits } from "@polkadot-agent-kit/common"
+import type { PolkadotApi } from "@polkadot-agent-kit/core"
 import { submitTxWithPolkadotSigner, transferNativeCall } from "@polkadot-agent-kit/core"
 import type { PolkadotSigner } from "polkadot-api/signer"
 import type { z } from "zod"
 
 import type { TransferToolResult, transferToolSchema } from "../types"
 import { ToolNames } from "../types/common"
 import { toolConfigTransferNative } from "../types/transfer"
-import { executeTool, getApiForChain, validateAndFormatMultiAddress } from "../utils"
+import { executeTool, validateAndFormatMultiAddress } from "../utils"
 
 /**
  * Returns a tool that transfers native tokens to a specific address
  * @param api - The API instance to use for the transfer
  * @returns A dynamic structured tool that transfers native tokens to the specified address
  */
-export const transferNativeTool = (
-  apis: Map<KnownChainId, Api<KnownChainId>>,
-  signer: PolkadotSigner
-) => {
+export const transferNativeTool = (polkadotApi: PolkadotApi, signer: PolkadotSigner) => {
   return tool(async ({ amount, to, chain }: z.infer<typeof transferToolSchema>) => {
     return executeTool<TransferToolResult>(
       ToolNames.TRANSFER_NATIVE,
       async () => {
-        const api = getApiForChain(apis, chain)
+        const api = polkadotApi.getApi(chain as KnownChainId)
         const formattedAddress = validateAndFormatMultiAddress(to, chain as KnownChainId)
         const parsedAmount = parseUnits(amount, getDecimalsByChainId(chain))
         const tx = await submitTxWithPolkadotSigner(

packages/llm/src/utils/tools.ts

@@ -1,4 +1,4 @@
-import type { Api, KnownChainId } from "@polkadot-agent-kit/common"
+import type { KnownChainId } from "@polkadot-agent-kit/common"
 import { convertAddress, toMultiAddress } from "@polkadot-agent-kit/core"
 import type { MultiAddress } from "@polkadot-api/descriptors"
 
@@ -12,18 +12,6 @@ import { ChainNotAvailableError, ErrorCodes, InvalidAddressError, isAnyToolError
  */
 const generateToolCallId = (prefix: string): string => `${prefix}_${Date.now()}`
 
-export const getApiForChain = (
-  apis: Map<KnownChainId, Api<KnownChainId>>,
-  chain: string
-): Api<KnownChainId> => {
-  const api = apis.get(chain as KnownChainId)
-  if (!api) {
-    const availableChains = Array.from(apis.keys())
-    throw new ChainNotAvailableError(chain, availableChains)
-  }
-  return api
-}
-
 export const validateAndFormatAddress = (address: string, chain: KnownChainId): string => {
   const formattedAddress = convertAddress(address, chain)
   if (!formattedAddress) {

packages/sdk/src/api.ts

@@ -1,8 +1,4 @@
-import type {
-  AgentConfig,
-  Api,
-  KnownChainId
-} from "@polkadot-agent-kit/common"
+import type { AgentConfig, Api, KnownChainId } from "@polkadot-agent-kit/common"
 import { getAllSupportedChains, getChainById } from "@polkadot-agent-kit/common"
 import type { IPolkadotApi } from "@polkadot-agent-kit/core"
 import { PolkadotApi } from "@polkadot-agent-kit/core"