453 lines
13 KiB
TypeScript
453 lines
13 KiB
TypeScript
import { Injectable, Logger } from '@nestjs/common';
|
|
import axios, { AxiosProxyConfig, AxiosRequestConfig, AxiosResponse } from 'axios';
|
|
import { isNil } from 'lodash-es';
|
|
import { GeminiRequest, GeminiResponse } from '../interfaces/request-interfaces';
|
|
import { GeminiInternalRequest } from '../../../../lib/antigravity/types';
|
|
import { getServerConfig } from '../../../server-config';
|
|
import { resolveRequestUserAgent } from '../request-user-agent';
|
|
import { UpstreamRequestError } from './upstream-error';
|
|
|
|
@Injectable()
|
|
export class GeminiClient {
|
|
private readonly logger = new Logger(GeminiClient.name);
|
|
// Default to v1beta for most features
|
|
private readonly baseUrl = 'https://generativelanguage.googleapis.com/v1beta';
|
|
private readonly defaultInternalBaseUrls = [
|
|
'https://cloudcode-pa.googleapis.com/v1internal',
|
|
'https://daily-cloudcode-pa.googleapis.com/v1internal',
|
|
];
|
|
|
|
async streamGenerate(
|
|
model: string,
|
|
content: GeminiRequest,
|
|
accessToken: string,
|
|
upstreamProxyUrl?: string,
|
|
): Promise<NodeJS.ReadableStream> {
|
|
const url = `${this.baseUrl}/models/${model}:streamGenerateContent?alt=sse`;
|
|
const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl);
|
|
|
|
try {
|
|
const response = await axios.post(url, content, {
|
|
headers: {
|
|
Authorization: `Bearer ${accessToken}`,
|
|
'Content-Type': 'application/json',
|
|
},
|
|
responseType: 'stream',
|
|
timeout: 60000,
|
|
proxy: axiosProxy,
|
|
});
|
|
return response.data;
|
|
} catch (error) {
|
|
if (axios.isAxiosError(error)) {
|
|
this.logger.error(`Gemini stream request failed: ${error.message}`);
|
|
throw new UpstreamRequestError({
|
|
message: error.response?.data?.error?.message || error.message,
|
|
status: error.response?.status,
|
|
headers: {
|
|
retryAfter: this.extractRetryAfterHeader(error.response?.headers),
|
|
},
|
|
body: this.describeAxiosErrorData(error.response?.data),
|
|
});
|
|
}
|
|
this.throwAsCleanError(error);
|
|
}
|
|
}
|
|
|
|
async generate(
|
|
model: string,
|
|
content: GeminiRequest,
|
|
accessToken: string,
|
|
upstreamProxyUrl?: string,
|
|
): Promise<GeminiResponse> {
|
|
const url = `${this.baseUrl}/models/${model}:generateContent`;
|
|
const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl);
|
|
|
|
try {
|
|
const response = await axios.post<GeminiResponse>(url, content, {
|
|
headers: {
|
|
Authorization: `Bearer ${accessToken}`,
|
|
'Content-Type': 'application/json',
|
|
},
|
|
timeout: 60000, // 60s timeout
|
|
proxy: axiosProxy,
|
|
});
|
|
return response.data;
|
|
} catch (error) {
|
|
if (axios.isAxiosError(error)) {
|
|
this.logger.error(
|
|
`Gemini request failed: ${error.message} - ${this.safeStringify(error.response?.data)}`,
|
|
);
|
|
throw new UpstreamRequestError({
|
|
message: error.response?.data?.error?.message || error.message,
|
|
status: error.response?.status,
|
|
headers: {
|
|
retryAfter: this.extractRetryAfterHeader(error.response?.headers),
|
|
},
|
|
body: this.describeAxiosErrorData(error.response?.data),
|
|
});
|
|
}
|
|
this.throwAsCleanError(error);
|
|
}
|
|
}
|
|
|
|
// --- Internal Gateway API Support ---
|
|
|
|
async streamGenerateInternal(
|
|
body: GeminiInternalRequest,
|
|
accessToken: string,
|
|
upstreamProxyUrl?: string,
|
|
extraHeaders?: Record<string, string>,
|
|
): Promise<NodeJS.ReadableStream> {
|
|
const response = await this.executeRequestWithEndpointFailover<NodeJS.ReadableStream>(
|
|
':streamGenerateContent?alt=sse',
|
|
body,
|
|
accessToken,
|
|
upstreamProxyUrl,
|
|
{
|
|
responseType: 'stream',
|
|
},
|
|
'stream-generate',
|
|
extraHeaders,
|
|
);
|
|
|
|
return response.data;
|
|
}
|
|
|
|
async generateInternal(
|
|
body: GeminiInternalRequest,
|
|
accessToken: string,
|
|
upstreamProxyUrl?: string,
|
|
extraHeaders?: Record<string, string>,
|
|
): Promise<GeminiResponse> {
|
|
const response = await this.executeRequestWithEndpointFailover<
|
|
GeminiResponse | { response: GeminiResponse }
|
|
>(
|
|
':generateContent',
|
|
body,
|
|
accessToken,
|
|
upstreamProxyUrl,
|
|
{},
|
|
'generate-content',
|
|
extraHeaders,
|
|
);
|
|
const payload = response.data;
|
|
if (payload && typeof payload === 'object' && 'response' in payload) {
|
|
return payload.response;
|
|
}
|
|
return payload;
|
|
}
|
|
|
|
private getInternalBaseUrls(): string[] {
|
|
const fromEnv =
|
|
process.env.PROXY_INTERNAL_BASE_URLS ?? process.env.ANTIGRAVITY_INTERNAL_BASE_URLS;
|
|
const configuredBaseUrls = fromEnv
|
|
?.split(',')
|
|
.map((item) => item.trim())
|
|
.filter(Boolean);
|
|
|
|
if (configuredBaseUrls && configuredBaseUrls.length > 0) {
|
|
return configuredBaseUrls.map((url) => url.replace(/\/+$/, ''));
|
|
}
|
|
|
|
return this.defaultInternalBaseUrls.map((url) => url.replace(/\/+$/, ''));
|
|
}
|
|
|
|
private getInternalTimeoutMs(): number {
|
|
const config = getServerConfig();
|
|
const timeoutSeconds = config?.request_timeout ?? 60;
|
|
return Math.max(1, timeoutSeconds) * 1000;
|
|
}
|
|
|
|
private shouldFailoverToNextEndpoint(error: unknown): boolean {
|
|
if (!axios.isAxiosError(error)) {
|
|
return false;
|
|
}
|
|
|
|
if (!error.response) {
|
|
return true;
|
|
}
|
|
|
|
const status = error.response.status;
|
|
|
|
// Permanent auth errors should fail fast for current token.
|
|
if (status === 401 || status === 403) {
|
|
return false;
|
|
}
|
|
|
|
return status === 408 || status === 429 || status >= 500;
|
|
}
|
|
|
|
private async executeRequestWithEndpointFailover<T>(
|
|
path: string,
|
|
body: GeminiInternalRequest,
|
|
accessToken: string,
|
|
upstreamProxyUrl: string | undefined,
|
|
config: AxiosRequestConfig,
|
|
operation: string,
|
|
extraHeaders?: Record<string, string>,
|
|
): Promise<AxiosResponse<T>> {
|
|
const baseUrls = this.getInternalBaseUrls();
|
|
const timeout = this.getInternalTimeoutMs();
|
|
const requestUserAgent = await resolveRequestUserAgent();
|
|
const axiosProxy = this.resolveUpstreamAxiosProxy(upstreamProxyUrl);
|
|
let lastError: unknown = null;
|
|
|
|
for (let index = 0; index < baseUrls.length; index++) {
|
|
const baseUrl = baseUrls[index];
|
|
const url = `${baseUrl}${path}`;
|
|
|
|
try {
|
|
return await axios.post<T>(url, body, {
|
|
headers: {
|
|
Authorization: `Bearer ${accessToken}`,
|
|
'Content-Type': 'application/json',
|
|
'User-Agent': requestUserAgent,
|
|
...(extraHeaders ?? {}),
|
|
},
|
|
timeout,
|
|
proxy: axiosProxy,
|
|
...config,
|
|
});
|
|
} catch (error) {
|
|
lastError = error;
|
|
const hasNextEndpoint = index < baseUrls.length - 1;
|
|
|
|
if (!hasNextEndpoint || !this.shouldFailoverToNextEndpoint(error)) {
|
|
await this.throwUpstreamRequestError(error, operation);
|
|
}
|
|
|
|
this.logger.warn(
|
|
`[${operation}] request failed at ${baseUrl}; trying next endpoint (${index + 2}/${
|
|
baseUrls.length
|
|
}).`,
|
|
);
|
|
}
|
|
}
|
|
|
|
await this.throwUpstreamRequestError(lastError, operation);
|
|
throw new Error(`[${operation}] unexpected control flow after upstream error handling`);
|
|
}
|
|
|
|
private async throwUpstreamRequestError(error: unknown, operation: string): Promise<never> {
|
|
if (axios.isAxiosError(error)) {
|
|
const responseData = error.response?.data;
|
|
const upstreamMessage = await this.extractAxiosErrorMessage(responseData);
|
|
this.logger.error(
|
|
`[${operation}] upstream request error: ${error.message} - ${this.describeAxiosErrorData(
|
|
responseData,
|
|
)}`,
|
|
);
|
|
throw new UpstreamRequestError({
|
|
message: upstreamMessage || error.message,
|
|
status: error.response?.status,
|
|
headers: {
|
|
retryAfter: this.extractRetryAfterHeader(error.response?.headers),
|
|
},
|
|
body: this.describeAxiosErrorData(responseData),
|
|
});
|
|
}
|
|
this.throwAsCleanError(error);
|
|
}
|
|
|
|
private async extractAxiosErrorMessage(responseData: unknown): Promise<string | null> {
|
|
const fromObject = this.extractAxiosErrorMessageFromObject(responseData);
|
|
if (fromObject) {
|
|
return fromObject;
|
|
}
|
|
|
|
if (typeof responseData === 'string') {
|
|
return this.extractAxiosErrorMessageFromText(responseData);
|
|
}
|
|
|
|
if (Buffer.isBuffer(responseData)) {
|
|
return this.extractAxiosErrorMessageFromText(responseData.toString('utf-8'));
|
|
}
|
|
|
|
if (this.isReadableStream(responseData)) {
|
|
const streamText = await this.readStreamAsText(responseData);
|
|
return streamText ? this.extractAxiosErrorMessageFromText(streamText) : null;
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
private extractAxiosErrorMessageFromObject(responseData: unknown): string | null {
|
|
if (!responseData || typeof responseData !== 'object' || this.isReadableStream(responseData)) {
|
|
return null;
|
|
}
|
|
|
|
const errorRecord = (responseData as { error?: unknown }).error;
|
|
if (errorRecord && typeof errorRecord === 'object') {
|
|
const message = (errorRecord as { message?: unknown }).message;
|
|
if (typeof message === 'string' && message.trim() !== '') {
|
|
return message.trim();
|
|
}
|
|
}
|
|
|
|
const message = (responseData as { message?: unknown }).message;
|
|
if (typeof message === 'string' && message.trim() !== '') {
|
|
return message.trim();
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
private extractAxiosErrorMessageFromText(rawText: string): string | null {
|
|
const text = rawText.trim();
|
|
if (!text) {
|
|
return null;
|
|
}
|
|
|
|
const sseLines = text
|
|
.split('\n')
|
|
.map((line) => line.trim())
|
|
.filter((line) => line.startsWith('data:'));
|
|
|
|
for (const line of sseLines) {
|
|
const payload = line.slice(5).trim();
|
|
if (!payload || payload === '[DONE]') {
|
|
continue;
|
|
}
|
|
const parsed = this.tryParseJson(payload);
|
|
const message = this.extractAxiosErrorMessageFromObject(parsed);
|
|
if (message) {
|
|
return message;
|
|
}
|
|
}
|
|
|
|
const parsed = this.tryParseJson(text);
|
|
const fromJson = this.extractAxiosErrorMessageFromObject(parsed);
|
|
if (fromJson) {
|
|
return fromJson;
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
private tryParseJson(value: string): unknown {
|
|
try {
|
|
return JSON.parse(value);
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
private isReadableStream(value: unknown): value is NodeJS.ReadableStream {
|
|
return (
|
|
Boolean(value) &&
|
|
typeof value === 'object' &&
|
|
typeof (value as { pipe?: unknown }).pipe === 'function'
|
|
);
|
|
}
|
|
|
|
private async readStreamAsText(stream: NodeJS.ReadableStream): Promise<string | null> {
|
|
return new Promise((resolve) => {
|
|
let buffer = '';
|
|
const maxChars = 512 * 1024;
|
|
|
|
stream.on('data', (chunk: Buffer | string) => {
|
|
const text = Buffer.isBuffer(chunk) ? chunk.toString('utf-8') : String(chunk);
|
|
if (buffer.length >= maxChars) {
|
|
return;
|
|
}
|
|
buffer += text;
|
|
if (buffer.length > maxChars) {
|
|
buffer = buffer.slice(0, maxChars);
|
|
}
|
|
});
|
|
|
|
stream.on('end', () => resolve(buffer));
|
|
stream.on('error', () => resolve(null));
|
|
});
|
|
}
|
|
|
|
private describeAxiosErrorData(responseData: unknown): string {
|
|
if (this.isReadableStream(responseData)) {
|
|
return '[stream]';
|
|
}
|
|
return this.safeStringify(responseData);
|
|
}
|
|
|
|
private resolveUpstreamAxiosProxy(
|
|
upstreamProxyUrl?: string,
|
|
): AxiosProxyConfig | false | undefined {
|
|
const config = getServerConfig();
|
|
const configuredProxyUrl =
|
|
upstreamProxyUrl ||
|
|
(config?.upstream_proxy?.enabled && config.upstream_proxy.url
|
|
? config.upstream_proxy.url
|
|
: '');
|
|
|
|
if (!configuredProxyUrl) {
|
|
return undefined;
|
|
}
|
|
|
|
try {
|
|
const parsed = new URL(configuredProxyUrl);
|
|
const port = parsed.port ? Number(parsed.port) : parsed.protocol === 'https:' ? 443 : 80;
|
|
|
|
const proxyConfig: AxiosProxyConfig = {
|
|
protocol: parsed.protocol.replace(':', ''),
|
|
host: parsed.hostname,
|
|
port,
|
|
};
|
|
|
|
if (parsed.username || parsed.password) {
|
|
proxyConfig.auth = {
|
|
username: decodeURIComponent(parsed.username),
|
|
password: decodeURIComponent(parsed.password),
|
|
};
|
|
}
|
|
|
|
return proxyConfig;
|
|
} catch {
|
|
this.logger.warn(`Upstream proxy URL is invalid: ${configuredProxyUrl}`);
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
private extractRetryAfterHeader(headers: unknown): string | undefined {
|
|
if (!headers || typeof headers !== 'object') {
|
|
return undefined;
|
|
}
|
|
|
|
const retryAfter = (headers as Record<string, unknown>)['retry-after'];
|
|
if (typeof retryAfter === 'string' && retryAfter.trim() !== '') {
|
|
return retryAfter.trim();
|
|
}
|
|
if (Array.isArray(retryAfter) && retryAfter.length > 0) {
|
|
const first = retryAfter[0];
|
|
if (typeof first === 'string' && first.trim() !== '') {
|
|
return first.trim();
|
|
}
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
private throwAsCleanError(error: unknown): never {
|
|
// Re-throw as clean Error to avoid circular reference issues.
|
|
throw error instanceof Error ? new Error(error.message) : new Error(String(error));
|
|
}
|
|
|
|
/**
|
|
* Safely stringify an object, handling circular references
|
|
*/
|
|
private safeStringify(obj: unknown): string {
|
|
if (isNil(obj)) {
|
|
return String(obj);
|
|
}
|
|
try {
|
|
const seen = new WeakSet();
|
|
return JSON.stringify(obj, (key, value) => {
|
|
if (typeof value === 'object' && value !== null) {
|
|
if (seen.has(value)) return '[Circular]';
|
|
seen.add(value);
|
|
}
|
|
return value;
|
|
});
|
|
} catch {
|
|
return '[Unserializable]';
|
|
}
|
|
}
|
|
}
|