export function useInference()

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,
  };
}