chore(example): add ollama model for example (#56)

CocDap · web-flow · commit c96b3af45828 · 2025-06-15T22:32:39.000+07:00
+ This PR will replace for #54 + Use model factory to create multiple models such as openai, ollama, ...
diff --git a/examples/telegram-bot/.env.example b/examples/telegram-bot/.env.example
@@ -1,5 +1,5 @@
 # OpenAI API Key
-OPENAI_API_KEY=
+#OPENAI_API_KEY=
 
 # Polkadot RPC Endpoint
 WS_ENDPOINT=
@@ -10,4 +10,6 @@ PRIVATE_KEY=
 # Telegram Bot Token
 TELEGRAM_BOT_TOKEN=
 
+# 
+
 
diff --git a/examples/telegram-bot/package.json b/examples/telegram-bot/package.json
@@ -23,7 +23,9 @@
     "telegraf": "^4.16.3",
     "telegraf-safe-md-reply": "^1.0.0",
     "@langchain/core": "^0.3.40",
-    "@langchain/openai": "^0.3.17"
+    "zod": "^3.22.0",
+    "@langchain/ollama":"^0.2.2",
+    "@langchain/openai":"^0.5.12"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",
diff --git a/examples/telegram-bot/src/TelegramBot.ts b/examples/telegram-bot/src/TelegramBot.ts
@@ -1,62 +1,62 @@
-import { Telegraf } from 'telegraf';
-import { ChatOpenAI } from '@langchain/openai';
-import { Tool } from '@langchain/core/tools';
-import { setupHandlers } from './handlers';
-import { PolkadotAgentKit } from '@polkadot-agent-kit/sdk';
-import { getChainByName, KnownChainId, getAllSupportedChains } from '@polkadot-agent-kit/common';
-
+import { Telegraf } from "telegraf";
+import { setupHandlers } from "./handlers";
+import { PolkadotAgentKit } from "@polkadot-agent-kit/sdk";
+import {
+  ChatModelFactory,
+  ChatModelOptions,
+  ChatModelWithTools,
+} from "./models";
 
 interface BotConfig {
   botToken: string;
   openAiApiKey?: string;
   privateKey?: string;
-  // delegatePrivateKey?: string;
-  // chains: { url: string; name: string; apiKey: string; type: 'RelayChain' | 'ParaChain'; paraId?: number }[];
 }
 
 export class TelegramBot {
   private bot: Telegraf;
   private agent: PolkadotAgentKit;
-  private llm: ChatOpenAI;
+  private llm: ChatModelWithTools;
+
+  private initializeLLM(openAiApiKey?: string): ChatModelWithTools {
+    const options: ChatModelOptions = {
+      provider: openAiApiKey ? ("openai" as const) : ("ollama" as const),
+      modelName: openAiApiKey ? "gpt-4o-mini" : "qwen3:latest",
+      temperature: 0.7,
+      verbose: false,
+    };
+    return ChatModelFactory.create(options);
+  }
 
   constructor(config: BotConfig) {
-    const {
-      botToken,
-      openAiApiKey,
-      privateKey,
-      // delegatePrivateKey,
-      // chains,
-    } = config;
+    const { botToken, openAiApiKey, privateKey } = config;
 
     if (!botToken) {
-      throw new Error('TELEGRAM_BOT_TOKEN must be provided!');
+      throw new Error("TELEGRAM_BOT_TOKEN must be provided!");
     }
 
     this.bot = new Telegraf(botToken);
 
-    this.agent = new PolkadotAgentKit(privateKey as string, {keyType: 'Sr25519'});
-
-    this.llm = new ChatOpenAI({
-      modelName: 'gpt-4',
-      temperature: 0.7,
-      openAIApiKey: openAiApiKey,
-      streaming: true,
+    this.agent = new PolkadotAgentKit(privateKey as string, {
+      keyType: "Sr25519",
     });
+
+    this.llm = this.initializeLLM(openAiApiKey);
   }
 
   async initialize() {
     console.log("Initializing bot...");
-    
+
     try {
       // Initialize APIs first
       await this.agent.initializeApi();
-    
-      // Set up tools 
+
+      // Set up tools
       // Get balance of agent account
       const checkBalance = this.agent.getNativeBalanceTool();
       // Transfer native tokens to a recipient address on a specific chain.
       const transferNative = this.agent.transferNativeTool();
-      
+
       setupHandlers(this.bot, this.llm, {
         checkBalance: checkBalance,
         transferNative: transferNative,
@@ -69,17 +69,13 @@ export class TelegramBot {
     }
   }
 
-
-
-
   public async start(): Promise<void> {
     try {
       await this.initialize();
       await this.bot.launch();
-      console.log('Bot is running!');
-      
+      console.log("Bot is running!");
     } catch (error) {
-      console.error('Failed to start bot:', error);
+      console.error("Failed to start bot:", error);
       throw error;
     }
   }
@@ -89,7 +85,7 @@ export class TelegramBot {
       await this.agent.disconnect();
       this.bot.stop();
     } catch (error) {
-      console.error('Error during shutdown:', error);
+      console.error("Error during shutdown:", error);
     }
   }
-}
+}
diff --git a/examples/telegram-bot/src/handlers.ts b/examples/telegram-bot/src/handlers.ts
@@ -1,7 +1,7 @@
-import { Telegraf } from 'telegraf';
-import { HumanMessage, SystemMessage } from '@langchain/core/messages';
-import { ChatOpenAI } from '@langchain/openai';
-import { DynamicStructuredTool, Tool } from '@langchain/core/tools';
+import { Telegraf } from "telegraf";
+import { HumanMessage, SystemMessage } from "@langchain/core/messages";
+import { DynamicStructuredTool } from "@langchain/core/tools";
+import { ChatModelWithTools } from "./models";
 
 const SYSTEM_PROMPT = `I am a Telegram bot powered by PolkadotAgentKit. I can assist you with:
 - Transferring native tokens on specific chain (e.g., "transfer 1 WND to 5CSox4ZSN4SGLKUG9NYPtfVK9sByXLtxP4hmoF4UgkM4jgDJ on westend_asset_hub")
@@ -24,29 +24,26 @@ Please provide instructions, and I will assist you!`;
 
 export function setupHandlers(
   bot: Telegraf,
-  llm: ChatOpenAI,
+  llm: ChatModelWithTools,
   toolsByName: Record<string, DynamicStructuredTool>,
 ): void {
-
   bot.start((ctx) => {
     ctx.reply(
-      'Welcome to Polkadot Bot!\n' +
-      'I can assist you with:\n' +
-      '- Transferring native tokens  (e.g., "transfer 1 token to westend_asset_hub to 5CSox4ZSN4SGLKUG9NYPtfVK9sByXLtxP4hmoF4UgkM4jgDJ")\n' +
-      '- Checking balance (e.g., "check balance on west/polkadot/kusama")\n' +
-      '- Checking proxies (e.g., "check proxies on westend" or "check proxies")\n' +
-      'Try asking something!',
+      "Welcome to Polkadot Bot!\n" +
+        "I can assist you with:\n" +
+        '- Transferring native tokens  (e.g., "transfer 1 token to westend_asset_hub to 5CSox4ZSN4SGLKUG9NYPtfVK9sByXLtxP4hmoF4UgkM4jgDJ")\n' +
+        '- Checking balance (e.g., "check balance on west/polkadot/kusama")\n' +
+        '- Checking proxies (e.g., "check proxies on westend" or "check proxies")\n' +
+        "Try asking something!",
     );
   });
 
-
-  bot.on('text', async (ctx) => {
+  bot.on("text", async (ctx) => {
     const message = ctx.message.text;
-    
-    if (message.startsWith('/')) return;
 
-    try {
+    if (message.startsWith("/")) return;
 
+    try {
       const llmWithTools = llm.bindTools(Object.values(toolsByName));
       const messages = [
         new SystemMessage({ content: SYSTEM_PROMPT }),
@@ -55,43 +52,43 @@ export function setupHandlers(
       const aiMessage = await llmWithTools.invoke(messages);
       if (aiMessage.tool_calls && aiMessage.tool_calls.length > 0) {
         for (const toolCall of aiMessage.tool_calls) {
-          
           const selectedTool = toolsByName[toCamelCase(toolCall.name)];
           if (selectedTool) {
             const toolMessage = await selectedTool.invoke(toolCall);
             if (!toolMessage || !toolMessage.content) {
-              await ctx.reply('Tool did not return a response.');
+              await ctx.reply("Tool did not return a response.");
               return;
             }
-            const response = JSON.parse(toolMessage.content || '{}');
-            
+            const response = JSON.parse(toolMessage.content || "{}");
             if (response.error) {
               await ctx.reply(`Error: ${response.message}`);
             } else {
-              await ctx.reply(response.message || response.content || 'No message from tool.');
+              const content = JSON.parse(response.content || "{}");
+              await ctx.reply(content.data || "No message from tool.");
             }
           } else {
             console.warn(`Tool not found: ${toolCall.name}`);
             await ctx.reply(`Tool ${toolCall.name} not found.`);
           }
         }
       } else {
-        const content = String(aiMessage.content || 'No response from LLM.');
+        const content = String(aiMessage.content || "No response from LLM.");
         await ctx.reply(content);
       }
     } catch (error) {
-      console.error('Error handling message:', error);
+      console.error("Error handling message:", error);
       if (error instanceof Error) {
-        console.error('Error stack:', error.stack);
+        console.error("Error stack:", error.stack);
       }
-      await ctx.reply(`Sorry, an error occurred: ${error instanceof Error ? error.message : String(error)}`);
+      await ctx.reply(
+        `Sorry, an error occurred: ${error instanceof Error ? error.message : String(error)}`,
+      );
     }
   });
 
-
   bot.catch((err, ctx) => {
     console.error(`Error for ${ctx.updateType}`, err);
-    ctx.reply('An error occurred. Please try again.');
+    ctx.reply("An error occurred. Please try again.");
   });
 }
 
diff --git a/examples/telegram-bot/src/index.ts b/examples/telegram-bot/src/index.ts
@@ -1,21 +1,19 @@
+import { TelegramBot } from "./TelegramBot";
 
-import { TelegramBot } from './TelegramBot';
-
-import * as dotenv from 'dotenv';
+import * as dotenv from "dotenv";
 dotenv.config();
 
 async function runBot() {
   const bot = new TelegramBot({
     botToken: process.env.TELEGRAM_BOT_TOKEN!,
     openAiApiKey: process.env.OPENAI_API_KEY!,
     privateKey: process.env.PRIVATE_KEY!,
-    // delegatePrivateKey: process.env.DELEGATE_PRIVATE_KEY!,
   });
 
   await bot.start();
 
-  process.once('SIGINT', () => bot.stop());
-  process.once('SIGTERM', () => bot.stop());
+  process.once("SIGINT", () => bot.stop());
+  process.once("SIGTERM", () => bot.stop());
 }
 
-runBot();
+runBot();
diff --git a/examples/telegram-bot/src/models.ts b/examples/telegram-bot/src/models.ts
@@ -0,0 +1,47 @@
+import { BaseChatModel } from "@langchain/core/language_models/chat_models";
+import { ChatOpenAI } from "@langchain/openai";
+import { ChatOllama } from "@langchain/ollama";
+
+export type ChatProvider = "ollama" | "openai";
+
+export type ChatModelWithTools = BaseChatModel & {
+  bindTools: (tools: any[]) => any;
+};
+
+export interface ChatModelOptions {
+  provider: ChatProvider;
+  temperature?: number;
+  modelName?: string;
+  verbose?: boolean;
+}
+
+const chatModelConstructors: Record<
+  ChatProvider,
+  (options: ChatModelOptions) => ChatModelWithTools
+> = {
+  openai: ({ modelName, temperature = 0.7, verbose = false }) =>
+    new ChatOpenAI({
+      modelName: modelName ?? "gpt-4o-mini",
+      temperature,
+      streaming: true,
+      openAIApiKey: process.env.OPENAI_API_KEY!,
+      verbose,
+    }),
+  ollama: ({ modelName, temperature = 0.7, verbose = false }) =>
+    new ChatOllama({
+      model: modelName ?? "llama3",
+      temperature,
+      verbose,
+    }),
+};
+
+export class ChatModelFactory {
+  static create(options: ChatModelOptions): ChatModelWithTools {
+    const { provider } = options;
+    const constructor = chatModelConstructors[provider];
+    if (!constructor) {
+      throw new Error(`Unsupported provider: ${provider}`);
+    }
+    return constructor(options);
+  }
+}
diff --git a/packages/llm/package.json b/packages/llm/package.json
@@ -37,10 +37,8 @@
   "dependencies": {
     "@polkadot-agent-kit/common": "workspace:*",
     "@polkadot-agent-kit/core": "workspace:*",
-    "@langchain/community": "^0.3.19",
     "@langchain/core": "^0.3.40",
     "@langchain/langgraph": "^0.2.33",
-    "@langchain/openai": "^0.3.17",
     "@noble/curves": "^1.6.0",
     "@polkadot-labs/hdkd": "^0.0.13",
     "@polkadot-labs/hdkd-helpers": "^0.0.13",
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
