Skip to content

Commit 0a276b7

Browse files
Copilotjulien-c
andcommitted
Add user agent with package version for hub calls
Co-authored-by: julien-c <[email protected]>
1 parent f44bd6c commit 0a276b7

File tree

5 files changed

+132
-0
lines changed

5 files changed

+132
-0
lines changed

packages/hub/src/consts.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,3 @@
11
export const HUB_URL = "https://huggingface.co";
2+
export const VERSION = "2.1.0";
3+
export const USER_AGENT = `@huggingface/hub/${VERSION}`;

packages/hub/src/index.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
export * from "./lib";
2+
export * from "./utils";
3+
export { USER_AGENT, VERSION } from "./consts";
4+
25
// Typescript 5 will add 'export type *'
36
export type {
47
AccessToken,
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import { describe, it, expect, vi } from "vitest";
2+
import { createFetch } from "./createFetch";
3+
import { USER_AGENT } from "../consts";
4+
5+
describe("createFetch", () => {
6+
it("should add user-agent header to fetch requests", async () => {
7+
// Create a mock fetch function
8+
const mockFetch = vi.fn().mockResolvedValue({
9+
ok: true,
10+
});
11+
12+
// Create a wrapped fetch with our utility
13+
const wrappedFetch = createFetch(mockFetch);
14+
15+
// Call the wrapped fetch
16+
const url = "https://huggingface.co/api/test";
17+
await wrappedFetch(url);
18+
19+
// Check if the mock was called with the correct headers
20+
expect(mockFetch).toHaveBeenCalledWith(url, expect.objectContaining({
21+
headers: expect.objectContaining({
22+
get: expect.any(Function),
23+
has: expect.any(Function),
24+
})
25+
}));
26+
27+
// Get the headers from the mock call
28+
const headers = mockFetch.mock.calls[0][1].headers;
29+
expect(headers.get("user-agent")).toBe(USER_AGENT);
30+
});
31+
32+
it("should not override existing user-agent header", async () => {
33+
// Create a mock fetch function
34+
const mockFetch = vi.fn().mockResolvedValue({
35+
ok: true,
36+
});
37+
38+
// Create a wrapped fetch with our utility
39+
const wrappedFetch = createFetch(mockFetch);
40+
41+
// Call the wrapped fetch with a custom user-agent header
42+
const url = "https://huggingface.co/api/test";
43+
const customUserAgent = "custom-user-agent";
44+
await wrappedFetch(url, {
45+
headers: {
46+
"user-agent": customUserAgent,
47+
},
48+
});
49+
50+
// Get the headers from the mock call
51+
const headers = mockFetch.mock.calls[0][1].headers;
52+
expect(headers.get("user-agent")).toBe(customUserAgent);
53+
});
54+
55+
it("should preserve other headers and request options", async () => {
56+
// Create a mock fetch function
57+
const mockFetch = vi.fn().mockResolvedValue({
58+
ok: true,
59+
});
60+
61+
// Create a wrapped fetch with our utility
62+
const wrappedFetch = createFetch(mockFetch);
63+
64+
// Call the wrapped fetch with additional headers and options
65+
const url = "https://huggingface.co/api/test";
66+
const options = {
67+
method: "POST",
68+
headers: {
69+
"Content-Type": "application/json",
70+
"Authorization": "token-value",
71+
},
72+
body: JSON.stringify({ data: "test" }),
73+
};
74+
75+
await wrappedFetch(url, options);
76+
77+
// Check if the mock was called with all the options preserved
78+
expect(mockFetch).toHaveBeenCalledWith(url, expect.objectContaining({
79+
method: "POST",
80+
body: JSON.stringify({ data: "test" }),
81+
headers: expect.objectContaining({
82+
get: expect.any(Function),
83+
has: expect.any(Function),
84+
})
85+
}));
86+
87+
// Get the headers from the mock call
88+
const headers = mockFetch.mock.calls[0][1].headers;
89+
expect(headers.get("Content-Type")).toBe("application/json");
90+
expect(headers.get("Authorization")).toBe("token-value");
91+
expect(headers.get("user-agent")).toBe(USER_AGENT);
92+
});
93+
});

packages/hub/src/utils/createFetch.ts

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
import { USER_AGENT } from "../consts";
2+
3+
/**
4+
* Creates a fetch wrapper that automatically adds the user-agent header with the package version
5+
*
6+
* @param fetch - The base fetch function to wrap
7+
* @returns A wrapped fetch function that includes the user-agent header
8+
*/
9+
export function createFetch(
10+
baseFetch: typeof fetch = typeof fetch !== "undefined" ? fetch : undefined as unknown as typeof fetch
11+
): typeof fetch {
12+
return function wrappedFetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
13+
const headers = new Headers(init?.headers);
14+
15+
// Only add the user-agent if it's not already set
16+
if (!headers.has("user-agent")) {
17+
headers.set("user-agent", USER_AGENT);
18+
}
19+
20+
const newInit: RequestInit = {
21+
...init,
22+
headers,
23+
};
24+
25+
return baseFetch(input, newInit);
26+
};
27+
}
28+
29+
/**
30+
* Default fetch instance with user-agent header included
31+
*/
32+
export const fetchWithUserAgent = createFetch();

packages/hub/src/utils/index.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
// Export utility functions
2+
export * from "./createFetch";

0 commit comments

Comments
 (0)