|
| 1 | +import * as providers from "./providers.js"; |
| 2 | +import { pairs_to_messages } from "#root/src/classes/message.js"; |
| 3 | +import { truncate_chat } from "#root/src/helpers/helper.js"; |
| 4 | + |
| 5 | +/** |
| 6 | + * Generate a response from a provider. |
| 7 | + * Returns an async generator that yields each chunk of the response. |
| 8 | + * |
| 9 | + * @param {Object} info - Provider info object (contains provider, model, etc.) |
| 10 | + * @param {Array} chat - Array of messages |
| 11 | + * @param {Object} options - Additional options |
| 12 | + * @param {Function} status - Optional function that returns true when generation should stop |
| 13 | + * @returns {AsyncGenerator} - Async generator that yields chunks of the response |
| 14 | + */ |
| 15 | +export const generateResponse = async function* (info, chat, options = {}, status = null) { |
| 16 | + const provider = info.provider; |
| 17 | + |
| 18 | + // Get options from info |
| 19 | + options = providers.get_options_from_info(info, options); |
| 20 | + |
| 21 | + try { |
| 22 | + if (provider.customStream) { |
| 23 | + // For providers that use stream_update callback: Use a promise-based signal. |
| 24 | + let response = ""; |
| 25 | + let queue = []; |
| 26 | + let signalPromiseResolve = null; |
| 27 | + let signalPromise = new Promise((resolve) => { |
| 28 | + signalPromiseResolve = resolve; |
| 29 | + }); |
| 30 | + let done = false; |
| 31 | + let error = null; |
| 32 | + |
| 33 | + // Start the provider's generate method in the background |
| 34 | + provider |
| 35 | + .generate(chat, options, { |
| 36 | + stream_update: (new_response) => { |
| 37 | + if (new_response === null || new_response === undefined) return; |
| 38 | + |
| 39 | + const content = new_response.substring(response.length); |
| 40 | + response = new_response; |
| 41 | + |
| 42 | + if (content) { |
| 43 | + queue.push(content); |
| 44 | + if (signalPromiseResolve) { |
| 45 | + // Signal that new data is available |
| 46 | + const resolveFunc = signalPromiseResolve; |
| 47 | + signalPromiseResolve = null; // Prevent resolving multiple times before next await |
| 48 | + signalPromise = new Promise((resolve) => { |
| 49 | + signalPromiseResolve = resolve; |
| 50 | + }); // Create next signal promise |
| 51 | + resolveFunc(); |
| 52 | + } |
| 53 | + } |
| 54 | + }, |
| 55 | + }) |
| 56 | + .then(() => { |
| 57 | + done = true; |
| 58 | + if (signalPromiseResolve) signalPromiseResolve(); // Signal completion |
| 59 | + }) |
| 60 | + .catch((e) => { |
| 61 | + console.error("Error in provider.generate:", e); |
| 62 | + error = e; |
| 63 | + done = true; |
| 64 | + if (signalPromiseResolve) signalPromiseResolve(); // Signal error/completion |
| 65 | + }); |
| 66 | + |
| 67 | + // Async generator loop consuming the queue, waiting for signals |
| 68 | + while (!done || queue.length > 0) { |
| 69 | + // Check status before yielding or waiting |
| 70 | + if (status && status()) { |
| 71 | + console.log("Generator stopping due to status check."); |
| 72 | + // TODO: Ideally, signal the provider.generate to stop if possible. |
| 73 | + return; |
| 74 | + } |
| 75 | + |
| 76 | + // Yield all currently queued chunks |
| 77 | + while (queue.length > 0) { |
| 78 | + yield queue.shift(); |
| 79 | + } |
| 80 | + |
| 81 | + // If not done and queue is empty, wait for the next signal (data or completion) |
| 82 | + if (!done) { |
| 83 | + await signalPromise; |
| 84 | + } |
| 85 | + } |
| 86 | + |
| 87 | + if (error) { |
| 88 | + console.log(error); |
| 89 | + } |
| 90 | + } else if (info.stream) { |
| 91 | + // For providers using async generators (standard streaming) |
| 92 | + const response = await provider.generate(chat, options, {}); |
| 93 | + |
| 94 | + // Yield chunks, checking status periodically |
| 95 | + let i = 0; |
| 96 | + for await (const chunk of response) { |
| 97 | + if ((i & 15) === 0 && status && status()) { |
| 98 | + return; |
| 99 | + } |
| 100 | + yield chunk; |
| 101 | + i++; |
| 102 | + } |
| 103 | + } else { |
| 104 | + // For non-streaming providers |
| 105 | + const response = await provider.generate(chat, options, {}); |
| 106 | + yield response; |
| 107 | + } |
| 108 | + } catch (e) { |
| 109 | + console.error("Error generating response:", e); |
| 110 | + yield `Error: ${e.message}`; |
| 111 | + } |
| 112 | +}; |
| 113 | + |
| 114 | +/** |
| 115 | + * Generate a complete response as a string. |
| 116 | + * |
| 117 | + * @param {Object} info - Provider info object |
| 118 | + * @param {Array} chat - Array of messages |
| 119 | + * @param {Object} options - Additional options |
| 120 | + * @returns {Promise<string>} - Promise that resolves to the complete response |
| 121 | + */ |
| 122 | +export const generateResponseSync = async (info, chat, options = {}) => { |
| 123 | + let completeResponse = ""; |
| 124 | + |
| 125 | + for await (const chunk of generateResponse(info, chat, options)) { |
| 126 | + completeResponse += chunk; |
| 127 | + } |
| 128 | + |
| 129 | + return completeResponse; |
| 130 | +}; |
| 131 | + |
| 132 | +/** |
| 133 | + * Format chat and generate a response. |
| 134 | + * |
| 135 | + * @param {Object} currentChat - The chat object |
| 136 | + * @param {string|null} query - Optional query to append |
| 137 | + * @param {Object} options - Additional options |
| 138 | + * @param {Function} status - Optional function that returns true when generation should stop |
| 139 | + * @returns {AsyncGenerator} - Async generator that yields chunks of the response |
| 140 | + */ |
| 141 | +export const generateChatResponse = async function* (currentChat, query = null, options = {}, status = null) { |
| 142 | + // Load provider and model |
| 143 | + const info = providers.get_provider_info(currentChat.provider); |
| 144 | + |
| 145 | + // Format chat |
| 146 | + let chat = pairs_to_messages(currentChat.messages, query); |
| 147 | + chat = truncate_chat(chat, info); |
| 148 | + |
| 149 | + // Merge options |
| 150 | + const mergedOptions = { |
| 151 | + ...providers.get_options_from_info(info, currentChat.options), |
| 152 | + ...options, |
| 153 | + }; |
| 154 | + |
| 155 | + // Generate response |
| 156 | + yield* generateResponse(info, chat, mergedOptions, status); |
| 157 | +}; |
| 158 | + |
| 159 | +/** |
| 160 | + * Get a complete chat response as a string |
| 161 | + * |
| 162 | + * @param {Object} currentChat - The chat object |
| 163 | + * @param {string|null} query - Optional query to append |
| 164 | + * @param {Object} options - Additional options |
| 165 | + * @returns {Promise<string>} - Promise that resolves to the complete response |
| 166 | + */ |
| 167 | +export const generateChatResponseSync = async (currentChat, query = null, options = {}) => { |
| 168 | + let completeResponse = ""; |
| 169 | + |
| 170 | + for await (const chunk of generateChatResponse(currentChat, query, options)) { |
| 171 | + completeResponse += chunk; |
| 172 | + } |
| 173 | + |
| 174 | + return completeResponse; |
| 175 | +}; |
0 commit comments