diff --git a/src/renderer/src/hooks/useQuickCompletion.ts b/src/renderer/src/hooks/useQuickCompletion.ts new file mode 100644 index 0000000000..1b049162d6 --- /dev/null +++ b/src/renderer/src/hooks/useQuickCompletion.ts @@ -0,0 +1,48 @@ +import { fetchChatCompletion } from '@renderer/services/ApiService' +import { getDefaultAssistant } from '@renderer/services/AssistantService' +import type { Assistant, FetchChatCompletionParams } from '@renderer/types' +import type { Chunk } from '@renderer/types/chunk' +import { useCallback } from 'react' + +import { useDefaultModel } from './useAssistant' + +/** + * Parameters for performing a quick completion using the quick model. + */ +export type QuickCompletionParams = { + /** + * The prompt text to send to the model. + */ + prompt: string + /** + * Callback invoked whenever a new chunk of the streaming response arrives. + */ + onChunk: (chunk: Chunk) => void + /** + * Optional partial assistant settings to override the default quick assistant. + */ + assistantUpdate?: Partial + /** + * Optional additional parameters to pass to the underlying fetchChatCompletion call. + * Excludes prompt, messages, assistant, and onChunkReceived which are handled internally. + */ + params?: Partial> +} + +export const useQuickCompletion = () => { + const { quickModel } = useDefaultModel() + + const completion = useCallback( + async ({ prompt, onChunk, assistantUpdate, params }: QuickCompletionParams) => { + const assistant = { + ...getDefaultAssistant(), + model: quickModel, + ...assistantUpdate + } satisfies Assistant + return fetchChatCompletion({ prompt, assistant, onChunkReceived: onChunk, ...params }) + }, + [quickModel] + ) + + return completion +}