Add a screen for custom prompts

src/utils/__tests__/enhance-prompt.test.ts

@@ -1,80 +1,126 @@
 import { enhancePrompt } from '../enhance-prompt'
-import { buildApiHandler } from '../../api'
 import { ApiConfiguration } from '../../shared/api'
-import { OpenRouterHandler } from '../../api/providers/openrouter'
+import { buildApiHandler, SingleCompletionHandler } from '../../api'
+import { defaultPrompts } from '../../shared/modes'
 
-// Mock the buildApiHandler function
+// Mock the API handler
 jest.mock('../../api', () => ({
-    buildApiHandler: jest.fn()
+  buildApiHandler: jest.fn()
 }))
 
 describe('enhancePrompt', () => {
-    const mockApiConfig: ApiConfiguration = {
-        apiProvider: 'openrouter',
-        apiKey: 'test-key',
-        openRouterApiKey: 'test-key',
-        openRouterModelId: 'test-model'
-    }
+  const mockApiConfig: ApiConfiguration = {
+    apiProvider: 'openai',
+    openAiApiKey: 'test-key',
+    openAiBaseUrl: 'https://api.openai.com/v1'
+  }
 
-    // Create a mock handler that looks like OpenRouterHandler
-    const mockHandler = {
-        completePrompt: jest.fn(),
-        createMessage: jest.fn(),
-        getModel: jest.fn()
-    }
-
-    // Make instanceof check work
-    Object.setPrototypeOf(mockHandler, OpenRouterHandler.prototype)
-
-    beforeEach(() => {
-        jest.clearAllMocks()
-        ;(buildApiHandler as jest.Mock).mockReturnValue(mockHandler)
-    })
-
-    it('should throw error for non-OpenRouter providers', async () => {
-        const nonOpenRouterConfig: ApiConfiguration = {
-            apiProvider: 'anthropic',
-            apiKey: 'test-key',
-            apiModelId: 'claude-3'
+  beforeEach(() => {
+    jest.clearAllMocks()
+    
+    // Mock the API handler with a completePrompt method
+    ;(buildApiHandler as jest.Mock).mockReturnValue({
+      completePrompt: jest.fn().mockResolvedValue('Enhanced prompt'),
+      createMessage: jest.fn(),
+      getModel: jest.fn().mockReturnValue({ 
+        id: 'test-model',
+        info: {
+          maxTokens: 4096,
+          contextWindow: 8192,
+          supportsPromptCache: false
         }
-        await expect(enhancePrompt(nonOpenRouterConfig, 'test')).rejects.toThrow('Prompt enhancement is only available with OpenRouter')
+      })
+    } as unknown as SingleCompletionHandler)
+  })
+
+  it('enhances prompt using default enhancement prompt when no custom prompt provided', async () => {
+    const result = await enhancePrompt(mockApiConfig, 'Test prompt')
+    
+    expect(result).toBe('Enhanced prompt')
+    const handler = buildApiHandler(mockApiConfig)
+    expect((handler as any).completePrompt).toHaveBeenCalledWith(
+      `${defaultPrompts.enhance}\n\nTest prompt`
+    )
+  })
+
+  it('enhances prompt using custom enhancement prompt when provided', async () => {
+    const customEnhancePrompt = 'You are a custom prompt enhancer'
+    
+    const result = await enhancePrompt(mockApiConfig, 'Test prompt', customEnhancePrompt)
+    
+    expect(result).toBe('Enhanced prompt')
+    const handler = buildApiHandler(mockApiConfig)
+    expect((handler as any).completePrompt).toHaveBeenCalledWith(
+      `${customEnhancePrompt}\n\nTest prompt`
+    )
+  })
+
+  it('throws error for empty prompt input', async () => {
+    await expect(enhancePrompt(mockApiConfig, '')).rejects.toThrow('No prompt text provided')
+  })
+
+  it('throws error for missing API configuration', async () => {
+    await expect(enhancePrompt({} as ApiConfiguration, 'Test prompt')).rejects.toThrow('No valid API configuration provided')
+  })
+
+  it('throws error for API provider that does not support prompt enhancement', async () => {
+    (buildApiHandler as jest.Mock).mockReturnValue({
+      // No completePrompt method
+      createMessage: jest.fn(),
+      getModel: jest.fn().mockReturnValue({ 
+        id: 'test-model',
+        info: {
+          maxTokens: 4096,
+          contextWindow: 8192,
+          supportsPromptCache: false
+        }
+      })
     })
 
-    it('should enhance a valid prompt', async () => {
-        const inputPrompt = 'Write a function to sort an array'
-        const enhancedPrompt = 'Write a TypeScript function that implements an efficient sorting algorithm for a generic array, including error handling and type safety'
-        
-        mockHandler.completePrompt.mockResolvedValue(enhancedPrompt)
+    await expect(enhancePrompt(mockApiConfig, 'Test prompt')).rejects.toThrow('The selected API provider does not support prompt enhancement')
+  })
 
-        const result = await enhancePrompt(mockApiConfig, inputPrompt)
+  it('uses appropriate model based on provider', async () => {
+    const openRouterConfig: ApiConfiguration = {
+      apiProvider: 'openrouter',
+      openRouterApiKey: 'test-key',
+      openRouterModelId: 'test-model'
+    }
 
-        expect(result).toBe(enhancedPrompt)
-        expect(buildApiHandler).toHaveBeenCalledWith(mockApiConfig)
-        expect(mockHandler.completePrompt).toHaveBeenCalledWith(
-            expect.stringContaining(inputPrompt)
-        )
-    })
+    // Mock successful enhancement
+    ;(buildApiHandler as jest.Mock).mockReturnValue({
+      completePrompt: jest.fn().mockResolvedValue('Enhanced prompt'),
+      createMessage: jest.fn(),
+      getModel: jest.fn().mockReturnValue({ 
+        id: 'test-model',
+        info: {
+          maxTokens: 4096,
+          contextWindow: 8192,
+          supportsPromptCache: false
+        }
+      })
+    } as unknown as SingleCompletionHandler)
 
-    it('should throw error when no prompt text is provided', async () => {
-        await expect(enhancePrompt(mockApiConfig, '')).rejects.toThrow('No prompt text provided')
-        expect(mockHandler.completePrompt).not.toHaveBeenCalled()
-    })
+    const result = await enhancePrompt(openRouterConfig, 'Test prompt')
+    
+    expect(buildApiHandler).toHaveBeenCalledWith(openRouterConfig)
+    expect(result).toBe('Enhanced prompt')
+  })
 
-    it('should pass through API errors', async () => {
-        const inputPrompt = 'Test prompt'
-        mockHandler.completePrompt.mockRejectedValue('API error')
+  it('propagates API errors', async () => {
+    (buildApiHandler as jest.Mock).mockReturnValue({
+      completePrompt: jest.fn().mockRejectedValue(new Error('API Error')),
+      createMessage: jest.fn(),
+      getModel: jest.fn().mockReturnValue({ 
+        id: 'test-model',
+        info: {
+          maxTokens: 4096,
+          contextWindow: 8192,
+          supportsPromptCache: false
+        }
+      })
+    } as unknown as SingleCompletionHandler)
 
-        await expect(enhancePrompt(mockApiConfig, inputPrompt)).rejects.toBe('API error')
-    })
-
-    it('should pass the correct prompt format to the API', async () => {
-        const inputPrompt = 'Test prompt'
-        mockHandler.completePrompt.mockResolvedValue('Enhanced test prompt')
-        
-        await enhancePrompt(mockApiConfig, inputPrompt)
-
-        expect(mockHandler.completePrompt).toHaveBeenCalledWith(
-            'Generate an enhanced version of this prompt (reply with only the enhanced prompt - no conversation, explanations, lead-in, bullet points, placeholders, or surrounding quotes):\n\nTest prompt'
-        )
-    })
+    await expect(enhancePrompt(mockApiConfig, 'Test prompt')).rejects.toThrow('API Error')
+  })
 })

src/utils/enhance-prompt.ts

@@ -1,26 +1,27 @@
 import { ApiConfiguration } from "../shared/api"
-import { buildApiHandler } from "../api"
-import { OpenRouterHandler } from "../api/providers/openrouter"
+import { buildApiHandler, SingleCompletionHandler } from "../api"
+import { defaultPrompts } from "../shared/modes"
 
 /**
- * Enhances a prompt using the OpenRouter API without creating a full Cline instance or task history.
+ * Enhances a prompt using the configured API without creating a full Cline instance or task history.
  * This is a lightweight alternative that only uses the API's completion functionality.
  */
-export async function enhancePrompt(apiConfiguration: ApiConfiguration, promptText: string): Promise<string> {
+export async function enhancePrompt(apiConfiguration: ApiConfiguration, promptText: string, enhancePrompt?: string): Promise<string> {
     if (!promptText) {
         throw new Error("No prompt text provided")
     }
-    if (apiConfiguration.apiProvider !== "openrouter") {
-        throw new Error("Prompt enhancement is only available with OpenRouter")
+    if (!apiConfiguration || !apiConfiguration.apiProvider) {
+        throw new Error("No valid API configuration provided")
     }
     
     const handler = buildApiHandler(apiConfiguration)
     
-    // Type guard to check if handler is OpenRouterHandler
-    if (!(handler instanceof OpenRouterHandler)) {
-        throw new Error("Expected OpenRouter handler")
+    // Check if handler supports single completions
+    if (!('completePrompt' in handler)) {
+        throw new Error("The selected API provider does not support prompt enhancement")
     }
     
-    const prompt = `Generate an enhanced version of this prompt (reply with only the enhanced prompt - no conversation, explanations, lead-in, bullet points, placeholders, or surrounding quotes):\n\n${promptText}`
-    return handler.completePrompt(prompt)
+    const enhancePromptText = enhancePrompt ?? defaultPrompts.enhance
+    const prompt = `${enhancePromptText}\n\n${promptText}`
+    return (handler as SingleCompletionHandler).completePrompt(prompt)
 }