Skip to content

Commit

Permalink
fix: #3513 - anthropic extension does not forward the system prompt (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
louis-jan authored Sep 24, 2024
1 parent c0b59ec commit 6af17c6
Show file tree
Hide file tree
Showing 5 changed files with 95 additions and 2 deletions.
9 changes: 9 additions & 0 deletions extensions/inference-anthropic-extension/jest.config.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
/** @type {import('ts-jest').JestConfigWithTsJest} */
module.exports = {
preset: 'ts-jest',
testEnvironment: 'node',
transform: {
'node_modules/@janhq/core/.+\\.(j|t)s?$': 'ts-jest',
},
transformIgnorePatterns: ['node_modules/(?!@janhq/core/.*)'],
}
1 change: 1 addition & 0 deletions extensions/inference-anthropic-extension/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
"author": "Jan <service@jan.ai>",
"license": "AGPL-3.0",
"scripts": {
"test": "jest test",
"build": "tsc -b . && webpack --config webpack.config.js",
"build:publish": "rimraf *.tgz --glob && yarn build && npm pack && cpx *.tgz ../../pre-install",
"sync:core": "cd ../.. && yarn build:core && cd extensions && rm yarn.lock && cd inference-anthropic-extension && yarn && yarn build:publish"
Expand Down
77 changes: 77 additions & 0 deletions extensions/inference-anthropic-extension/src/anthropic.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
// Import necessary modules
import JanInferenceAnthropicExtension, { Settings } from '.'
import { PayloadType, ChatCompletionRole } from '@janhq/core'

// Mocks
jest.mock('@janhq/core', () => ({
RemoteOAIEngine: jest.fn().mockImplementation(() => ({
registerSettings: jest.fn(),
registerModels: jest.fn(),
getSetting: jest.fn(),
onChange: jest.fn(),
onSettingUpdate: jest.fn(),
onLoad: jest.fn(),
headers: jest.fn(),
})),
PayloadType: jest.fn(),
ChatCompletionRole: {
User: 'user' as const,
Assistant: 'assistant' as const,
System: 'system' as const,
},
}))

// Helper functions
const createMockPayload = (): PayloadType => ({
messages: [
{ role: ChatCompletionRole.System, content: 'Meow' },
{ role: ChatCompletionRole.User, content: 'Hello' },
{ role: ChatCompletionRole.Assistant, content: 'Hi there' },
],
model: 'claude-v1',
stream: false,
})

describe('JanInferenceAnthropicExtension', () => {
let extension: JanInferenceAnthropicExtension

beforeEach(() => {
extension = new JanInferenceAnthropicExtension('', '')
extension.apiKey = 'mock-api-key'
extension.inferenceUrl = 'mock-endpoint'
jest.clearAllMocks()
})

it('should initialize with correct settings', async () => {
await extension.onLoad()
expect(extension.apiKey).toBe('mock-api-key')
expect(extension.inferenceUrl).toBe('mock-endpoint')
})

it('should transform payload correctly', () => {
const payload = createMockPayload()
const transformedPayload = extension.transformPayload(payload)

expect(transformedPayload).toEqual({
max_tokens: 4096,
model: 'claude-v1',
stream: false,
system: 'Meow',
messages: [
{ role: 'user', content: 'Hello' },
{ role: 'assistant', content: 'Hi there' },
],
})
})

it('should transform response correctly', () => {
const nonStreamResponse = { content: [{ text: 'Test response' }] }
const streamResponse =
'data: {"type":"content_block_delta","delta":{"text":"Hello"}}'

expect(extension.transformResponse(nonStreamResponse)).toBe('Test response')
expect(extension.transformResponse(streamResponse)).toBe('Hello')
expect(extension.transformResponse('')).toBe('')
expect(extension.transformResponse('event: something')).toBe('')
})
})
7 changes: 6 additions & 1 deletion extensions/inference-anthropic-extension/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import { ChatCompletionRole } from '@janhq/core'
declare const SETTINGS: Array<any>
declare const MODELS: Array<any>

enum Settings {
export enum Settings {
apiKey = 'anthropic-api-key',
chatCompletionsEndPoint = 'chat-completions-endpoint',
}
Expand All @@ -23,6 +23,7 @@ type AnthropicPayloadType = {
model?: string
max_tokens?: number
messages?: Array<{ role: string; content: string }>
system?: string
}

/**
Expand Down Expand Up @@ -113,6 +114,10 @@ export default class JanInferenceAnthropicExtension extends RemoteOAIEngine {
role: 'assistant',
content: item.content as string,
})
} else if (item.role === ChatCompletionRole.System) {
// When using Claude, you can dramatically improve its performance by using the system parameter to give it a role.
// This technique, known as role prompting, is the most powerful way to use system prompts with Claude.
convertedData.system = item.content as string
}
})

Expand Down
3 changes: 2 additions & 1 deletion extensions/inference-anthropic-extension/tsconfig.json
Original file line number Diff line number Diff line change
Expand Up @@ -10,5 +10,6 @@
"skipLibCheck": true,
"rootDir": "./src"
},
"include": ["./src"]
"include": ["./src"],
"exclude": ["**/*.test.ts"]
}

0 comments on commit 6af17c6

Please sign in to comment.