Skip to content

[WIP] Add user agent with package version for hub calls #1490

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions packages/hub/src/consts.ts
Original file line number Diff line number Diff line change
@@ -1 +1,3 @@
export const HUB_URL = "https://huggingface.co";
export const VERSION = "2.1.0";
export const USER_AGENT = `@huggingface/hub/${VERSION}`;
3 changes: 3 additions & 0 deletions packages/hub/src/index.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
93 changes: 93 additions & 0 deletions packages/hub/src/utils/createFetch.spec.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});
32 changes: 32 additions & 0 deletions packages/hub/src/utils/createFetch.ts
Original file line number Diff line number Diff line change
@@ -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<Response> {
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();
2 changes: 2 additions & 0 deletions packages/hub/src/utils/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
// Export utility functions
export * from "./createFetch";