diff --git a/eslint.config.ts b/eslint.config.ts index eacd510..a2f4245 100644 --- a/eslint.config.ts +++ b/eslint.config.ts @@ -9,7 +9,7 @@ export default tseslint.config( ...tseslint.configs.strictTypeChecked, ...tseslint.configs.stylisticTypeChecked, { - ignores: ['**/node_modules/*', '**/*.mjs', '**/*.js', 'src/mapper.ts'], + ignores: ['**/node_modules/*', '**/*.mjs', '**/*.js'], }, { languageOptions: { diff --git a/src/chatwrapper.ts b/src/chatwrapper.ts index 37810a4..d98d434 100644 --- a/src/chatwrapper.ts +++ b/src/chatwrapper.ts @@ -12,6 +12,7 @@ import { } from '@google/gemini-cli-core/dist/src/config/models.js'; import { Content, GeminiResponse, Model } from './types.js'; +import consola from 'consola'; /* ------------------------------------------------------------------ */ /* 1. Build the ContentGenerator exactly like the CLI does */ @@ -32,6 +33,33 @@ const generatorPromise: Promise = (async () => { /* ------------------------------------------------------------------ */ type GenConfig = Record; +const MAX_RETRIES = 3; +const INITIAL_RETRY_DELAY = 1000; // 1 second + +async function withRetry(operation: () => Promise): Promise { + let retries = 0; + while (true) { + try { + return await operation(); + } catch (error) { + // Check if it's an Error object with a message property + if (!(error instanceof Error) || + !error.message.includes('RESOURCE_EXHAUSTED') || + retries >= MAX_RETRIES) { + throw error; + } + retries++; + const delay = INITIAL_RETRY_DELAY * Math.pow(2, retries - 1); + + consola.error( + `Rate limit hit, retrying in ${delay}ms ` + + `(attempt ${retries}/${MAX_RETRIES})`, + ); + await new Promise(resolve => setTimeout(resolve, delay)); + } + } +} + export async function sendChat({ contents, generationConfig = {}, @@ -41,11 +69,11 @@ export async function sendChat({ tools?: unknown, // accepted but ignored for now }): Promise { const generator = await generatorPromise; - const gResp = await generator.generateContent({ + const gResp = await withRetry(() => generator.generateContent({ model: modelName, contents, config: generationConfig, - }); + })); return { text: gResp.text ?? '', usageMetadata: { @@ -65,11 +93,11 @@ export async function* sendChatStream({ tools?: unknown, }) { const generator = await generatorPromise; - const stream = await generator.generateContentStream({ + const stream = await withRetry(() => generator.generateContentStream({ model: modelName, contents, config: generationConfig, - }); + })); for await (const chunk of stream) yield chunk; } @@ -94,3 +122,4 @@ export function listModels(): Model[] { // export async function embed(_input: unknown) { // throw new Error('Embeddings endpoint not implemented yet.'); // } + diff --git a/src/mapper.ts b/src/mapper.ts index e746509..f917560 100644 --- a/src/mapper.ts +++ b/src/mapper.ts @@ -2,16 +2,25 @@ /* mapper.ts – OpenAI ⇆ Gemini (with reasoning/1 M context) */ /* ------------------------------------------------------------------ */ import { fetchAndEncode } from './remoteimage'; -import { z } from 'zod'; -import { ToolRegistry } from '@google/gemini-cli-core/dist/src/tools/tool-registry.js'; -import { RequestBody } from './types'; +import { z, ZodRawShape } from 'zod'; +import { ToolRegistry } + from '@google/gemini-cli-core/dist/src/tools/tool-registry.js'; +import { Config } from '@google/gemini-cli-core/dist/src/config/config.js'; +import { Tool } from '@google/gemini-cli-core/dist/src/tools/tools.js'; +import { + Part, + RequestBody, + GeminiResponse, + GeminiStreamChunk, +} from './types'; -/* ------------------------------------------------------------------ */ -interface Part { text?: string; inlineData?: { mimeType: string, data: string } } - -/* ------------------------------------------------------------------ */ -function callLocalFunction(_name: string, _args: unknown) { - return { ok: true }; +/* ----------------------------------------------------------------- */ +async function callLocalFunction(/*_name: string, _args: unknown*/) { + return Promise.resolve({ + ok: true, + llmContent: [], + returnDisplay: 'Function executed successfully', + }); } /* ================================================================== */ @@ -61,21 +70,29 @@ export async function mapRequest(body: RequestBody) { }; /* ---- Tool / function mapping ----------------------------------- */ - const tools = new ToolRegistry({} as any); + // Note: ToolRegistry expects a complex Config object that we don't have + // access to. Casting to `Config` is a workaround. + const tools = new ToolRegistry({} as Config); if (body.functions?.length) { - const reg = tools as any; - body.functions.forEach((fn: any) => - reg.registerTool( - fn.name, - { - title: fn.name, - description: fn.description ?? '', - inputSchema: z.object(fn.parameters?.properties ?? {}), - }, - (args: unknown) => callLocalFunction(fn.name, args), - ), - ); + for (const fn of body.functions) { + tools.registerTool({ + name: fn.name, + displayName: fn.name, + description: fn.description ?? '', + schema: z.object( + (fn.parameters?.properties as ZodRawShape) ?? {}, + ), + isOutputMarkdown: false, + canUpdateOutput: false, + validateToolParams: () => null, + getDescription: (params: unknown) => + `Executing ${fn.name} with parameters: ` + + JSON.stringify(params), + shouldConfirmExecute: () => Promise.resolve(false), + execute: () => callLocalFunction(), + } as Tool); + } } return { geminiReq, tools }; @@ -84,8 +101,12 @@ export async function mapRequest(body: RequestBody) { /* ================================================================== */ /* Non-stream response: Gemini ➞ OpenAI */ /* ================================================================== */ -export function mapResponse(gResp: any) { - const usage = gResp.usageMetadata ?? {}; +export function mapResponse(gResp: GeminiResponse) { + const usage = gResp.usageMetadata ?? { + promptTokens: 0, + candidatesTokens: 0, + totalTokens: 0, + }; return { id: `chatcmpl-${Date.now()}`, object: 'chat.completion', @@ -99,9 +120,9 @@ export function mapResponse(gResp: any) { }, ], usage: { - prompt_tokens: usage.promptTokens ?? 0, - completion_tokens: usage.candidatesTokens ?? 0, - total_tokens: usage.totalTokens ?? 0, + prompt_tokens: usage.promptTokens, + completion_tokens: usage.candidatesTokens, + total_tokens: usage.totalTokens, }, }; } @@ -110,9 +131,9 @@ export function mapResponse(gResp: any) { /* Stream chunk mapper: Gemini ➞ OpenAI */ /* ================================================================== */ -export function mapStreamChunk(chunk: any) { +export function mapStreamChunk(chunk: GeminiStreamChunk) { const part = chunk?.candidates?.[0]?.content?.parts?.[0] ?? {}; - const delta: any = { role: 'assistant' }; + const delta: { role: 'assistant', content?: string } = { role: 'assistant' }; if (part.thought === true) { delta.content = `${part.text ?? ''}`; // ST renders grey bubble diff --git a/src/server.ts b/src/server.ts index a74240d..bb84385 100644 --- a/src/server.ts +++ b/src/server.ts @@ -2,7 +2,7 @@ import consola from 'consola'; import http from 'http'; import { listModels, sendChat, sendChatStream } from './chatwrapper'; import { mapRequest, mapResponse, mapStreamChunk } from './mapper.js'; -import { RequestBody, GeminiResponse } from './types'; +import { RequestBody, GeminiResponse, GeminiStreamChunk, Part } from './types'; import { config } from './config'; /* ── basic config ─────────────────────────────────────────────────── */ @@ -41,8 +41,11 @@ function readJSON( error: { message: 'Request body is missing for POST request' }, }), ); + resolve(null); + return; } - return resolve(null); + resolve(null); + return; } try { resolve(JSON.parse(data) as RequestBody); @@ -51,6 +54,7 @@ function readJSON( res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(JSON.stringify({ error: { message: 'Malformed JSON' } })); resolve(null); + return; } }); }); @@ -101,7 +105,35 @@ http }); for await (const chunk of sendChatStream({ ...geminiReq, tools })) { - res.write(`data: ${JSON.stringify(mapStreamChunk(chunk))}\n\n`); + // Transform the chunk to match our expected type + const transformedParts = + chunk.candidates?.[0]?.content?.parts?.map(part => { + const transformedPart: Part = { + text: part.text, + thought: part.text?.startsWith?.('') ?? false, + }; + + if (part.inlineData?.data) { + transformedPart.inlineData = { + mimeType: part.inlineData.mimeType ?? 'text/plain', + data: part.inlineData.data, + }; + } + + return transformedPart; + }) ?? []; + + const streamChunk: GeminiStreamChunk = { + candidates: [{ + content: { + parts: transformedParts, + }, + }], + }; + + res.write( + `data: ${JSON.stringify(mapStreamChunk(streamChunk))}\n\n`, + ); } res.end('data: [DONE]\n\n'); } else { @@ -112,10 +144,23 @@ http } catch (err) { const error = err as Error; consola.error('Proxy error ➜', error); - res.writeHead(500, { 'Content-Type': 'application/json' }); - res.end(JSON.stringify({ error: { message: error.message } })); + + // For streaming responses, send error in stream format + if (body.stream && res.headersSent) { + res.write(`data: ${JSON.stringify({ + error: { + message: error.message, + type: 'error', + }, + })}\n\n`); + res.end('data: [DONE]\n\n'); + return; + } else { + // For non-streaming responses or if headers haven't been sent yet + res.writeHead(500, { 'Content-Type': 'application/json' }); + res.end(JSON.stringify({ error: { message: error.message } })); + } } - return; } /* ---- anything else ---------- */ diff --git a/src/types.ts b/src/types.ts index d3cef57..f09cb64 100644 --- a/src/types.ts +++ b/src/types.ts @@ -6,9 +6,16 @@ export interface Model { object: 'model'; owned_by: 'google'; } + +interface InlineData { + mimeType: string; + data: string; +} + export interface Part { - text?: string; - inlineData?: { mimeType: string, data: string }; + text?: string; + inlineData?: InlineData; + thought?: boolean; } export interface Content { @@ -16,6 +23,14 @@ export interface Content { parts: Part[]; } +interface FunctionDef { + name: string; + description?: string; + parameters?: { + properties?: Record, + }; +} + export interface RequestBody { messages: { content: @@ -28,13 +43,7 @@ export interface RequestBody { generationConfig?: Record; include_reasoning?: boolean; stream?: boolean; - functions?: { - name: string, - description?: string, - parameters?: { - properties?: Record, - }, - }[]; + functions?: FunctionDef[]; } export interface GeminiResponse { @@ -44,4 +53,12 @@ export interface GeminiResponse { candidatesTokens: number, totalTokens: number, }; -} \ No newline at end of file +} + +export interface GeminiStreamChunk { + candidates?: { + content?: { + parts?: Part[], + }, + }[]; +}