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/memoryPruner.ts

@@ -0,0 +1,16 @@
+import { createLogger } from '../../../utils/logger.js';
+import { config } from '../../../config/index.js';
+
+const logger = createLogger('memory-pruning');
+const { MAX_PROCESSED_IDS } = config.memoryConfig;
+
+export const pruneProcessedIds = (ids: Set<string>): Set<string> => {
+  const prunedIds = new Set(Array.from(ids).slice(-MAX_PROCESSED_IDS));
+
+  logger.info('Pruned processed IDs', {
+    originalSize: ids.size,
+    newSize: prunedIds.size,
+  });
+
+  return prunedIds;
+};

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';
@@ -14,6 +13,7 @@ import { trendSchema, summarySchema } from './schemas.js';
 import { z } from 'zod';
 import { createPrompts } from './prompts.js';
 import { LLMFactory } from '../../../services/llm/factory.js';
+import { pruneProcessedIds } from './memoryPruner.js';
 
 export const logger = createLogger('agent-workflow');
 
@@ -27,41 +27,41 @@ 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) => pruneProcessedIds(new Set([...curr, ...update])),
   }),
   repliedToTweetIds: Annotation<Set<string>>({
     default: () => new Set(),
-    reducer: (curr, update) => new Set([...curr, ...update]),
+    reducer: (_, update) => new Set([...update]),
   }),
 });
 
@@ -128,10 +128,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 +146,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 +164,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/pruneMemory.test.ts

@@ -0,0 +1,41 @@
+import { createMockState } from './__fixtures__/mockState';
+import { pruneProcessedIds } from '../../../../src/agents/workflows/kol/memoryPruner';
+import { config } from '../../../../src/config/index';
+
+describe('Memory Pruning', () => {
+  let mockState: ReturnType<typeof createMockState>;
+
+  beforeEach(() => {
+    mockState = createMockState();
+    jest.clearAllMocks();
+  });
+
+  it('should prune processed IDs according to memory config', () => {
+    const processedIds = new Set(['1', '2', '3', '4', '5']);
+    const prunedIds = pruneProcessedIds(processedIds);
+    expect(prunedIds.size).toBe(config.memoryConfig.MAX_PROCESSED_IDS);
+    expect(Array.from(prunedIds)).toEqual(['4', '5']);
+  });
+
+  it('should not prune if under the limit', () => {
+    jest.replaceProperty(config.memoryConfig, 'MAX_PROCESSED_IDS', 5);
+    const processedIds = new Set(['1', '2']);
+    const prunedIds = pruneProcessedIds(processedIds);
+    expect(prunedIds.size).toBe(2);
+    expect(Array.from(prunedIds)).toEqual(['1', '2']);
+  });
+
+  it('should handle empty sets', () => {
+    const processedIds = new Set<string>();
+    const prunedIds = pruneProcessedIds(processedIds);
+
+    expect(prunedIds.size).toBe(0);
+    expect(Array.from(prunedIds)).toEqual([]);
+  });
+
+  it('should maintain order of IDs when pruning', () => {
+    const processedIds = new Set(['1', '2', '3', '4']);
+    const prunedIds = pruneProcessedIds(processedIds);
+    expect(Array.from(prunedIds)).toEqual(['3', '4']);
+  });
+});

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 });