Update dsn cfg with network id - clean coding

autonomys · Feb 3, 2025 · fa0ebef · fa0ebef
1 parent 7820686
commit fa0ebef
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 13 deletions.
diff --git a/characters/character.example/config/config.example.yaml b/characters/character.example/config/config.example.yaml
@@ -12,6 +12,7 @@ twitter:
 
 auto_drive:
   upload: false
+  network: 'taurus' # or 'mainnet'
 
 llm:
   nodes:

diff --git a/src/agents/tools/utils/dsn/api.ts b/src/agents/tools/utils/dsn/api.ts
@@ -0,0 +1,7 @@
+import { createAutoDriveApi } from '@autonomys/auto-drive';
+import { config } from '../../../../config/index.js';
+
+export const dsnApi = createAutoDriveApi({
+  apiKey: config.autoDriveConfig.AUTO_DRIVE_API_KEY || '',
+  network: config.autoDriveConfig.AUTO_DRIVE_NETWORK,
+});
diff --git a/src/agents/tools/utils/dsn/dsnDownload.ts b/src/agents/tools/utils/dsn/dsnDownload.ts
@@ -1,6 +1,5 @@
-import { createAutoDriveApi, downloadFile } from '@autonomys/auto-drive';
-import { NetworkId } from '@autonomys/auto-utils';
-import { config } from '../../../../config/index.js';
+import { downloadFile } from '@autonomys/auto-drive';
+import { dsnApi } from './api.js';
 import { createLogger } from '../../../../utils/logger.js';
 import { withRetry } from './retry.js';
 
@@ -16,12 +15,8 @@ interface BaseMemory {
 export const download = async (cid: string): Promise<BaseMemory> => {
   return withRetry(
     async () => {
-      const api = createAutoDriveApi({
-        apiKey: config.autoDriveConfig.AUTO_DRIVE_API_KEY || '',
-        network: NetworkId.TAURUS,
-      });
       logger.info(`Downloading file: ${cid}`);
-      const stream = await downloadFile(api, cid);
+      const stream = await downloadFile(dsnApi, cid);
 
       const chunks: Uint8Array[] = [];
       for await (const chunk of stream) {

diff --git a/src/agents/tools/utils/dsn/dsnUpload.ts b/src/agents/tools/utils/dsn/dsnUpload.ts
@@ -6,13 +6,9 @@ import { agentVersion, config } from '../../../../config/index.js';
 import { signMessage, wallet } from '../blockchain/agentWallet.js';
 import { getLastMemoryCid, setLastMemoryHash } from '../blockchain/agentMemoryContract.js';
 import { withRetry } from './retry.js';
-import { NetworkId } from '@autonomys/auto-utils';
+import { dsnApi } from './api.js';
 
 const logger = createLogger('dsn-upload-tool');
-const dsnApi = createAutoDriveApi({
-  apiKey: config.autoDriveConfig.AUTO_DRIVE_API_KEY || '',
-  network: NetworkId.TAURUS,
-});
 
 let currentNonce = await wallet.getNonce();
 

diff --git a/src/config/index.ts b/src/config/index.ts
@@ -108,6 +108,7 @@ export const config = (() => {
         AUTO_DRIVE_API_KEY: process.env.AUTO_DRIVE_API_KEY,
         AUTO_DRIVE_ENCRYPTION_PASSWORD: process.env.AUTO_DRIVE_ENCRYPTION_PASSWORD,
         AUTO_DRIVE_UPLOAD: yamlConfig.auto_drive.upload ?? false,
+        AUTO_DRIVE_NETWORK: yamlConfig.auto_drive.network,
       },
       blockchainConfig: {
         RPC_URL: process.env.RPC_URL || undefined,

diff --git a/src/config/schema.ts b/src/config/schema.ts
@@ -1,5 +1,6 @@
 import { z } from 'zod';
 import { LLMProvider } from '../services/llm/types.js';
+import { NetworkId } from '@autonomys/auto-utils';
 
 const twitterConfigSchema = z.object({
   USERNAME: z.string().min(1, 'Twitter username is required'),
@@ -98,6 +99,10 @@ const autoDriveConfigSchema = z.object({
   AUTO_DRIVE_API_KEY: z.string().optional(),
   AUTO_DRIVE_ENCRYPTION_PASSWORD: z.string().optional(),
   AUTO_DRIVE_UPLOAD: z.boolean(),
+  AUTO_DRIVE_NETWORK: z
+    .enum(['mainnet', 'taurus'])
+    .transform(val => NetworkId[val.toUpperCase() as 'MAINNET' | 'TAURUS'])
+    .default('taurus'),
 });
 
 const blockchainConfigSchema = z.object({
-Original file line number
+Diff line change
@@ Expand Up / @@ -12,6 +12,7 @@ twitter: @@
     auto_drive:
       upload: false
+      network: 'taurus' # or 'mainnet'
     llm:
       nodes:
@@ Expand Down @@