Merge pull request #128 from autonomys/feat/af-memory-pruning

Feat/af memory pruning

auto-agents-framework/.gitignore

@@ -6,4 +6,5 @@ config/
 .cookies/
 dsn-kol-schemas.json
 memories/
+!src/config/
 

auto-agents-framework/config/config.example.yaml

@@ -31,4 +31,7 @@ llm:
       temperature: 0.8
     response:
       size: "small"
-      temperature: 0.8 
+      temperature: 0.8 
+
+memory:
+  MAX_PROCESSED_IDS: 5000

auto-agents-framework/src/agents/workflows/kol/workflow.ts

@@ -1,6 +1,5 @@
 import { END, MemorySaver, StateGraph, START, Annotation } from '@langchain/langgraph';
 import { BaseMessage } from '@langchain/core/messages';
-import { ChatOpenAI } from '@langchain/openai';
 import { parseMessageContent } from '../utils.js';
 import { config } from '../../../config/index.js';
 import { createLogger } from '../../../utils/logger.js';
@@ -27,41 +26,44 @@ export const State = Annotation.Root({
   }),
   timelineTweets: Annotation<ReadonlySet<Tweet>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...update]),
+    reducer: (_, update) => new Set([...update]),
   }),
   mentionsTweets: Annotation<ReadonlySet<Tweet>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...update]),
+    reducer: (_, update) => new Set([...update]),
   }),
   myRecentTweets: Annotation<ReadonlySet<Tweet>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...update]),
+    reducer: (_, update) => new Set([...update]),
   }),
   myRecentReplies: Annotation<ReadonlySet<Tweet>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...update]),
+    reducer: (_, update) => new Set([...update]),
   }),
   summary: Annotation<Summary>,
   engagementDecisions: Annotation<EngagementDecision[]>({
     default: () => [],
-    reducer: (curr, update) => update,
+    reducer: (_, update) => update,
   }),
   trendAnalysisTweets: Annotation<ReadonlySet<Tweet>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...update]),
+    reducer: (_, update) => new Set([...update]),
   }),
   trendAnalysis: Annotation<TrendAnalysis>,
   dsnData: Annotation<Record<string, any>[]>({
     default: () => [],
-    reducer: (curr, update) => update,
+    reducer: (_, update) => update,
   }),
   processedTweetIds: Annotation<Set<string>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...curr, ...update]),
+    reducer: (curr, update) => {
+      const newSet = new Set([...curr, ...update]);
+      return new Set(Array.from(newSet).slice(-config.memoryConfig.MAX_PROCESSED_IDS));
+    },
   }),
   repliedToTweetIds: Annotation<Set<string>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...curr, ...update]),
+    reducer: (_, update) => new Set([...update]),
   }),
 });
 
@@ -128,10 +130,7 @@ const shouldContinue = (state: typeof State.State) => {
   const hasDsnData = state.dsnData && Object.keys(state.dsnData).length > 0;
 
   if (hasDsnData && config.autoDriveConfig.AUTO_DRIVE_UPLOAD) return 'uploadToDsnNode';
-  else {
-    state.repliedToTweetIds = new Set();
-    return END;
-  }
+  else return END;
 };
 
 // Workflow creation function
@@ -149,9 +148,9 @@ export const createWorkflow = async (nodes: Awaited<ReturnType<typeof createNode
     .addEdge('engagementNode', 'analyzeTrendNode')
     .addEdge('analyzeTrendNode', 'generateTweetNode')
     .addConditionalEdges('generateTweetNode', shouldContinue);
+
   return workflow;
 };
