import { useCompletion } from 'ai/react';
import { useState, useCallback } from 'react';
interface MorphStreamOptions {
onComplete?: (result: string) => void;
onError?: (error: Error) => void;
defaultModel?: string;
}
export function useMorphStream(options: MorphStreamOptions = {}) {
const [originalCode, setOriginalCode] = useState('');
const [appliedCode, setAppliedCode] = useState<string | null>(null);
const [selectedModel, setSelectedModel] = useState(options.defaultModel || 'morph/morph-v3-large');
const {
completion,
isLoading,
error,
complete,
stop,
} = useCompletion({
api: '/api/morph',
onFinish: (prompt, completion) => {
setAppliedCode(completion);
options.onComplete?.(completion);
},
onError: options.onError,
});
const applyEdit = useCallback(async (code: string, edit: string, model?: string) => {
setOriginalCode(code);
setAppliedCode(null);
await complete('', {
body: {
originalCode: code,
updateSnippet: edit,
model: model || selectedModel, // Allow model override
},
});
}, [complete, selectedModel]);
const acceptChanges = useCallback(() => {
if (appliedCode) {
setOriginalCode(appliedCode);
setAppliedCode(null);
}
}, [appliedCode]);
const rejectChanges = useCallback(() => {
setAppliedCode(null);
}, []);
return {
originalCode,
streamingCode: completion,
appliedCode,
selectedModel,
setSelectedModel,
isLoading,
error,
applyEdit,
acceptChanges,
rejectChanges,
stop,
};
}