JimLiu-baoyu-skills/skills/baoyu-danger-gemini-web/scripts/gemini-webapi/client.ts

628 lines
21 KiB
TypeScript

import { Endpoint, ErrorCode, Headers, Model } from './constants.js';
import { GemMixin } from './components/gem-mixin.js';
import {
APIError,
AuthError,
GeminiError,
ImageGenerationError,
ModelInvalid,
TemporarilyBlocked,
TimeoutError,
UsageLimitExceeded,
} from './exceptions.js';
import { Candidate, Gem, GeneratedImage, ModelOutput, RPCData, WebImage } from './types/index.js';
import {
extract_json_from_response,
get_access_token,
get_nested_value,
logger,
parse_file_name,
rotate_1psidts,
rotate_tasks,
fetch_with_timeout,
sleep,
upload_file,
write_cookie_file,
resolveGeminiWebCookiePath,
} from './utils/index.js';
type InitOptions = {
timeout?: number;
auto_close?: boolean;
close_delay?: number;
auto_refresh?: boolean;
refresh_interval?: number;
verbose?: boolean;
};
type RequestKwargs = RequestInit & { timeout_ms?: number };
function normalize_headers(h?: HeadersInit): Record<string, string> {
if (!h) return {};
if (Array.isArray(h)) return Object.fromEntries(h.map(([k, v]) => [k, v]));
if (h instanceof Headers) {
const out: Record<string, string> = {};
h.forEach((v, k) => {
out[k] = v;
});
return out;
}
return { ...(h as Record<string, string>) };
}
function collect_strings(root: unknown, accept: (s: string) => boolean, limit: number = 20): string[] {
const out: string[] = [];
const seen = new Set<string>();
const stack: unknown[] = [root];
while (stack.length > 0 && out.length < limit) {
const v = stack.pop();
if (typeof v === 'string') {
if (accept(v) && !seen.has(v)) {
seen.add(v);
out.push(v);
}
continue;
}
if (Array.isArray(v)) {
for (let i = 0; i < v.length; i++) stack.push(v[i]);
continue;
}
if (v && typeof v === 'object') {
for (const val of Object.values(v as Record<string, unknown>)) stack.push(val);
}
}
return out;
}
export class GeminiClient extends GemMixin {
public cookies: Record<string, string> = {};
public proxy: string | null = null;
public _running: boolean = false;
public access_token: string | null = null;
public timeout: number = 300;
public auto_close: boolean = false;
public close_delay: number = 300;
public auto_refresh: boolean = true;
public refresh_interval: number = 540;
public kwargs: RequestInit;
private close_timer: ReturnType<typeof setTimeout> | null = null;
private refresh_abort: AbortController | null = null;
constructor(
secure_1psid: string | null = null,
secure_1psidts: string | null = null,
proxy: string | null = null,
kwargs: RequestInit = {},
) {
super();
this.proxy = proxy;
this.kwargs = kwargs;
if (secure_1psid) {
this.cookies['__Secure-1PSID'] = secure_1psid;
if (secure_1psidts) this.cookies['__Secure-1PSIDTS'] = secure_1psidts;
}
}
async init(
timeoutOrOpts: number | InitOptions = 300,
auto_close: boolean = false,
close_delay: number = 300,
auto_refresh: boolean = true,
refresh_interval: number = 540,
verbose: boolean = true,
): Promise<void> {
const opts: InitOptions =
typeof timeoutOrOpts === 'object'
? timeoutOrOpts
: { timeout: timeoutOrOpts, auto_close, close_delay, auto_refresh, refresh_interval, verbose };
const timeout = opts.timeout ?? 300;
const ac = opts.auto_close ?? false;
const cd = opts.close_delay ?? 300;
const ar = opts.auto_refresh ?? true;
const ri = opts.refresh_interval ?? 540;
const vb = opts.verbose ?? true;
try {
const [token, valid] = await get_access_token(this.cookies, this.proxy, vb);
this.access_token = token;
this.cookies = valid;
this._running = true;
this.timeout = timeout;
this.auto_close = ac;
this.close_delay = cd;
if (this.auto_close) await this.reset_close_task();
this.auto_refresh = ar;
this.refresh_interval = ri;
const sid = this.cookies['__Secure-1PSID'];
if (sid) {
const existing = rotate_tasks.get(sid);
if (existing && existing instanceof AbortController) existing.abort();
rotate_tasks.delete(sid);
}
if (this.auto_refresh && sid) {
const ctl = new AbortController();
this.refresh_abort?.abort();
this.refresh_abort = ctl;
rotate_tasks.set(sid, ctl);
void this.start_auto_refresh(ctl.signal);
}
await write_cookie_file(this.cookies, resolveGeminiWebCookiePath(), 'client').catch(() => {});
if (vb) logger.success('Gemini client initialized successfully.');
} catch (e) {
await this.close();
throw e;
}
}
async close(delay: number = 0): Promise<void> {
if (delay > 0) await sleep(delay * 1000);
this._running = false;
if (this.close_timer) {
clearTimeout(this.close_timer);
this.close_timer = null;
}
this.refresh_abort?.abort();
this.refresh_abort = null;
const sid = this.cookies['__Secure-1PSID'];
const t = sid ? rotate_tasks.get(sid) : null;
if (t && t instanceof AbortController) t.abort();
if (sid) rotate_tasks.delete(sid);
}
async reset_close_task(): Promise<void> {
if (this.close_timer) {
clearTimeout(this.close_timer);
this.close_timer = null;
}
this.close_timer = setTimeout(() => {
void this.close(0);
}, this.close_delay * 1000);
this.close_timer.unref?.();
}
async start_auto_refresh(signal: AbortSignal): Promise<void> {
while (!signal.aborted) {
let newTs: string | null = null;
try {
newTs = await rotate_1psidts(this.cookies, this.proxy);
} catch (e) {
if (e instanceof AuthError) {
logger.warning('AuthError: Failed to refresh cookies. Auto refresh task canceled.');
return;
}
logger.warning(`Unexpected error while refreshing cookies: ${e instanceof Error ? e.message : String(e)}`);
}
if (newTs) {
this.cookies['__Secure-1PSIDTS'] = newTs;
await write_cookie_file(this.cookies, resolveGeminiWebCookiePath(), 'refresh').catch(() => {});
logger.debug('Cookies refreshed. New __Secure-1PSIDTS applied.');
}
await sleep(this.refresh_interval * 1000, signal);
}
}
protected async _run<T>(fn: () => Promise<T>, retry: number): Promise<T> {
try {
if (!this._running) {
await this.init({
timeout: this.timeout,
auto_close: this.auto_close,
close_delay: this.close_delay,
auto_refresh: this.auto_refresh,
refresh_interval: this.refresh_interval,
verbose: false,
});
if (!this._running) {
throw new APIError('Client initialization failed.');
}
}
return await fn();
} catch (e) {
let r = retry;
if (e instanceof ImageGenerationError) r = Math.min(1, r);
if (e instanceof APIError && r > 0) {
await sleep(1000);
return await this._run(fn, r - 1);
}
throw e;
}
}
async generate_content(
prompt: string,
files: string[] | null = null,
model: Model | string | Record<string, unknown> = Model.UNSPECIFIED,
gem: Gem | string | null = null,
chat: ChatSession | null = null,
kwargs: RequestKwargs = {},
): Promise<ModelOutput> {
return await this._run(async () => {
if (!prompt) throw new Error('Prompt cannot be empty.');
let mdl: Model;
if (typeof model === 'string') mdl = Model.from_name(model);
else if (model instanceof Model) mdl = model;
else if (model && typeof model === 'object') mdl = Model.from_dict(model);
else throw new TypeError(`'model' must be a Model instance, string, or dictionary; got ${typeof model}`);
const gem_id = gem instanceof Gem ? gem.id : gem;
if (this.auto_close) await this.reset_close_task();
if (!this.access_token) throw new APIError('Missing access token.');
const f = files?.length ? files : null;
const uploaded =
f &&
(await Promise.all(
f.map(async (p) => [[await upload_file(p, this.proxy)], parse_file_name(p)] as [string[], string]),
));
const first = uploaded ? [prompt, 0, null, uploaded] : [prompt];
const inner: unknown[] = [first, null, chat ? chat.metadata : null];
if (gem_id) {
for (let i = 0; i < 16; i++) inner.push(null);
inner.push(gem_id);
}
const f_req = JSON.stringify([null, JSON.stringify(inner)]);
const body = new URLSearchParams({ at: this.access_token, 'f.req': f_req }).toString();
const h0 = { ...Headers.GEMINI, ...mdl.model_header, Cookie: Object.entries(this.cookies).map(([k, v]) => `${k}=${v}`).join('; ') };
const h1 = { ...h0, ...normalize_headers(kwargs.headers) };
let res: Response;
try {
const timeout_ms = typeof kwargs.timeout_ms === 'number' ? kwargs.timeout_ms : this.timeout * 1000;
const { timeout_ms: _t, ...rest } = kwargs;
res = await fetch_with_timeout(Endpoint.GENERATE, {
method: 'POST',
headers: h1,
body,
redirect: 'follow',
...this.kwargs,
...rest,
timeout_ms,
});
} catch (e) {
throw new TimeoutError(
`Generate content request timed out, please try again. If the problem persists, consider setting a higher 'timeout' value when initializing GeminiClient. (${e instanceof Error ? e.message : String(e)})`,
);
}
if (res.status !== 200) {
await this.close();
throw new APIError(`Failed to generate contents. Request failed with status code ${res.status}`);
}
const txt = await res.text();
const response_json = extract_json_from_response(txt);
let body_json: unknown[] | null = null;
let body_index = 0;
try {
if (!Array.isArray(response_json)) throw new Error('Invalid JSON');
for (let part_index = 0; part_index < response_json.length; part_index++) {
const part = response_json[part_index];
if (!Array.isArray(part)) continue;
const part_body = get_nested_value<string | null>(part, [2], null);
if (!part_body) continue;
try {
const part_json = JSON.parse(part_body) as unknown[];
if (get_nested_value(part_json, [4], null)) {
body_index = part_index;
body_json = part_json;
break;
}
} catch {}
}
if (!body_json) throw new Error('No body');
} catch {
await this.close();
try {
const code = get_nested_value<number>(response_json, [0, 5, 2, 0, 1, 0], -1);
if (code === ErrorCode.USAGE_LIMIT_EXCEEDED) {
throw new UsageLimitExceeded(
`Failed to generate contents. Usage limit of ${mdl.model_name} model has exceeded. Please try switching to another model.`,
);
}
if (code === ErrorCode.MODEL_INCONSISTENT) {
throw new ModelInvalid(
'Failed to generate contents. The specified model is inconsistent with the chat history. Please make sure to pass the same `model` parameter when starting a chat session with previous metadata.',
);
}
if (code === ErrorCode.MODEL_HEADER_INVALID) {
throw new ModelInvalid(
'Failed to generate contents. The specified model is not available. Please update gemini_webapi to the latest version. If the error persists and is caused by the package, please report it on GitHub.',
);
}
if (code === ErrorCode.IP_TEMPORARILY_BLOCKED) {
throw new TemporarilyBlocked(
'Failed to generate contents. Your IP address is temporarily blocked by Google. Please try using a proxy or waiting for a while.',
);
}
} catch (e) {
if (e instanceof GeminiError) throw e;
}
logger.debug(`Invalid response: ${txt.slice(0, 500)}`);
throw new APIError('Failed to generate contents. Invalid response data received. Client will try to re-initialize on next request.');
}
try {
const candidate_list = get_nested_value<unknown[]>(body_json, [4], []);
const out: Candidate[] = [];
for (let candidate_index = 0; candidate_index < candidate_list.length; candidate_index++) {
const candidate = candidate_list[candidate_index];
if (!Array.isArray(candidate)) continue;
const rcid = get_nested_value<string | null>(candidate, [0], null);
if (!rcid) continue;
let text = String(get_nested_value(candidate, [1, 0], ''));
if (/^http:\/\/googleusercontent\.com\/card_content\/\d+/.test(text)) {
text = String(get_nested_value(candidate, [22, 0], text));
}
const thoughts = get_nested_value<string | null>(candidate, [37, 0, 0], null);
const web_images: WebImage[] = [];
for (const w of get_nested_value<unknown[]>(candidate, [12, 1], [])) {
if (!Array.isArray(w)) continue;
const url = get_nested_value<string | null>(w, [0, 0, 0], null);
if (!url) continue;
web_images.push(new WebImage(url, String(get_nested_value(w, [7, 0], '')), String(get_nested_value(w, [0, 4], '')), this.proxy));
}
const generated_images: GeneratedImage[] = [];
const wants_generated =
get_nested_value(candidate, [12, 7, 0], null) != null ||
/http:\/\/googleusercontent\.com\/image_generation_content\/\d+/.test(text);
if (wants_generated) {
let img_body: unknown[] | null = null;
for (let part_index = body_index; part_index < (response_json as unknown[]).length; part_index++) {
const part = (response_json as unknown[])[part_index];
if (!Array.isArray(part)) continue;
const part_body = get_nested_value<string | null>(part, [2], null);
if (!part_body) continue;
try {
const part_json = JSON.parse(part_body) as unknown[];
const cand = get_nested_value<unknown>(part_json, [4, candidate_index], null);
if (!cand) continue;
const urls = collect_strings(cand, (s) => s.startsWith('https://lh3.googleusercontent.com/gg-dl/'), 1);
if (urls.length > 0) {
img_body = part_json;
break;
}
} catch {}
}
if (!img_body) {
throw new ImageGenerationError(
'Failed to parse generated images. Please update gemini_webapi to the latest version. If the error persists and is caused by the package, please report it on GitHub.',
);
}
const img_candidate = get_nested_value<unknown[]>(img_body, [4, candidate_index], []);
const finished = get_nested_value<string | null>(img_candidate, [1, 0], null);
if (finished) {
text = finished.replace(/http:\/\/googleusercontent\.com\/image_generation_content\/\d+/g, '').trimEnd();
}
const gen = get_nested_value<unknown[]>(img_candidate, [12, 7, 0], []);
for (let img_index = 0; img_index < gen.length; img_index++) {
const g = gen[img_index];
if (!Array.isArray(g)) continue;
const url = get_nested_value<string | null>(g, [0, 3, 3], null);
if (!url) continue;
const img_num = get_nested_value<number | null>(g, [3, 6], null);
const title = img_num ? `[Generated Image ${img_num}]` : '[Generated Image]';
const alt_list = get_nested_value<unknown[]>(g, [3, 5], []);
const alt =
(typeof alt_list[img_index] === 'string' ? (alt_list[img_index] as string) : null) ??
(typeof alt_list[0] === 'string' ? (alt_list[0] as string) : '') ??
'';
generated_images.push(new GeneratedImage(url, title, alt, this.proxy, this.cookies));
}
if (generated_images.length === 0) {
const urls = collect_strings(img_candidate, (s) => s.startsWith('https://lh3.googleusercontent.com/gg-dl/'), 4);
for (const url of urls) {
generated_images.push(new GeneratedImage(url, '[Generated Image]', '', this.proxy, this.cookies));
}
}
}
out.push(new Candidate({ rcid, text, thoughts, web_images, generated_images }));
}
if (out.length === 0) {
throw new GeminiError('Failed to generate contents. No output data found in response.');
}
const metadata = get_nested_value<string[]>(body_json, [1], []);
const output = new ModelOutput({ metadata, candidates: out });
if (chat instanceof ChatSession) chat.last_output = output;
return output;
} catch (e) {
if (e instanceof GeminiError || e instanceof APIError) throw e;
throw new APIError('Failed to parse response body. Data structure is invalid.');
}
}, 2);
}
async generateContent(
prompt: string,
files?: string[] | null,
model?: Model | string | Record<string, unknown>,
gem?: Gem | string | null,
chat?: ChatSession | null,
kwargs?: RequestKwargs,
): Promise<ModelOutput> {
return await this.generate_content(prompt, files ?? null, model ?? Model.UNSPECIFIED, gem ?? null, chat ?? null, kwargs ?? {});
}
start_chat(opts?: ConstructorParameters<typeof ChatSession>[1]): ChatSession {
return new ChatSession(this, opts);
}
startChat(opts?: ConstructorParameters<typeof ChatSession>[1]): ChatSession {
return this.start_chat(opts);
}
protected async _batch_execute(payloads: RPCData[], opts: RequestInit = {}): Promise<Response> {
if (!this.access_token) throw new APIError('Missing access token.');
const f_req = JSON.stringify([payloads.map((p) => p.serialize())]);
const body = new URLSearchParams({ at: this.access_token, 'f.req': f_req }).toString();
const h0 = { ...Headers.GEMINI, Cookie: Object.entries(this.cookies).map(([k, v]) => `${k}=${v}`).join('; ') };
const h1 = { ...h0, ...normalize_headers(opts.headers) };
const res = await fetch_with_timeout(Endpoint.BATCH_EXEC, {
method: 'POST',
headers: h1,
body,
redirect: 'follow',
...this.kwargs,
...opts,
timeout_ms: this.timeout * 1000,
});
if (res.status !== 200) {
await this.close();
throw new APIError(`Batch execution failed with status code ${res.status}`);
}
return res;
}
}
export class ChatSession {
private __metadata: Array<string | null> = [null, null, null];
public geminiclient: GeminiClient;
private _last_output: ModelOutput | null = null;
public model: Model | string | Record<string, unknown>;
public gem: Gem | string | null;
constructor(
geminiclient: GeminiClient,
opts: {
metadata?: Array<string | null>;
cid?: string | null;
rid?: string | null;
rcid?: string | null;
model?: Model | string | Record<string, unknown>;
gem?: Gem | string | null;
} = {},
) {
this.geminiclient = geminiclient;
this.model = opts.model ?? Model.UNSPECIFIED;
this.gem = opts.gem ?? null;
if (opts.metadata) this.metadata = opts.metadata;
if (opts.cid) this.cid = opts.cid;
if (opts.rid) this.rid = opts.rid;
if (opts.rcid) this.rcid = opts.rcid;
}
toString(): string {
return `ChatSession(cid='${this.cid}', rid='${this.rid}', rcid='${this.rcid}')`;
}
get last_output(): ModelOutput | null {
return this._last_output;
}
set last_output(v: ModelOutput | null) {
this._last_output = v;
if (v) {
this.metadata = (v.metadata ?? []) as Array<string | null>;
this.rcid = v.rcid;
}
}
async send_message(prompt: string, files: string[] | null = null, kwargs: RequestKwargs = {}): Promise<ModelOutput> {
return await this.geminiclient.generate_content(prompt, files, this.model, this.gem, this, kwargs);
}
async sendMessage(prompt: string, files?: string[] | null, kwargs?: RequestKwargs): Promise<ModelOutput> {
return await this.send_message(prompt, files ?? null, kwargs ?? {});
}
choose_candidate(index: number): ModelOutput {
if (!this.last_output) throw new Error('No previous output data found in this chat session.');
if (index >= this.last_output.candidates.length) {
throw new Error(`Index ${index} exceeds the number of candidates in last model output.`);
}
this.last_output.chosen = index;
this.rcid = this.last_output.rcid;
return this.last_output;
}
chooseCandidate(index: number): ModelOutput {
return this.choose_candidate(index);
}
get metadata(): Array<string | null> {
return this.__metadata;
}
set metadata(v: Array<string | null>) {
if (v.length > 3) throw new Error('metadata cannot exceed 3 elements');
this.__metadata = [null, null, null];
for (let i = 0; i < v.length; i++) this.__metadata[i] = v[i] ?? null;
}
get cid(): string | null {
return this.__metadata[0];
}
set cid(v: string | null) {
this.__metadata[0] = v;
}
get rid(): string | null {
return this.__metadata[1];
}
set rid(v: string | null) {
this.__metadata[1] = v;
}
get rcid(): string | null {
return this.__metadata[2];
}
set rcid(v: string | null) {
this.__metadata[2] = v;
}
}