in src/lib/inference.tsx [92:153]
export function useInference({ apiKey }) {
const [isLoading, setIsLoading] = useState(false);
const [partialText, setPartialText] = useState("");
const [inferenceResult, setInferenceResult] = useState("");
const [error, setError] = useState<string | null>(null);
const inferenceInternal = async ({
prompt,
model,
maxTokens,
}: {
prompt: string;
model: string;
maxTokens: number;
}) => {
setIsLoading(true);
setPartialText("");
const client = new InferenceClient(apiKey);
try {
const stream = client.chatCompletionStream({
provider: "nebius",
model,
maxTokens,
messages: [
{
role: "user",
content: prompt,
},
],
});
let result = "";
for await (const chunk of stream) {
result += chunk.choices[0].delta.content;
setPartialText(result);
}
setIsLoading(false);
setInferenceResult(result);
return {status: "success", result};
} catch (error) {
console.error("Error in inference", error);
setError(error.message);
setIsLoading(false);
return {status: "error", result: error.message};
}
};
const status = isLoading ? "thinking" : error ? "error" : "done";
return {
status,
partialText,
inferenceResult,
error,
inference: inferenceInternal,
};
}