Skip to content

Commit 4b727e7

Browse files
CocDapCocDap
committed
chore(integration-test): fix test case, query by natural language
1 parent fe28083 commit 4b727e7

File tree

6 files changed

+1222
-107
lines changed

6 files changed

+1222
-107
lines changed

‎packages/sdk/package.json‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,7 @@
5050
"polkadot-api": "^1.9.13",
5151
"zod": "^3.24.3",
5252
"@polkadot/keyring": "^13.5.1",
53-
"@polkadot/util": "^13.5.1",
54-
"@langchain/ollama": "^0.2.2"
53+
"@polkadot/util": "^13.5.1"
5554
},
5655
"devDependencies": {
5756
"@babel/plugin-syntax-import-attributes": "^7.26.0",
@@ -66,6 +65,9 @@
6665
"dotenv": "^16.4.7",
6766
"prettier": "^3.5.3",
6867
"rollup": "^4.37.0",
69-
"rollup-plugin-dts": "^6.2.1"
68+
"rollup-plugin-dts": "^6.2.1",
69+
"@langchain/ollama": "^0.2.2",
70+
"langchain": "^0.1.21",
71+
"@agentic/langchain": "^7.6.9"
7072
}
7173
}

‎packages/sdk/tests/integration-tests/ollamaAgent.ts‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import { ChatOllama } from "@langchain/ollama";
2+
import { AgentExecutor, createToolCallingAgent } from "langchain/agents";
3+
import { SYSTEM_PROMPT, sleep } from "./utils";
4+
import { PolkadotAgentKit } from "../../src/api";
5+
import { DynamicStructuredTool } from "@langchain/core/tools";
6+
import { z } from "zod";
7+
import { ChatPromptTemplate } from '@langchain/core/prompts'
8+
9+
// Wrap PolkadotAgentKit tools as LangChain-compatible tools
10+
function wrapBalanceTool(agentKit: PolkadotAgentKit) {
11+
return new DynamicStructuredTool({
12+
name: "getNativeBalanceTool",
13+
description: "Get native token balance for an address on a specific chain",
14+
schema: z.object({
15+
chain: z.string().describe("The chain to check balance on")
16+
}),
17+
func: async ({ chain }) => {
18+
const tool = agentKit.getNativeBalanceTool();
19+
const result = await tool.call({ chain });
20+
return JSON.stringify(result);
21+
}
22+
});
23+
}
24+
25+
function wrapTransferTool(agentKit: PolkadotAgentKit) {
26+
return new DynamicStructuredTool({
27+
name: "transferNativeTool",
28+
description: "Transfer native tokens to an address on a specific chain",
29+
schema: z.object({
30+
to: z.string(),
31+
amount: z.string(),
32+
chain: z.string()
33+
}),
34+
func: async ({ to, amount, chain }) => {
35+
const tool = agentKit.transferNativeTool();
36+
const result = await tool.call({ to, amount, chain });
37+
return JSON.stringify(result);
38+
}
39+
});
40+
}
41+
42+
function wrapXcmTool(agentKit: PolkadotAgentKit) {
43+
return new DynamicStructuredTool({
44+
name: "xcmTransferNativeTool",
45+
description: "Transfer native tokens across chains using XCM",
46+
schema: z.object({
47+
to: z.string(),
48+
amount: z.string(),
49+
sourceChain: z.string(),
50+
destChain: z.string()
51+
}),
52+
func: async ({ to, amount, sourceChain, destChain }) => {
53+
const tool = agentKit.xcmTransferNativeTool();
54+
const result = await tool.call({ to, amount, sourceChain, destChain });
55+
return JSON.stringify(result);
56+
}
57+
});
58+
}
59+
60+
export class OllamaAgent {
61+
private agentExecutor: AgentExecutor | undefined;
62+
63+
constructor(
64+
private agentKit: PolkadotAgentKit,
65+
private model: string = "qwen3:latest"
66+
) {}
67+
68+
async init() {
69+
// Initialize the Ollama LLM
70+
const llm = new ChatOllama({
71+
model: this.model,
72+
temperature: 0,
73+
});
74+
75+
// Prepare tools as LangChain-compatible DynamicStructuredTool instances
76+
const tools = [
77+
wrapBalanceTool(this.agentKit),
78+
wrapTransferTool(this.agentKit),
79+
wrapXcmTool(this.agentKit),
80+
];
81+
82+
// Use SYSTEM_PROMPT as the system message
83+
const agentPrompt = createToolCallingAgent({
84+
llm: llm as any,
85+
tools: tools as any,
86+
prompt: ChatPromptTemplate.fromMessages([
87+
['system', SYSTEM_PROMPT],
88+
['placeholder', '{chat_history}'],
89+
['human', '{input}'],
90+
['placeholder', '{agent_scratchpad}']
91+
]) as any
92+
});
93+
94+
this.agentExecutor = new AgentExecutor({
95+
agent: agentPrompt,
96+
tools: tools as any,
97+
verbose: true,
98+
});
99+
100+
await sleep(2000);
101+
}
102+
103+
async ask(query: string) {
104+
if (!this.agentExecutor) {
105+
throw new Error("OllamaAgent not initialized. Call init() first.");
106+
}
107+
return this.agentExecutor.invoke({ input: query });
108+
}
109+
}

‎packages/sdk/tests/integration-tests/sdk.itest.ts‎

Lines changed: 34 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -1,82 +1,53 @@
11
import { describe, it, expect, beforeAll, afterAll } from 'vitest';
22
import { PolkadotAgentKit } from '../../src/api';
33
import { RECIPIENT, AGENT_PRIVATE_KEY, sleep } from './utils';
4-
import { parseUnits} from '@polkadot-agent-kit/common';
5-
import { ChatOllama } from "@langchain/ollama";
4+
import { OllamaAgent } from './ollamaAgent';
65

