|
| 1 | +import { getUserConfig } from '../../config/index.mjs' |
| 2 | +import { getChatSystemPromptBase, pushRecord, setAbortController } from './shared.mjs' |
| 3 | +import { getConversationPairs } from '../../utils/get-conversation-pairs.mjs' |
| 4 | +import { fetchSSE } from '../../utils/fetch-sse.mjs' |
| 5 | +import { isEmpty } from 'lodash-es' |
| 6 | + |
| 7 | +/** |
| 8 | + * @param {Runtime.Port} port |
| 9 | + * @param {string} question |
| 10 | + * @param {Session} session |
| 11 | + */ |
| 12 | +export async function generateAnswersWithClaudeApi(port, question, session) { |
| 13 | + const { controller, messageListener, disconnectListener } = setAbortController(port) |
| 14 | + const config = await getUserConfig() |
| 15 | + |
| 16 | + const prompt = getConversationPairs( |
| 17 | + session.conversationRecords.slice(-config.maxConversationContextLength), |
| 18 | + false, |
| 19 | + ) |
| 20 | + prompt.unshift({ role: 'Assistant', content: await getChatSystemPromptBase() }) |
| 21 | + prompt.push({ role: 'Human', content: question }) |
| 22 | + |
| 23 | + let answer = '' |
| 24 | + await fetchSSE( |
| 25 | + `https://api.anthropic.com/v1/complete`, |
| 26 | + { |
| 27 | + method: 'POST', |
| 28 | + signal: controller.signal, |
| 29 | + headers: { |
| 30 | + 'Content-Type': 'application/json', |
| 31 | + 'accept': 'application/json', |
| 32 | + 'anthropic-version': '2023-06-01', |
| 33 | + 'x-api-key': config.claudeApiKey, |
| 34 | + }, |
| 35 | + body: JSON.stringify({ |
| 36 | + model: "claude-2", |
| 37 | + prompt: "\n\nHuman: " + question + "\n\nAssistant:", |
| 38 | + stream: true, |
| 39 | + max_tokens_to_sample: config.maxResponseTokenLength, |
| 40 | + temperature: config.temperature, |
| 41 | + }), |
| 42 | + onMessage(message) { |
| 43 | + console.debug('sse message', message); |
| 44 | + |
| 45 | + let data; |
| 46 | + try { |
| 47 | + data = JSON.parse(message); |
| 48 | + } catch (error) { |
| 49 | + console.debug('json error', error); |
| 50 | + return; |
| 51 | + } |
| 52 | + |
| 53 | + // The Claude v2 API may send metadata fields, handle them here |
| 54 | + if (data.conversationId) session.conversationId = data.conversationId; |
| 55 | + if (data.parentMessageId) session.parentMessageId = data.parentMessageId; |
| 56 | + |
| 57 | + // In Claude's case, the "completion" key holds the text |
| 58 | + if (data.completion) { |
| 59 | + answer += data.completion; |
| 60 | + port.postMessage({ answer: answer, done: false, session: null }); |
| 61 | + } |
| 62 | + |
| 63 | + // Check if the message indicates that Claude is done |
| 64 | + if (data.stop_reason === 'stop_sequence') { |
| 65 | + pushRecord(session, question, answer); |
| 66 | + console.debug('conversation history', { content: session.conversationRecords }); |
| 67 | + port.postMessage({ answer: null, done: true, session: session }); |
| 68 | + } |
| 69 | + }, |
| 70 | + async onStart() {}, |
| 71 | + async onEnd() { |
| 72 | + port.postMessage({ done: true }) |
| 73 | + port.onMessage.removeListener(messageListener) |
| 74 | + port.onDisconnect.removeListener(disconnectListener) |
| 75 | + }, |
| 76 | + async onError(resp) { |
| 77 | + port.onMessage.removeListener(messageListener) |
| 78 | + port.onDisconnect.removeListener(disconnectListener) |
| 79 | + if (resp instanceof Error) throw resp |
| 80 | + const error = await resp.json().catch(() => ({})) |
| 81 | + throw new Error( |
| 82 | + !isEmpty(error) ? JSON.stringify(error) : `${resp.status} ${resp.statusText}`, |
| 83 | + ) |
| 84 | + }, |
| 85 | + }, |
| 86 | + ) |
| 87 | +} |
0 commit comments