40 lines
1 KiB
TypeScript
40 lines
1 KiB
TypeScript
'use client';
|
|
|
|
import React, { createContext, useContext, useMemo, useState } from 'react';
|
|
import type { DataUIPart } from 'ai';
|
|
import type { CustomUIDataTypes } from '@/lib/types';
|
|
|
|
interface DataStreamContextValue {
|
|
dataStream: DataUIPart<CustomUIDataTypes>[];
|
|
setDataStream: React.Dispatch<
|
|
React.SetStateAction<DataUIPart<CustomUIDataTypes>[]>
|
|
>;
|
|
}
|
|
|
|
const DataStreamContext = createContext<DataStreamContextValue | null>(null);
|
|
|
|
export function DataStreamProvider({
|
|
children,
|
|
}: {
|
|
children: React.ReactNode;
|
|
}) {
|
|
const [dataStream, setDataStream] = useState<DataUIPart<CustomUIDataTypes>[]>(
|
|
[],
|
|
);
|
|
|
|
const value = useMemo(() => ({ dataStream, setDataStream }), [dataStream]);
|
|
|
|
return (
|
|
<DataStreamContext.Provider value={value}>
|
|
{children}
|
|
</DataStreamContext.Provider>
|
|
);
|
|
}
|
|
|
|
export function useDataStream() {
|
|
const context = useContext(DataStreamContext);
|
|
if (!context) {
|
|
throw new Error('useDataStream must be used within a DataStreamProvider');
|
|
}
|
|
return context;
|
|
}
|