diff --git a/packages/hub/src/consts.ts b/packages/hub/src/consts.ts index 5d34e9caec..bbbd1f1f01 100644 --- a/packages/hub/src/consts.ts +++ b/packages/hub/src/consts.ts @@ -1 +1,3 @@ export const HUB_URL = "https://huggingface.co"; +export const VERSION = "2.1.0"; +export const USER_AGENT = `@huggingface/hub/${VERSION}`; diff --git a/packages/hub/src/index.ts b/packages/hub/src/index.ts index a73655c797..9183b68c53 100644 --- a/packages/hub/src/index.ts +++ b/packages/hub/src/index.ts @@ -1,4 +1,7 @@ export * from "./lib"; +export * from "./utils"; +export { USER_AGENT, VERSION } from "./consts"; + // Typescript 5 will add 'export type *' export type { AccessToken, diff --git a/packages/hub/src/utils/createFetch.spec.ts b/packages/hub/src/utils/createFetch.spec.ts new file mode 100644 index 0000000000..0e51fb7f9e --- /dev/null +++ b/packages/hub/src/utils/createFetch.spec.ts @@ -0,0 +1,93 @@ +import { describe, it, expect, vi } from "vitest"; +import { createFetch } from "./createFetch"; +import { USER_AGENT } from "../consts"; + +describe("createFetch", () => { + it("should add user-agent header to fetch requests", async () => { + // Create a mock fetch function + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + }); + + // Create a wrapped fetch with our utility + const wrappedFetch = createFetch(mockFetch); + + // Call the wrapped fetch + const url = "https://huggingface.co/api/test"; + await wrappedFetch(url); + + // Check if the mock was called with the correct headers + expect(mockFetch).toHaveBeenCalledWith(url, expect.objectContaining({ + headers: expect.objectContaining({ + get: expect.any(Function), + has: expect.any(Function), + }) + })); + + // Get the headers from the mock call + const headers = mockFetch.mock.calls[0][1].headers; + expect(headers.get("user-agent")).toBe(USER_AGENT); + }); + + it("should not override existing user-agent header", async () => { + // Create a mock fetch function + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + }); + + // Create a wrapped fetch with our utility + const wrappedFetch = createFetch(mockFetch); + + // Call the wrapped fetch with a custom user-agent header + const url = "https://huggingface.co/api/test"; + const customUserAgent = "custom-user-agent"; + await wrappedFetch(url, { + headers: { + "user-agent": customUserAgent, + }, + }); + + // Get the headers from the mock call + const headers = mockFetch.mock.calls[0][1].headers; + expect(headers.get("user-agent")).toBe(customUserAgent); + }); + + it("should preserve other headers and request options", async () => { + // Create a mock fetch function + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + }); + + // Create a wrapped fetch with our utility + const wrappedFetch = createFetch(mockFetch); + + // Call the wrapped fetch with additional headers and options + const url = "https://huggingface.co/api/test"; + const options = { + method: "POST", + headers: { + "Content-Type": "application/json", + "Authorization": "token-value", + }, + body: JSON.stringify({ data: "test" }), + }; + + await wrappedFetch(url, options); + + // Check if the mock was called with all the options preserved + expect(mockFetch).toHaveBeenCalledWith(url, expect.objectContaining({ + method: "POST", + body: JSON.stringify({ data: "test" }), + headers: expect.objectContaining({ + get: expect.any(Function), + has: expect.any(Function), + }) + })); + + // Get the headers from the mock call + const headers = mockFetch.mock.calls[0][1].headers; + expect(headers.get("Content-Type")).toBe("application/json"); + expect(headers.get("Authorization")).toBe("token-value"); + expect(headers.get("user-agent")).toBe(USER_AGENT); + }); +}); \ No newline at end of file diff --git a/packages/hub/src/utils/createFetch.ts b/packages/hub/src/utils/createFetch.ts new file mode 100644 index 0000000000..42ac5ca41e --- /dev/null +++ b/packages/hub/src/utils/createFetch.ts @@ -0,0 +1,32 @@ +import { USER_AGENT } from "../consts"; + +/** + * Creates a fetch wrapper that automatically adds the user-agent header with the package version + * + * @param fetch - The base fetch function to wrap + * @returns A wrapped fetch function that includes the user-agent header + */ +export function createFetch( + baseFetch: typeof fetch = typeof fetch !== "undefined" ? fetch : undefined as unknown as typeof fetch +): typeof fetch { + return function wrappedFetch(input: RequestInfo | URL, init?: RequestInit): Promise { + const headers = new Headers(init?.headers); + + // Only add the user-agent if it's not already set + if (!headers.has("user-agent")) { + headers.set("user-agent", USER_AGENT); + } + + const newInit: RequestInit = { + ...init, + headers, + }; + + return baseFetch(input, newInit); + }; +} + +/** + * Default fetch instance with user-agent header included + */ +export const fetchWithUserAgent = createFetch(); \ No newline at end of file diff --git a/packages/hub/src/utils/index.ts b/packages/hub/src/utils/index.ts new file mode 100644 index 0000000000..73fefc9038 --- /dev/null +++ b/packages/hub/src/utils/index.ts @@ -0,0 +1,2 @@ +// Export utility functions +export * from "./createFetch"; \ No newline at end of file