keystone-cli 1.3.0 → 2.0.1

package/src/runner/executors/memory-executor.ts

@@ -1,12 +1,15 @@
-import type { MemoryDb } from '../../db/memory-db.ts';
+import { embed } from 'ai';
+import { MemoryDb } from '../../db/memory-db.ts';
 import type { ExpressionContext } from '../../expression/evaluator.ts';
 import { ExpressionEvaluator } from '../../expression/evaluator.ts';
 import type { MemoryStep } from '../../parser/schema.ts';
+import { ConfigLoader } from '../../utils/config-loader.ts';
 import type { Logger } from '../../utils/logger.ts';
+import { getEmbeddingModel } from '../llm-adapter.ts';
 import type { StepExecutorOptions, StepResult } from './types.ts';
 
 /**
- * Execute a memory step (setting/deleting keys in memory)
+ * Execute a memory step (storing/searching embeddings in memory)
  */
 export async function executeMemoryStep(
   step: MemoryStep,
@@ -18,32 +21,45 @@ export async function executeMemoryStep(
   if (abortSignal?.aborted) {
     throw new Error('Memory operation aborted');
   }
-  const memoryDb = options.memoryDb;
-  if (!memoryDb) {
-    throw new Error('Memory database not initialized');
+  const memoryDbFromOptions = options.memoryDb;
+  if (!memoryDbFromOptions && !options.memoryDb) {
+    // We'll initialize it below if not provided
   }
 
-  const requestedModel = step.model || 'local';
-  if (!requestedModel.toLowerCase().startsWith('local')) {
-    throw new Error(`Memory steps only support local embeddings (requested: ${requestedModel})`);
-  }
+  // Get embedding model and dimension from config or step
+  const config = ConfigLoader.load();
+  const modelName = step.model || config.embedding_model;
 
-  const adapterResult = options.getAdapter ? options.getAdapter(requestedModel) : null;
-  const adapter = adapterResult?.adapter;
-  const resolvedModel = adapterResult?.resolvedModel ?? requestedModel;
-  if (!resolvedModel.toLowerCase().startsWith('local')) {
-    throw new Error(`Memory steps only support local embeddings (requested: ${resolvedModel})`);
-  }
-  if (!adapter || !adapter.embed) {
-    throw new Error(`Model ${resolvedModel} does not support embeddings`);
+  if (!modelName) {
+    throw new Error(
+      'No embedding model configured. Set embedding_model in config or specify model in step.'
+    );
   }
 
+  // Resolve provider dimension
+  const providerName = ConfigLoader.getProviderForModel(modelName);
+  const providerConfig = config.providers[providerName];
+  const dimension = providerConfig?.embedding_dimension || config.embedding_dimension || 384;
+
+  const memoryDb = memoryDbFromOptions || new MemoryDb('.keystone/memory.db', dimension);
+
+  // Helper to get embedding using AI SDK
+  const getEmbedding = async (text: string): Promise<number[]> => {
+    const model = await getEmbeddingModel(modelName);
+    const result = await embed({
+      model,
+      value: text,
+      abortSignal,
+    });
+    return result.embedding;
+  };
+
   switch (step.op) {
     case 'store': {
       const text = ExpressionEvaluator.evaluateString(step.text || '', context);
       if (!text) throw new Error('Text is required for memory store operation');
 
-      const embedding = await adapter.embed(text, undefined, { signal: abortSignal });
+      const embedding = await getEmbedding(text);
       const metadata = step.metadata || {};
       const id = await memoryDb.store(text, embedding, metadata as Record<string, unknown>);
 
@@ -56,7 +72,7 @@ export async function executeMemoryStep(
       const query = ExpressionEvaluator.evaluateString(step.query || '', context);
       if (!query) throw new Error('Query is required for memory search operation');
 
-      const embedding = await adapter.embed(query, undefined, { signal: abortSignal });
+      const embedding = await getEmbedding(query);
       const limit = step.limit || 5;
       const results = await memoryDb.search(embedding, limit);
 

package/src/runner/executors/plan-executor.ts

@@ -97,7 +97,6 @@ export async function executePlanStep(
     options.mcpManager,
     options.artifactRoot, // Note: using artifactRoot as fallback for workflowDir if not explicit
     options.abortSignal,
-    options.getAdapter,
     options.emitEvent,
     options.runId ? { runId: options.runId } : undefined
   );

package/src/runner/executors/types.ts

@@ -1,11 +1,11 @@
 import type { MemoryDb } from '../../db/memory-db.ts';
 import type { WorkflowDb } from '../../db/workflow-db.ts';
 import type { ExpressionContext } from '../../expression/evaluator.ts';
-import type { WorkflowStep } from '../../parser/schema.ts';
+import type { Step, WorkflowStep } from '../../parser/schema.ts';
 import type { Logger } from '../../utils/logger.ts';
 import type { SafeSandbox } from '../../utils/sandbox.ts';
 import type { WorkflowEvent } from '../events.ts';
-import type { getAdapter } from '../llm-adapter.ts';
+
 import type { MCPManager } from '../mcp-manager.ts';
 import type { executeLlmStep } from './llm-executor.ts';
 
@@ -62,8 +62,8 @@ export interface StepExecutorOptions {
   debug?: boolean;
   allowInsecure?: boolean;
   emitEvent?: (event: WorkflowEvent) => void;
-  getAdapter?: typeof getAdapter;
-  executeStep?: any; // To avoid circular dependency
+
+  executeStep?: (step: Step, context: ExpressionContext) => Promise<StepResult>; // To avoid circular dependency
   executeLlmStep?: typeof executeLlmStep;
   sandbox?: typeof SafeSandbox;
 }

package/src/runner/llm-adapter.integration.test.ts

@@ -0,0 +1,151 @@
+import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test';
+import { AuthManager } from '../utils/auth-manager';
+import { ConfigLoader } from '../utils/config-loader';
+import { resetLlmMocks, setupLlmMocks } from './__test__/llm-test-setup';
+import {
+  DynamicProviderRegistry,
+  getEmbeddingModel,
+  getModel,
+  resetProviderRegistry,
+} from './llm-adapter';
+
+// Mocks for AI SDK models
+const mockLanguageModel = {
+  specificationVersion: 'v1',
+  provider: 'test-provider',
+  modelId: 'test-model',
+  doGenerate: async () => ({}),
+  doStream: async () => ({}),
+} as any;
+
+const mockEmbeddingModel = {
+  specificationVersion: 'v1',
+  provider: 'test-provider',
+  modelId: 'test-embedding-model',
+  doEmbed: async () => ({}),
+  doEmbedMany: async () => ({}),
+} as any;
+
+describe('LLM Adapter (AI SDK)', () => {
+  beforeEach(() => {
+    setupLlmMocks();
+    ConfigLoader.clear();
+    resetProviderRegistry();
+    // Reset AuthManager mocks if any
+    mock.restore();
+  });
+
+  afterEach(() => {
+    resetLlmMocks();
+    mock.restore();
+  });
+
+  describe('getModel', () => {
+    it('should load a provider and return a language model', async () => {
+      // Mock Config
+      ConfigLoader.setConfig({
+        default_provider: 'test-provider',
+        providers: {
+          'test-provider': {
+            type: 'openai', // Use standard type to trigger logic
+            package: '@ai-sdk/openai', // Use real package name to match mock
+          },
+        },
+        model_mappings: {},
+      } as any);
+
+      // With shared setupLlmMocks, we expect 'mock' provider
+      const model = (await getModel('model-name')) as any;
+      expect(model.modelId).toBe('mock-model');
+      expect(model.provider).toBe('mock');
+    });
+
+    it('should handle auth token retrieval for standard providers', async () => {
+      ConfigLoader.setConfig({
+        default_provider: 'openai',
+        providers: {
+          openai: {
+            type: 'openai',
+            package: '@ai-sdk/openai',
+            api_key_env: 'OPENAI_API_KEY',
+          },
+        },
+        model_mappings: {},
+      } as any);
+
+      spyOn(ConfigLoader, 'getSecret').mockReturnValue('fake-token');
+
+      const model = (await getModel('gpt-4')) as any;
+      // With global mock, we mostly check it didn't throw and loaded the 'mock' provider
+      expect(model.provider).toBe('mock');
+      expect(ConfigLoader.getSecret).toHaveBeenCalledWith('OPENAI_API_KEY');
+    });
+  });
+
+  describe('getEmbeddingModel', () => {
+    it('should return an embedding model if supported by provider', async () => {
+      ConfigLoader.setConfig({
+        default_provider: 'embed-provider',
+        providers: {
+          'embed-provider': { type: 'custom', package: 'pkg' },
+        },
+        model_mappings: {},
+      } as any);
+
+      const mockProvider = (modelId: string) => mockLanguageModel;
+      mockProvider.textEmbeddingModel = (modelId: string) => mockEmbeddingModel;
+
+      spyOn(DynamicProviderRegistry, 'getProvider').mockResolvedValue(() => mockProvider);
+
+      const model = (await getEmbeddingModel('text-embedding-3')) as any;
+      expect(model.modelId).toBe(mockEmbeddingModel.modelId);
+    });
+
+    it('should throw if provider does not support embeddings', async () => {
+      ConfigLoader.setConfig({
+        default_provider: 'bad-provider',
+        providers: {
+          'bad-provider': { type: 'custom', package: 'pkg' },
+        },
+        model_mappings: {},
+      } as any);
+
+      const mockProvider = (modelId: string) => mockLanguageModel;
+      // No textEmbeddingModel method
+
+      spyOn(DynamicProviderRegistry, 'getProvider').mockResolvedValue(() => mockProvider);
+
+      // Use a non-default model name to avoid fallback to LocalEmbeddingModel
+      await expect(getEmbeddingModel('non-default-model')).rejects.toThrow(
+        /does not support embeddings/
+      );
+    });
+  });
+
+  describe('Tool Cases', () => {
+    it('should handle assistant response with tool calls and NO content', async () => {
+      ConfigLoader.setConfig({
+        default_provider: 'test-provider',
+        providers: { 'test-provider': { type: 'openai', package: 'test-pkg' } },
+        model_mappings: {},
+      } as any);
+
+      const mockProvider = (modelId: string) => ({
+        ...mockLanguageModel,
+        doGenerate: async () => ({
+          content: [
+            { type: 'tool-call', toolCallId: '1', toolName: 'test', args: {}, input: '{}' },
+          ],
+          finishReason: 'tool-calls',
+          usage: { promptTokens: 1, completionTokens: 1 },
+        }),
+      });
+      spyOn(DynamicProviderRegistry, 'getProvider').mockResolvedValue(() => mockProvider);
+
+      const model = (await getModel('model')) as any;
+      const result = await model.doGenerate({ input: [], prompt: [], mode: { type: 'regular' } });
+      expect(result.content).toHaveLength(1);
+      expect(result.content[0].type).toBe('tool-call');
+    });
+  });
+});