From 57cc2cdec22d43326e0719bf3efe6369dadc25b1 Mon Sep 17 00:00:00 2001 From: Ayoubhm07 Date: Tue, 6 Oct 2026 00:20:13 +0100 Subject: [PATCH] fix(huggingface): pass stop sequences to the Inference API provider HuggingFaceInferenceAPI._convertArgs forwarded max_new_tokens, temperature, top_k and top_p but dropped options.stop. Its sibling HuggingFaceTGI speaks the same TGI protocol, posts the same { inputs, parameters } payload and does forward stop, capped by maxStopWords. Autocomplete templates always set stop, and stopAtStopTokens only trims the visible text inside the autocomplete pipeline, so chat and edit lost the sequences entirely while autocomplete kept generating up to max_new_tokens instead of stopping. Apply the same capped forwarding the sibling already uses. Add a provider test in the shape the repository already uses for non-OpenAI providers. --- core/llm/llms/HuggingFaceInferenceAPI.ts | 5 + .../llms/HuggingFaceInferenceAPI.vitest.ts | 230 ++++++++++++++++++ 2 files changed, 235 insertions(+) create mode 100644 core/llm/llms/HuggingFaceInferenceAPI.vitest.ts diff --git a/core/llm/llms/HuggingFaceInferenceAPI.ts b/core/llm/llms/HuggingFaceInferenceAPI.ts index 208f6fd495d..3d0fb801c20 100644 --- a/core/llm/llms/HuggingFaceInferenceAPI.ts +++ b/core/llm/llms/HuggingFaceInferenceAPI.ts @@ -3,6 +3,7 @@ import { CompletionOptions } from "../../index.js"; import { BaseLLM } from "../index.js"; class HuggingFaceInferenceAPI extends BaseLLM { + private static MAX_STOP_TOKENS = 4; static providerName = "huggingface-inference-api"; private _convertArgs(options: CompletionOptions) { @@ -11,6 +12,10 @@ class HuggingFaceInferenceAPI extends BaseLLM { temperature: options.temperature, top_k: options.topK, top_p: options.topP, + stop: options.stop?.slice( + 0, + this.maxStopWords ?? HuggingFaceInferenceAPI.MAX_STOP_TOKENS, + ), }; } diff --git a/core/llm/llms/HuggingFaceInferenceAPI.vitest.ts b/core/llm/llms/HuggingFaceInferenceAPI.vitest.ts new file mode 100644 index 00000000000..cfe99c2de1e --- /dev/null +++ b/core/llm/llms/HuggingFaceInferenceAPI.vitest.ts @@ -0,0 +1,230 @@ +import { afterEach, describe, expect, test, vi } from "vitest"; +import { ILLM } from "../../index.js"; +import HuggingFaceInferenceAPI from "./HuggingFaceInferenceAPI.js"; + +interface LlmTestCase { + llm: ILLM; + methodToTest: keyof ILLM; + params: any[]; + expectedRequest: { + url: string; + method: string; + headers?: Record; + body?: Record; + }; + mockResponse?: any; + mockStream?: any[]; +} + +function createMockStream(mockStream: any[]) { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + for (const chunk of mockStream) { + controller.enqueue( + encoder.encode( + `data: ${ + typeof chunk === "string" ? chunk : JSON.stringify(chunk) + }\n\n`, + ), + ); + } + controller.close(); + }, + }); +} + +function setupMockFetch(mockResponse?: any, mockStream?: any[]) { + const mockFetch = vi.fn(); + + if (mockStream) { + const stream = createMockStream(mockStream); + mockFetch.mockResolvedValue( + new Response(stream, { + headers: { + "Content-Type": "text/event-stream", + }, + }), + ); + } else { + mockFetch.mockResolvedValue( + new Response(JSON.stringify(mockResponse), { + headers: { "Content-Type": "application/json" }, + }), + ); + } + + return mockFetch; +} + +function setupReadableStreamPolyfill() { + // This can be removed if https://github.com/nodejs/undici/issues/2888 is resolved + // @ts-ignore + const originalFrom = ReadableStream.from; + // @ts-ignore + ReadableStream.from = (body) => { + if (body?.source) { + return body; + } + return originalFrom(body); + }; +} + +async function executeLlmMethod( + llm: ILLM, + methodToTest: keyof ILLM, + params: any[], +) { + if (typeof (llm as any)[methodToTest] !== "function") { + throw new Error( + `Method ${String(methodToTest)} does not exist on the LLM instance.`, + ); + } + + const result = await (llm as any)[methodToTest](...params); + if (result?.next) { + for await (const _ of result) { + } + } +} + +function assertFetchCall(mockFetch: any, expectedRequest: any) { + expect(mockFetch).toHaveBeenCalledTimes(1); + const [url, options] = mockFetch.mock.calls[0]; + + expect(url.toString()).toBe(expectedRequest.url); + expect(options.method).toBe(expectedRequest.method); + + if (expectedRequest.headers) { + expect(options.headers).toEqual( + expect.objectContaining(expectedRequest.headers), + ); + } + + if (expectedRequest.body) { + const actualBody = JSON.parse(options.body as string); + expect(actualBody).toEqual(expectedRequest.body); + } +} + +async function runLlmTest(testCase: LlmTestCase) { + const { + llm, + methodToTest, + params, + expectedRequest, + mockResponse, + mockStream, + } = testCase; + + const mockFetch = setupMockFetch(mockResponse, mockStream); + setupReadableStreamPolyfill(); + + (llm as any).fetch = mockFetch; + + await executeLlmMethod(llm, methodToTest, params); + assertFetchCall(mockFetch, expectedRequest); +} + +describe("HuggingFaceInferenceAPI", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + test("streamComplete should forward stop sequences", async () => { + const hf = new HuggingFaceInferenceAPI({ + apiKey: "test-api-key", + model: "bigcode/starcoder2-3b", + apiBase: "https://example.endpoints.huggingface.cloud", + }); + + await runLlmTest({ + llm: hf, + methodToTest: "streamComplete", + params: [ + "def add(a, b): return ", + new AbortController().signal, + { raw: true, maxTokens: 1024, stop: ["<|endoftext|>", ""] }, + ], + expectedRequest: { + url: "https://example.endpoints.huggingface.cloud/", + method: "POST", + headers: { + Authorization: "Bearer test-api-key", + "Content-Type": "application/json", + }, + body: { + inputs: "def add(a, b): return ", + stream: true, + parameters: { + max_new_tokens: 1024, + stop: ["<|endoftext|>", ""], + }, + }, + }, + mockStream: [{ token: { text: "a + b" } }], + }); + }); + + test("streamComplete should cap stop sequences at maxStopWords", async () => { + const hf = new HuggingFaceInferenceAPI({ + apiKey: "test-api-key", + model: "bigcode/starcoder2-3b", + apiBase: "https://example.endpoints.huggingface.cloud", + maxStopWords: 2, + }); + + await runLlmTest({ + llm: hf, + methodToTest: "streamComplete", + params: [ + "hello", + new AbortController().signal, + { raw: true, maxTokens: 1024, stop: ["a", "b", "c", "d"] }, + ], + expectedRequest: { + url: "https://example.endpoints.huggingface.cloud/", + method: "POST", + body: { + inputs: "hello", + stream: true, + parameters: { + max_new_tokens: 1024, + stop: ["a", "b"], + }, + }, + }, + mockStream: [{ token: { text: "ok" } }], + }); + }); + + test("streamComplete should omit stop when none are given", async () => { + const hf = new HuggingFaceInferenceAPI({ + apiKey: "test-api-key", + model: "bigcode/starcoder2-3b", + apiBase: "https://example.endpoints.huggingface.cloud", + }); + + await runLlmTest({ + llm: hf, + methodToTest: "streamComplete", + params: [ + "hello", + new AbortController().signal, + { raw: true, maxTokens: 1024 }, + ], + expectedRequest: { + url: "https://example.endpoints.huggingface.cloud/", + method: "POST", + body: { + inputs: "hello", + stream: true, + parameters: { + max_new_tokens: 1024, + }, + }, + }, + mockStream: [{ token: { text: "ok" } }], + }); + }); +});