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" } }], + }); + }); +});