Skip to content

Commit

Permalink
Merge pull request #1809 from SciPhi-AI/feature/add-auth-callback
Browse files Browse the repository at this point in the history 
Feature/add auth callback
  • Loading branch information
@emrgnt-cmplxty
emrgnt-cmplxty authored Jan 13, 2025
2 parents 9ec211f + 7f79424 commit aabd114
Show file tree
Hide file tree
Showing 6 changed files with 197 additions and 47 deletions.
2 changes: 1 addition & 1 deletion js/sdk/package.json
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "r2r-js",
"version": "0.4.14",
"version": "0.4.15",
"description": "",
"main": "dist/index.js",
"browser": "dist/index.browser.js",
Expand Down
27 changes: 21 additions & 6 deletions js/sdk/src/baseClient.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,12 +42,23 @@ export abstract class BaseClient {
protected refreshToken: string | null;
protected anonymousTelemetry: boolean;

constructor(baseURL: string, prefix: string = "", anonymousTelemetry = true) {
// NEW: declare enableAutoRefresh
protected enableAutoRefresh: boolean;

constructor(
baseURL: string,
prefix: string = "",
anonymousTelemetry = true,
enableAutoRefresh = false
) {
this.baseUrl = `${baseURL}${prefix}`;
this.accessToken = null;
this.refreshToken = null;
this.anonymousTelemetry = anonymousTelemetry;

// Add this assignment
this.enableAutoRefresh = enableAutoRefresh;

this.axiosInstance = axios.create({
baseURL: this.baseUrl,
headers: {
Expand All @@ -60,7 +71,7 @@ export abstract class BaseClient {
method: Method,
endpoint: string,
options: any = {},
version: "v3" = "v3",
version: "v3" = "v3"
): Promise<T> {
const url = `/${version}/${endpoint}`;
const config: AxiosRequestConfig = {
Expand All @@ -81,11 +92,13 @@ export abstract class BaseClient {
if (Array.isArray(value)) {
return value
.map(
(v) => `${encodeURIComponent(key)}=${encodeURIComponent(v)}`,
(v) => `${encodeURIComponent(key)}=${encodeURIComponent(v)}`
)
.join("&");
}
return `${encodeURIComponent(key)}=${encodeURIComponent(String(value))}`;
return `${encodeURIComponent(key)}=${encodeURIComponent(
String(value)
)}`;
})
.join("&");
};
Expand All @@ -102,7 +115,9 @@ export abstract class BaseClient {
config.data = Object.keys(options.data)
.map(
(key) =>
`${encodeURIComponent(key)}=${encodeURIComponent(options.data[key])}`,
`${encodeURIComponent(key)}=${encodeURIComponent(
options.data[key]
)}`
)
.join("&");
} else {
Expand Down Expand Up @@ -144,7 +159,7 @@ export abstract class BaseClient {
throw new Error(
`HTTP error! status: ${response.status}: ${
ensureCamelCase(errorData).message || "Unknown error"
}`,
}`
);
}

Expand Down
161 changes: 123 additions & 38 deletions js/sdk/src/r2rClient.ts
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import axios, { Method } from "axios";

import axios, { Method, AxiosError } from "axios";
import { BaseClient } from "./baseClient";

import { ChunksClient } from "./v3/clients/chunks";
Expand All @@ -13,15 +12,25 @@ import { RetrievalClient } from "./v3/clients/retrieval";
import { SystemClient } from "./v3/clients/system";
import { UsersClient } from "./v3/clients/users";

let fs: any;
if (typeof window === "undefined") {
import("fs").then((module) => {
fs = module;
});
}

import { initializeTelemetry } from "./feature";

type RefreshTokenResponse = {
results: {
accessToken: string;
refreshToken: string;
};
};

interface R2RClientOptions {
enableAutoRefresh?: boolean;
getTokensCallback?: () => {
accessToken: string | null;
refreshToken: string | null;
};
setTokensCallback?: (accessToken: string | null, refreshToken: string | null) => void;
onRefreshFailedCallback?: () => void;
}

export class r2rClient extends BaseClient {
public readonly chunks: ChunksClient;
public readonly collections: CollectionsClient;
Expand All @@ -34,8 +43,18 @@ export class r2rClient extends BaseClient {
public readonly system: SystemClient;
public readonly users: UsersClient;

constructor(baseURL: string, anonymousTelemetry = true) {
super(baseURL, "", anonymousTelemetry);
private getTokensCallback?: R2RClientOptions["getTokensCallback"];
private setTokensCallback?: R2RClientOptions["setTokensCallback"];
private onRefreshFailedCallback?: R2RClientOptions["onRefreshFailedCallback"];

constructor(
baseURL: string,
anonymousTelemetry = true,
options: R2RClientOptions = {}
) {
super(baseURL, "", anonymousTelemetry, options.enableAutoRefresh);

console.log("[r2rClient] Creating new client with baseURL =", baseURL);

this.chunks = new ChunksClient(this);
this.collections = new CollectionsClient(this);
Expand All @@ -55,48 +74,114 @@ export class r2rClient extends BaseClient {
headers: {
"Content-Type": "application/json",
},
paramsSerializer: (params) => {
const parts: string[] = [];
Object.entries(params).forEach(([key, value]) => {
if (Array.isArray(value)) {
value.forEach((v) =>
parts.push(`${encodeURIComponent(key)}=${encodeURIComponent(v)}`),
);
} else {
parts.push(
`${encodeURIComponent(key)}=${encodeURIComponent(value)}`,
);
}
});
return parts.join("&");
});

this.getTokensCallback = options.getTokensCallback;
this.setTokensCallback = options.setTokensCallback;
this.onRefreshFailedCallback = options.onRefreshFailedCallback;

// 1) Request interceptor: attach current access token (if any)
this.axiosInstance.interceptors.request.use(
(config) => {
const tokenData = this.getTokensCallback?.();
const accessToken = tokenData?.accessToken || null;
if (accessToken) {
console.log(`[r2rClient] Attaching access token to request: ${accessToken.slice(0, 15)}...`);
config.headers["Authorization"] = `Bearer ${accessToken}`;
} else {
console.log("[r2rClient] No access token found, sending request without Authorization header");
}
return config;
},
transformRequest: [
(data) => {
if (typeof data === "string") {
return data;
(error) => {
console.error("[r2rClient] Request interceptor error:", error);
return Promise.reject(error);
}
);

// 2) Response interceptor: see if we got 401/403 => attempt to refresh
this.setupResponseInterceptor();
}

private setupResponseInterceptor() {
this.axiosInstance.interceptors.response.use(
(response) => response,
async (error: AxiosError) => {
// Some logs to see what's going on
console.warn("[r2rClient] Response interceptor caught an error:", error);
const status = error.response?.status;
const failingUrl = error.config?.url;
console.warn("[r2rClient] Failing request URL:", failingUrl, "status =", status);
console.warn("failingUrl?.includes('/v3/users/refresh-token') = ", failingUrl?.includes("/v3/users/refresh-token"))

// 1) If the refresh endpoint itself fails => don't try again
if (failingUrl?.includes("/v3/users/refresh-token")) {
console.error("[r2rClient] Refresh call itself returned 401/403 => logging out");
this.onRefreshFailedCallback?.();
return Promise.reject(error);
}

// 2) If normal request => attempt refresh if 401/403
if ((status === 401 || status === 403) && this.getTokensCallback) {
const { refreshToken } = this.getTokensCallback();
if (!refreshToken) {
console.error("[r2rClient] No refresh token found => logout");
this.onRefreshFailedCallback?.();
return Promise.reject(error);
}
return JSON.stringify(data);
},
],
});

// Attempt refresh
try {
console.log("[r2rClient] Attempting token refresh...");
const refreshResponse = (await this.users.refreshAccessToken()) as RefreshTokenResponse;
const newAccessToken = refreshResponse.results.accessToken;
const newRefreshToken = refreshResponse.results.refreshToken;
console.log("[r2rClient] Refresh call succeeded; new access token:", newAccessToken.slice(0, 15), "...");

// set new tokens
this.setTokens(newAccessToken, newRefreshToken);

// Re-try the original request
if (error.config) {
error.config.headers["Authorization"] = `Bearer ${newAccessToken}`;
console.log("[r2rClient] Retrying original request with new access token...");
return this.axiosInstance.request(error.config);
} else {
console.warn("[r2rClient] No request config found to retry. Possibly manual re-fetch needed");
}
} catch (refreshError) {
console.error("[r2rClient] Refresh attempt failed => logging out. Error was:", refreshError);
this.onRefreshFailedCallback?.();
return Promise.reject(refreshError);
}
}

// 3) If not a 401/403, or no refresh logic => just reject
console.log("[r2rClient] Non-401/403 error => rejecting request");
return Promise.reject(error);
}
);
}

public makeRequest<T = any>(
method: Method,
endpoint: string,
options: any = {},
options: any = {}
): Promise<T> {
console.log(`[r2rClient] makeRequest: ${method.toUpperCase()} ${endpoint}`);
return this._makeRequest(method, endpoint, options, "v3");
}

public getRefreshToken(): string | null {
return this.refreshToken;
}

setTokens(accessToken: string | null, refreshToken: string | null): void {
this.accessToken = accessToken;
this.refreshToken = refreshToken;
public setTokens(accessToken: string | null, refreshToken: string | null): void {
// Optional: log the changes, but be careful not to log full tokens in prod
console.log("[r2rClient] Setting tokens. Access token:", accessToken?.slice(0, 15), "... refresh token:", refreshToken?.slice(0, 15), "...");
super.setTokens(accessToken || "", refreshToken || "");
this.setTokensCallback?.(accessToken, refreshToken);
}
}

export default r2rClient;
export default r2rClient;
7 changes: 7 additions & 0 deletions py/core/configs/r2r_azure.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,10 @@ vision_pdf_model = "azure/gpt-4o"

[ingestion.chunk_enrichment_settings]
generation_config = { model = "azure/gpt-4o-mini" }


[auth]
provider = "r2r"
access_token_lifetime_in_minutes = 1
refresh_token_lifetime_in_days = 3
require_authentication = false
2 changes: 1 addition & 1 deletion py/r2r/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
__all__ += core.__all__
except ImportError as e:
logger.warning(
f"Warning: encountered ImportError: `{e}`, likely due to core dependencies not being installed. This will not affect you use of SDK, but use of `serve` method will not be available."
f"Warning: encountered ImportError: `{e}`, likely due to core dependencies not being installed. This will not affect your use of SDK, but use of `r2r serve` may not be available."
)


Expand Down
45 changes: 44 additions & 1 deletion py/r2r/r2r.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,49 @@ default_max_documents_per_user = 10_000
default_max_chunks_per_user = 10_000_000
default_max_collections_per_user = 5_000

# Set the default max upload size to 2 GB for local testing
default_max_upload_size = 2147483648 # 2 GB for anything not explicitly listed

[app.max_upload_size_by_type]
# Common text-based formats
txt = 2147483648 # 2 GB
md = 2147483648
tsv = 2147483648
csv = 2147483648
xml = 2147483648
html = 2147483648

# Office docs
doc = 2147483648
docx = 2147483648
ppt = 2147483648
pptx = 2147483648
xls = 2147483648
xlsx = 2147483648
odt = 2147483648

# PDFs
pdf = 2147483648

# E-mail
eml = 2147483648
msg = 2147483648
p7s = 2147483648

# Images
bmp = 2147483648
heic = 2147483648
jpeg = 2147483648
jpg = 2147483648
png = 2147483648
tiff = 2147483648

# E-books and other formats
epub = 2147483648
rtf = 2147483648
rst = 2147483648
org = 2147483648

[agent]
system_instruction_name = "rag_agent"
# tool_names = ["local_search", "web_search"] # uncomment to enable web search
Expand All @@ -15,7 +58,7 @@ tool_names = ["local_search"]

[auth]
provider = "r2r"
access_token_lifetime_in_minutes = 60
access_token_lifetime_in_minutes = 60000 # set a very high default value, for easier testing
refresh_token_lifetime_in_days = 7
require_authentication = false
require_email_verification = false
Expand Down

0 comments on commit aabd114

Please sign in to comment.