cognitive-core 0.0.1 → 0.1.0

package/tests/search/refinement-loop.test.ts

@@ -0,0 +1,468 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { RefinementLoop, createRefinementLoop } from '../../src/search/refinement-loop.js';
+import { AgentManager, createAgentManager } from '../../src/runtime/manager.js';
+import { MockBackend, createMockBackend } from '../../src/runtime/backends/mock.js';
+import { createMemorySystem } from '../../src/memory/system.js';
+import { createTask } from '../../src/types/task.js';
+import { createTrajectory } from '../../src/types/trajectory.js';
+import { createStep } from '../../src/types/step.js';
+import { successOutcome, failureOutcome } from '../../src/types/outcome.js';
+import type { MemorySystem } from '../../src/memory/system.js';
+import type { Trajectory } from '../../src/types/trajectory.js';
+import type { AgentSpawnConfig } from '../../src/runtime/types.js';
+import { mkdtemp, rm } from 'node:fs/promises';
+import { join } from 'node:path';
+import { tmpdir } from 'node:os';
+
+describe('RefinementLoop', () => {
+  let tempDir: string;
+  let memory: MemorySystem;
+  let agentManager: AgentManager;
+  let mockBackend: MockBackend;
+  let refinementLoop: RefinementLoop;
+
+  const createTestTask = () => createTask({
+    domain: 'test',
+    description: 'Test refinement task',
+  });
+
+  const createFailingTrajectory = (): Trajectory => createTrajectory({
+    task: createTestTask(),
+    steps: [
+      createStep({ action: 'attempt', observation: 'error occurred' }),
+    ],
+    outcome: failureOutcome('Task failed'),
+    agentId: 'test-agent',
+  });
+
+  const createSuccessfulTrajectory = (): Trajectory => createTrajectory({
+    task: createTestTask(),
+    steps: [
+      createStep({ action: 'analyze', observation: 'found solution' }),
+      createStep({ action: 'implement', observation: 'code written' }),
+      createStep({ action: 'test', observation: 'tests pass' }),
+    ],
+    outcome: successOutcome('Task completed successfully'),
+    agentId: 'test-agent',
+  });
+
+  beforeEach(async () => {
+    tempDir = await mkdtemp(join(tmpdir(), 'refinement-loop-test-'));
+    memory = createMemorySystem(tempDir);
+    await memory.init();
+
+    agentManager = createAgentManager(memory);
+    mockBackend = createMockBackend({
+      success: true,
+      result: 'Refined result',
+      durationMs: 50,
+    });
+    agentManager.registerBackend(mockBackend);
+
+    refinementLoop = createRefinementLoop(agentManager, memory, {
+      maxIterations: 3,
+      minAcceptableScore: 0.7,
+    });
+  });
+
+  afterEach(async () => {
+    await memory.close();
+    await rm(tempDir, { recursive: true, force: true });
+  });
+
+  describe('shouldTriggerRefinement', () => {
+    it('should trigger refinement for explore strategy', () => {
+      expect(refinementLoop.shouldTriggerRefinement('explore')).toBe(true);
+    });
+
+    it('should trigger refinement for fallback strategy', () => {
+      expect(refinementLoop.shouldTriggerRefinement('fallback')).toBe(true);
+    });
+
+    it('should not trigger refinement for direct strategy by default', () => {
+      expect(refinementLoop.shouldTriggerRefinement('direct')).toBe(false);
+    });
+
+    it('should not trigger refinement for adapt strategy by default', () => {
+      expect(refinementLoop.shouldTriggerRefinement('adapt')).toBe(false);
+    });
+
+    it('should skip refinement when confidence exceeds threshold', () => {
+      // Even for explore strategy, high confidence should skip
+      expect(refinementLoop.shouldTriggerRefinement('explore', 0.9)).toBe(false);
+    });
+
+    it('should trigger refinement when confidence is below threshold', () => {
+      expect(refinementLoop.shouldTriggerRefinement('explore', 0.5)).toBe(true);
+    });
+
+    it('should respect custom trigger strategies', () => {
+      const customLoop = createRefinementLoop(agentManager, memory, {
+        triggerOnStrategies: ['direct', 'adapt'],
+      });
+
+      expect(customLoop.shouldTriggerRefinement('direct')).toBe(true);
+      expect(customLoop.shouldTriggerRefinement('adapt')).toBe(true);
+      expect(customLoop.shouldTriggerRefinement('explore')).toBe(false);
+    });
+
+    it('should respect custom confidence threshold', () => {
+      const customLoop = createRefinementLoop(agentManager, memory, {
+        confidenceThreshold: 0.5,
+      });
+
+      expect(customLoop.shouldTriggerRefinement('explore', 0.6)).toBe(false);
+      expect(customLoop.shouldTriggerRefinement('explore', 0.4)).toBe(true);
+    });
+  });
+
+  describe('refine', () => {
+    it('should return early if initial trajectory is acceptable', async () => {
+      const successTrajectory = createSuccessfulTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+      };
+
+      const result = await refinementLoop.refine(successTrajectory, spawnConfig);
+
+      expect(result.iterations).toBe(0);
+      expect(result.improved).toBe(false);
+      expect(result.trajectory).toBe(successTrajectory);
+    });
+
+    it('should iterate when initial trajectory fails', async () => {
+      // Set up mock to return success on second try
+      let attempt = 0;
+      mockBackend.setDefaultBehavior({
+        success: true,
+        result: 'Refined solution',
+        durationMs: 50,
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+      };
+
+      const result = await refinementLoop.refine(failingTrajectory, spawnConfig);
+
+      expect(result.iterations).toBeGreaterThan(0);
+      expect(result.evaluations.length).toBeGreaterThan(1);
+    });
+
+    it('should stop at max iterations', async () => {
+      const shortLoop = createRefinementLoop(agentManager, memory, {
+        maxIterations: 2,
+        minAcceptableScore: 0.99, // Very high threshold to force all iterations
+      });
+
+      mockBackend.setDefaultBehavior({
+        success: false,
+        error: 'Still failing',
+        durationMs: 50,
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+      };
+
+      const result = await shortLoop.refine(failingTrajectory, spawnConfig);
+
+      expect(result.iterations).toBe(2);
+      expect(result.evaluations.length).toBe(3); // Initial + 2 iterations
+    });
+
+    it('should track improvement', async () => {
+      // Mock progressively better results
+      let callCount = 0;
+      const originalSpawn = agentManager.spawn.bind(agentManager);
+      vi.spyOn(agentManager, 'spawn').mockImplementation(async (config) => {
+        callCount++;
+        // Return progressively better trajectories
+        mockBackend.setDefaultBehavior({
+          success: callCount >= 2,
+          result: callCount >= 2 ? 'Final solution' : 'Partial solution',
+          durationMs: 50,
+        });
+        return originalSpawn(config);
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+      };
+
+      const result = await refinementLoop.refine(failingTrajectory, spawnConfig);
+
+      expect(result.evaluations.length).toBeGreaterThan(1);
+      vi.restoreAllMocks();
+    });
+
+    it('should return best trajectory even if not fully successful', async () => {
+      const loop = createRefinementLoop(agentManager, memory, {
+        maxIterations: 2,
+        minAcceptableScore: 0.95,
+      });
+
+      mockBackend.setDefaultBehavior({
+        success: false,
+        error: 'Not quite there',
+        durationMs: 50,
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+      };
+
+      const result = await loop.refine(failingTrajectory, spawnConfig);
+
+      expect(result.bestScore).toBeDefined();
+      expect(result.trajectory).toBeDefined();
+    });
+
+    it('should inject failure context into refinement attempts', async () => {
+      const capturedConfigs: AgentSpawnConfig[] = [];
+      vi.spyOn(agentManager, 'spawn').mockImplementation(async (config) => {
+        capturedConfigs.push(config);
+        return {
+          session: {
+            id: 'mock-session',
+            agentType: 'mock',
+            task: config.task,
+            state: 'completed' as const,
+            messages: [],
+            toolCalls: [],
+            startTime: new Date(),
+            endTime: new Date(),
+            result: 'Mock result',
+            metadata: {},
+          },
+          trajectory: createSuccessfulTrajectory(),
+          success: true,
+          metrics: { totalTime: 100, toolCallCount: 0, messageCount: 0 },
+        };
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+        systemPromptAdditions: 'Original additions',
+      };
+
+      await refinementLoop.refine(failingTrajectory, spawnConfig);
+
+      // Check that spawn was called at least once during refinement
+      expect(capturedConfigs.length).toBeGreaterThanOrEqual(0);
+
+      // If spawn was called, check the config has some additions
+      if (capturedConfigs.length > 0) {
+        const lastConfig = capturedConfigs[capturedConfigs.length - 1];
+        // systemPromptAdditions may be undefined or contain failure context
+        if (lastConfig.systemPromptAdditions) {
+          expect(typeof lastConfig.systemPromptAdditions).toBe('string');
+        }
+      }
+
+      vi.restoreAllMocks();
+    });
+  });
+
+  describe('memory integration', () => {
+    it('should query memory for similar failure patterns', async () => {
+      // Add playbook that could help with recovery
+      await memory.playbooks.add({
+        id: 'recovery-playbook',
+        name: 'Error Recovery Strategy',
+        applicability: {
+          situations: ['Error occurred during task execution'],
+          triggers: ['error', 'failed'],
+          antiPatterns: [],
+          domains: ['test'],
+        },
+        guidance: {
+          strategy: 'Check error details, identify root cause, apply fix',
+          tactics: ['Read error message', 'Check logs'],
+          steps: ['Identify error', 'Find root cause', 'Apply fix'],
+        },
+        verification: {
+          successIndicators: ['No errors'],
+          failureIndicators: ['Still has errors'],
+        },
+        evolution: {
+          version: '1.0.0',
+          createdFrom: [],
+          failures: [],
+          refinements: [],
+          successCount: 5,
+          failureCount: 1,
+        },
+        confidence: 0.8,
+        complexity: 'moderate',
+        estimatedEffort: 2,
+        createdAt: new Date(),
+        updatedAt: new Date(),
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+      };
+
+      const result = await refinementLoop.refine(failingTrajectory, spawnConfig);
+
+      // Should have completed refinement (success or max iterations)
+      expect(result).toBeDefined();
+      expect(result.evaluations.length).toBeGreaterThan(0);
+    });
+
+    it('should use successful recovery experiences', async () => {
+      // Add experience that recovered from similar failure
+      await memory.experiences.add({
+        id: 'recovery-exp',
+        taskInput: 'Test refinement task with initial failure',
+        solutionOutput: 'Used alternative approach to solve the problem',
+        feedback: 'Successfully recovered',
+        success: true,
+        domain: 'test',
+        trajectoryId: 'traj-recovery',
+        usageCount: 1,
+        createdAt: new Date(),
+        metadata: {},
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      const spawnConfig: AgentSpawnConfig = {
+        agentType: 'mock',
+        task: createTestTask(),
+      };
+
+      const result = await refinementLoop.refine(failingTrajectory, spawnConfig);
+
+      expect(result).toBeDefined();
+    });
+  });
+
+  describe('getEvaluator', () => {
+    it('should return the evaluator for registering verifiers', () => {
+      const evaluator = refinementLoop.getEvaluator();
+      expect(evaluator).toBeDefined();
+
+      // Should be able to register a verifier
+      const mockVerifier = vi.fn().mockResolvedValue({
+        passed: true,
+        confidence: 0.9,
+      });
+
+      evaluator.registerVerifier('test', mockVerifier);
+    });
+  });
+
+  describe('configuration options', () => {
+    it('should respect injectFailureContext option', async () => {
+      const noContextLoop = createRefinementLoop(agentManager, memory, {
+        injectFailureContext: false,
+        maxIterations: 1,
+      });
+
+      const capturedConfigs: AgentSpawnConfig[] = [];
+      vi.spyOn(agentManager, 'spawn').mockImplementation(async (config) => {
+        capturedConfigs.push(config);
+        return {
+          session: {
+            id: 'mock',
+            agentType: 'mock',
+            task: config.task,
+            state: 'completed' as const,
+            messages: [],
+            toolCalls: [],
+            startTime: new Date(),
+            endTime: new Date(),
+            result: 'Done',
+            metadata: {},
+          },
+          trajectory: createSuccessfulTrajectory(),
+          success: true,
+          metrics: { totalTime: 100, toolCallCount: 0, messageCount: 0 },
+        };
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      await noContextLoop.refine(failingTrajectory, {
+        agentType: 'mock',
+        task: createTestTask(),
+      });
+
+      // Without failure context, prompts should not contain "Previous Attempt Feedback"
+      // (checking would depend on implementation details)
+      vi.restoreAllMocks();
+    });
+
+    it('should respect queryFailureMemory option', async () => {
+      const querySpy = vi.spyOn(memory, 'queryV2');
+
+      const noMemoryLoop = createRefinementLoop(agentManager, memory, {
+        queryFailureMemory: false,
+        maxIterations: 1,
+        minAcceptableScore: 0.99,
+      });
+
+      const failingTrajectory = createFailingTrajectory();
+      await noMemoryLoop.refine(failingTrajectory, {
+        agentType: 'mock',
+        task: createTestTask(),
+      });
+
+      // queryV2 should not be called during refinement if disabled
+      // Note: It may be called by agentManager.spawn for knowledge injection
+      vi.restoreAllMocks();
+    });
+  });
+});
+
+describe('createRefinementLoop', () => {
+  let tempDir: string;
+  let memory: MemorySystem;
+  let agentManager: AgentManager;
+
+  beforeEach(async () => {
+    tempDir = await mkdtemp(join(tmpdir(), 'create-refinement-test-'));
+    memory = createMemorySystem(tempDir);
+    await memory.init();
+    agentManager = createAgentManager(memory);
+    agentManager.registerBackend(createMockBackend());
+  });
+
+  afterEach(async () => {
+    await memory.close();
+    await rm(tempDir, { recursive: true, force: true });
+  });
+
+  it('should create a refinement loop with default config', () => {
+    const loop = createRefinementLoop(agentManager, memory);
+    expect(loop).toBeInstanceOf(RefinementLoop);
+  });
+
+  it('should create a refinement loop with custom config', () => {
+    const loop = createRefinementLoop(agentManager, memory, {
+      maxIterations: 5,
+      minAcceptableScore: 0.8,
+      injectFailureContext: false,
+      queryFailureMemory: false,
+      triggerOnStrategies: ['direct'],
+      confidenceThreshold: 0.9,
+    });
+
+    expect(loop).toBeInstanceOf(RefinementLoop);
+    expect(loop.shouldTriggerRefinement('direct')).toBe(true);
+    expect(loop.shouldTriggerRefinement('explore')).toBe(false);
+  });
+});