Solana Agent Kit + AiMo Integration

# AiMo Network Configuration
AIMO_API_KEY=your_aimo_api_key_here
AIMO_BASE_URL=https://devnet.aimo.network/api/v1
AIMO_MODEL=kimi-k2-0711-preview
 
# Your existing Solana Agent Kit variables
SOLANA_RPC_URL=https://api.mainnet-beta.solana.com
SOLANA_PRIVATE_KEY=your_private_key
# ... other agent kit vars

interface AimoMessage {
  role: 'system' | 'user' | 'assistant' | 'tool';
  content: string;
  name?: string;
}
 
export function createAimoClient(config: {
  apiKey: string;
  baseUrl: string;
  model: string;
}) {
  return {
    async chat(messages: AimoMessage[], options?: {
      temperature?: number;
      maxTokens?: number;
    }) {
      const response = await fetch(`${config.baseUrl}/chat/completions`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${config.apiKey}`,
        },
        body: JSON.stringify({
          model: config.model,
          messages,
          temperature: options?.temperature ?? 0.3,
          max_tokens: options?.maxTokens,
        }),
      });
 
      if (!response.ok) {
        const error = await response.text();
        throw new Error(`AiMo API error (${response.status}): ${error}`);
      }
 
      const data = await response.json();
      return data.choices?.[0]?.message?.content || '';
    },
 
    async chatStream(messages: AimoMessage[], options?: {
      temperature?: number;
      maxTokens?: number;
    }) {
      const response = await fetch(`${config.baseUrl}/chat/completions`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${config.apiKey}`,
        },
        body: JSON.stringify({
          model: config.model,
          messages,
          temperature: options?.temperature ?? 0.3,
          max_tokens: options?.maxTokens,
          stream: true,
        }),
      });
 
      if (!response.ok) {
        throw new Error(`AiMo API error: ${response.status}`);
      }
 
      return response.body; // Return stream for consumption
    }
  };
}

import { config } from 'dotenv';
config();
 
import { createAimoClient } from '../llm/aimoClient';
import { SolanaAgentKit } from 'solana-agent-kit'; // or your agent kit import
import { Connection, Keypair } from '@solana/web3.js';
 
export async function createAgent() {
  // Initialize AiMo client
  const aimoClient = createAimoClient({
    apiKey: process.env.AIMO_API_KEY!,
    baseUrl: process.env.AIMO_BASE_URL!,
    model: process.env.AIMO_MODEL!,
  });
 
  // Initialize Solana connection and wallet
  const connection = new Connection(process.env.SOLANA_RPC_URL!);
  const wallet = loadKeypair(process.env.SOLANA_PRIVATE_KEY!);
 
  // Create agent with AiMo as LLM provider
  // Note: Actual integration depends on your agent kit's API
  // This is a conceptual example - adapt to your kit's interface
  const agent = new SolanaAgentKit({
    connection,
    wallet,
    llm: aimoClient, // Pass AiMo client as LLM provider
  });
 
  return { agent, aimoClient };
}
 
function loadKeypair(privateKey: string): Keypair {
  // Your existing keypair loading logic
  // ...
}

import { createAgent } from './agent';
 
async function runAgent(userPrompt: string) {
  const { agent, aimoClient } = await createAgent();
  
  const messages = [
    {
      role: 'system' as const,
      content: `You are a Solana blockchain agent with these tools:
- get_sol_balance: Get SOL balance for an address
- transfer_sol: Send SOL to a recipient
- get_spl_balance: Get SPL token balance
 
Respond with JSON in this format:
{"thought": "your reasoning", "action": "tool_name", "args": {...}}
Or for final answer:
{"thought": "summary", "final": "answer to user"}`
    },
    {
      role: 'user' as const,
      content: userPrompt
    }
  ];
 
  const maxSteps = 8;
  
  for (let step = 0; step < maxSteps; step++) {
    // Get reasoning from AiMo
    const response = await aimoClient.chat(messages);
    
    try {
      const parsed = JSON.parse(response);
      
      if (parsed.final) {
        console.log('Agent:', parsed.final);
        return parsed.final;
      }
      
      if (parsed.action) {
        console.log(`Step ${step + 1}: ${parsed.thought}`);
        console.log(`Action: ${parsed.action}`, parsed.args);
        
        // Execute tool via your agent kit
        const result = await agent.executeTool(parsed.action, parsed.args);
        
        messages.push({
          role: 'assistant',
          content: response
        });
        messages.push({
          role: 'tool',
          name: parsed.action,
          content: JSON.stringify(result)
        });
      }
    } catch (err) {
      console.error('Failed to parse agent response:', err);
      break;
    }
  }
  
  return 'Agent exceeded max steps without final answer.';
}
 
// Example usage
runAgent('What is my SOL balance?').then(console.log);

// test-aimo.ts
import { createAimoClient } from './src/llm/aimoClient';
import { config } from 'dotenv';
config();
 
async function test() {
  const client = createAimoClient({
    apiKey: process.env.AIMO_API_KEY!,
    baseUrl: process.env.AIMO_BASE_URL!,
    model: process.env.AIMO_MODEL!,
  });
 
  const response = await client.chat([
    { role: 'user', content: 'Respond with "AiMo integration working" if you can read this.' }
  ]);
 
  console.log('AiMo response:', response);
}
 
test();

npx ts-node test-aimo.ts

// example-agent.ts
import { config } from 'dotenv';
config();
 
import { createAimoClient } from './src/llm/aimoClient';
import { Connection, PublicKey, LAMPORTS_PER_SOL } from '@solana/web3.js';
 
const aimoClient = createAimoClient({
  apiKey: process.env.AIMO_API_KEY!,
  baseUrl: process.env.AIMO_BASE_URL!,
  model: process.env.AIMO_MODEL!,
});
 
const connection = new Connection(process.env.SOLANA_RPC_URL!);
 
async function simpleAgent(prompt: string) {
  // Simple balance tool
  async function getBalance(address: string) {
    const pubkey = new PublicKey(address);
    const lamports = await connection.getBalance(pubkey);
    return { address, sol: lamports / LAMPORTS_PER_SOL };
  }
 
  const response = await aimoClient.chat([
    {
      role: 'system',
      content: 'You help users check Solana balances. Extract addresses from user messages.'
    },
    {
      role: 'user',
      content: prompt
    }
  ]);
 
  console.log('AiMo response:', response);
  return response;
}
 
simpleAgent('Check balance for address: 4Nd1mBQtrMJVYVfKf2PJy9NZUZdTAsp7D4xWLs4gDB4T');