7-
8-
9-
let agent: PolkadotAgentKit;
6+
let agentKit: PolkadotAgentKit;
7+
let ollamaAgent: OllamaAgent;
108

119
beforeAll(async () => {
12-
agent = new PolkadotAgentKit(AGENT_PRIVATE_KEY, { keyType: 'Sr25519' });
13-
await agent.initializeApi();
14-
// make sure the API is ready
15-
await sleep(5000);
10+
agentKit = new PolkadotAgentKit(AGENT_PRIVATE_KEY, { keyType: 'Sr25519' });
11+
await agentKit.initializeApi();
12+
ollamaAgent = new OllamaAgent(agentKit);
13+
await ollamaAgent.init();
1614
});
1715

1816
afterAll(async () => {
19-
await agent.disconnect();
17+
await agentKit.disconnect();
2018
});
2119

22-
2320
/// Note:
2421
/// We are using Ollama for testing purposes, but you can use any other model you want
2522
/// We are using Westend and Westend Asset Hub for integrations tests
26-
describe('PolkadotAgentKit Integration', () => {
27-
it('should get native balance on Westend', async () => {
28-
const tool = agent.getNativeBalanceTool();
29-
const result = await tool.call({ chain: 'west' });
30-
// make sure the transaction is confirmed
31-
await sleep(10000);
32-
const { success, data } = JSON.parse(result.content);
33-
expect(success).toBe(true);
34-
expect(data).toBeDefined();
23+
describe('PolkadotAgentKit Integration with OllamaAgent', () => {
24+
it('should get native balance on Westend using Ollama agent', async () => {
25+
const result = await ollamaAgent.ask('check balance on Westend');
26+
console.log('OllamaAgent Balance Query Result:', result);
27+
await sleep(20000);
28+
expect(result.output).toBeDefined();
3529
});
36-
37-
it('should transfer 0.001 WND to recipient on Westend', async () => {
38-
const tool = agent.transferNativeTool();
39-
const result = await tool.call({
40-
to: RECIPIENT,
41-
amount: '0.01',
42-
chain: 'west',
43-
});
44-
// make sure the transaction is confirmed
45-
await sleep(15000);
46-
const { success, data } = JSON.parse(result.content);
47-
expect(success).toBe(true);
48-
expect(data).toBeDefined();
49-
50-
30+
it('should transfer 0.001 WND to recipient on Westend using natural language', async () => {
31+
const userQuery = `transfer 0.001 WND to ${RECIPIENT} on Westend`;
32+
33+
const result = await ollamaAgent.ask(userQuery);
34+
console.log('Transfer Query Result:', result);
35+
await sleep(20000);
36+
expect(result.output).toBeDefined();
5137
});
52-
53-
it('should xcm transfer 0.001 WND from Westend to Westend Asset Hub', async () => {
54-
const tool = agent.xcmTransferNativeTool();
55-
const result = await tool.call({
56-
to: RECIPIENT,
57-
amount: '0.01',
58-
sourceChain: 'west',
59-
destChain: 'west_asset_hub',
60-
});
61-
// make sure the transaction is confirmed
62-
await sleep(15000);
63-
const { success, data } = JSON.parse(result.content);
64-
expect(success).toBe(true);
65-
expect(data).toBeDefined();
38+
it('should xcm transfer 0.001 WND from Westend to Westend Asset Hub using natural language', async () => {
39+
const userQuery = `transfer 0.001 WND to ${RECIPIENT} from Westend to Westend Asset Hub`;
40+
const result = await ollamaAgent.ask(userQuery);
41+
console.log('XCM Transfer Query Result (Westend → Asset Hub):', result);
42+
await sleep(20000);
43+
expect(result.output).toBeDefined();
6644
});
67-
68-
it('should xcm transfer 0.001 WND from Westend Asset Hub to Westend', async () => {
69-
const tool = agent.xcmTransferNativeTool();
70-
const result = await tool.call({
71-
to: RECIPIENT,
72-
amount: '0.01',
73-
sourceChain: 'west_asset_hub',
74-
destChain: 'west',
75-
});
76-
// make sure the transaction is confirmed
77-
await sleep(10000);
78-
const { success, data } = JSON.parse(result.content);
79-
expect(success).toBe(true);
80-
expect(data).toBeDefined();
45+
it('should xcm transfer 0.001 WND from Westend Asset Hub to Westend using natural language', async () => {
46+
const userQuery = `transfer 0.001 WND to ${RECIPIENT} from Westend Asset Hub to Westend`;
47+
const result = await ollamaAgent.ask(userQuery);
48+
console.log('XCM Transfer Query Result (Asset Hub → Westend):', result);
49+
await sleep(20000);
50+
expect(result.output).toBeDefined();
8151
});
52+
8253
});

‎packages/sdk/tests/integration-tests/utils.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,4 +63,5 @@ Please provide instructions, and I will assist you!`;
6363

6464
export function sleep(ms: number) {
6565
return new Promise(resolve => setTimeout(resolve, ms));
66-
}
66+
}
67+

‎packages/sdk/vitest.config.integration.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ export default defineConfig({
55
globals: true,
66
environment: 'node',
77
include: ['./tests/integration-tests/**/*.itest.ts'],
8-
testTimeout: 60000,
9-
hookTimeout: 60000
8+
testTimeout: 300000,
9+
hookTimeout: 300000
1010
}
1111
})
1212

0 commit comments

Comments
 (0)