-
 // Workflow runner type
 type WorkflowRunner = Readonly<{
   runWorkflow: () => Promise<unknown>;
@@ -167,7 +166,6 @@ const createWorkflowRunner = async (characterFile: string): Promise<WorkflowRunn
 
   return {
     runWorkflow: async () => {
-      // Use a fixed thread ID for shared state across runs
       const threadId = 'shared_workflow_state';
       logger.info('Starting tweet response workflow', { threadId });
 

auto-agents-framework/src/config/index.ts

@@ -6,6 +6,7 @@ import { fileURLToPath } from 'url';
 import { mkdir } from 'fs/promises';
 import { llmDefaultConfig } from './llm.js';
 import { twitterDefaultConfig } from './twitter.js';
+import { memoryDefaultConfig } from './memory.js';
 import yaml from 'yaml';
 import { readFileSync } from 'fs';
 
@@ -29,7 +30,7 @@ function formatZodError(error: z.ZodError) {
   });
   return `Missing or invalid environment variables:
     \n${missingVars.join('\n')}
-    \nPlease check your .env file and ensure all required variables are set correctly.`;
+    \nPlease check your .env file and config.yaml file and ensure all required variables are set correctly.`;
 }
 
 export const agentVersion = process.env.AGENT_VERSION || '1.0.0';
@@ -75,6 +76,10 @@ export const config = (() => {
         CONTRACT_ADDRESS: process.env.CONTRACT_ADDRESS || undefined,
         PRIVATE_KEY: process.env.PRIVATE_KEY || undefined,
       },
+      memoryConfig: {
+        ...memoryDefaultConfig,
+        ...(yamlConfig.memory || {}),
+      },
       SERPAPI_API_KEY: process.env.SERPAPI_API_KEY || '',
       NODE_ENV: process.env.NODE_ENV || 'development',
     };
@@ -85,7 +90,7 @@ export const config = (() => {
       console.error('\x1b[31m%s\x1b[0m', formatZodError(error));
       console.info(
         '\x1b[36m%s\x1b[0m',
-        '\nTip: Copy .env.sample to .env and fill in the required values.',
+        '\nTip: Copy .env.sample to .env and config.example.yaml to config.yaml and fill/change the required values.',
       );
       process.exit(1);
     }

auto-agents-framework/src/config/memory.ts

@@ -0,0 +1,3 @@
+export const memoryDefaultConfig = {
+  MAX_PROCESSED_IDS: 5000,
+};

auto-agents-framework/src/config/schema.ts

@@ -126,13 +126,18 @@ const blockchainConfigSchema = z.object({
   ]),
 });
 
+const memoryConfigSchema = z.object({
+  MAX_PROCESSED_IDS: z.number().int().positive().default(5000),
+});
+
 const SERPAPI_API_KEY = z.string().optional();
 
 export const configSchema = z.object({
   twitterConfig: twitterConfigSchema,
   llmConfig: llmConfigSchema,
   autoDriveConfig: autoDriveConfigSchema,
   blockchainConfig: blockchainConfigSchema,
+  memoryConfig: memoryConfigSchema,
   SERPAPI_API_KEY: SERPAPI_API_KEY,
   NODE_ENV: z.enum(['development', 'production', 'test']),
 });

auto-agents-framework/tests/agents/workflows/kol/__fixtures__/mockState.ts

@@ -9,10 +9,20 @@ export const createMockState = () => ({
   myRecentTweets: new Set<Tweet>(),
   myRecentReplies: new Set<Tweet>(),
   trendAnalysisTweets: new Set<Tweet>(),
-  dsnData: [],
-  summary: { patterns: [], commonWords: [] },
-  trendAnalysis: { summary: '', trends: [] },
-  engagementDecisions: [],
+  dsnData: [] as Array<{ type: string; data: any }>,
+  summary: { patterns: [] as string[], commonWords: [] as string[] },
+  trendAnalysis: {
+    summary: '',
+    trends: [] as Array<{
+      topic: string;
+      description: string;
+      trendStrength: number;
+    }>,
+  },
+  engagementDecisions: [] as Array<{
+    decision: { shouldEngage: boolean; reason: string };
+    tweet: Tweet;
+  }>,
 });
 
 export const createMockTweet = (overrides = {}): Tweet => ({

auto-agents-framework/tests/agents/workflows/kol/workflow.test.ts

@@ -1,7 +1,12 @@
-import { createWorkflow, getWorkflowConfig } from '../../../../src/agents/workflows/kol/workflow';
+import {
+  createWorkflow,
+  getWorkflowConfig,
+  State,
+} from '../../../../src/agents/workflows/kol/workflow';
 import { createNodes } from '../../../../src/agents/workflows/kol/nodes';
 import { WorkflowConfig } from '../../../../src/agents/workflows/kol/types';
 import { createMockState } from './__fixtures__/mockState';
+import { config } from '../../../../src/config';
 
 jest.mock('../../../../src/services/twitter/client', () => ({
   createTwitterApi: jest.fn().mockResolvedValue({
@@ -78,12 +83,23 @@ describe('KOL Workflow', () => {
   });
 
   // Test State annotations
-  it('should properly handle state annotations', () => {
-    const state = createMockState();
+  describe('State Annotations', () => {
+    it('should properly initialize state', () => {
+      const state = createMockState();
 
-    expect(state.messages).toEqual([]);
-    expect(state.timelineTweets).toBeInstanceOf(Set);
-    expect(state.mentionsTweets).toBeInstanceOf(Set);
-    expect(state.processedTweetIds).toBeInstanceOf(Set);
+      expect(state.messages).toEqual([]);
+      expect(state.timelineTweets).toBeInstanceOf(Set);
+      expect(state.mentionsTweets).toBeInstanceOf(Set);
+      expect(state.processedTweetIds).toBeInstanceOf(Set);
+    });
+
+    it('should prune processed tweet IDs when exceeding limit', () => {
+      const state = createMockState();
+      state.processedTweetIds = new Set(['1', '2', '3']);
+      state.processedTweetIds = new Set(['4', '5']);
+
+      expect(Array.from(state.processedTweetIds)).toEqual(['4', '5']);
+      expect(state.processedTweetIds.size).toBe(config.memoryConfig.MAX_PROCESSED_IDS);
+    });
   });
 });

auto-agents-framework/tests/setup.js

@@ -24,6 +24,9 @@ jest.mock('../src/config/index.js', () => ({
       RPC_URL: 'http://mock-rpc',
       PRIVATE_KEY: '1234567890123456789012345678901234567890123456789012345678901234',
       CONTRACT_ADDRESS: '0x1234567890123456789012345678901234567890'
+    },
+    memoryConfig: {
+      MAX_PROCESSED_IDS: 2,
     }
   }
 }), { virtual: true });