@just-every/ensemble 0.1.0

package/dist/model_providers/gemini.js

@@ -0,0 +1,711 @@
+import { GoogleGenAI, Type, FunctionCallingConfigMode, } from '@google/genai';
+import { v4 as uuidv4 } from 'uuid';
+import { costTracker } from '@just-every/ensemble/cost_tracker';
+import { log_llm_error, log_llm_request, log_llm_response, } from '../utils/llm_logger.js';
+import { isPaused } from '../utils/communication.js';
+import { extractBase64Image, resizeAndTruncateForGemini, } from '../utils/image_utils.js';
+import { bufferDelta, flushBufferedDeltas, } from '../utils/delta_buffer.js';
+function convertParameterToGeminiFormat(param) {
+    let type = Type.STRING;
+    switch (param.type) {
+        case 'string':
+            type = Type.STRING;
+            break;
+        case 'number':
+            type = Type.NUMBER;
+            break;
+        case 'boolean':
+            type = Type.BOOLEAN;
+            break;
+        case 'object':
+            type = Type.OBJECT;
+            break;
+        case 'array':
+            type = Type.ARRAY;
+            break;
+        case 'null':
+            type = Type.STRING;
+            console.warn(`Mapping 'null' type to STRING`);
+            break;
+        default:
+            console.warn(`Unsupported parameter type '${param.type}'. Defaulting to STRING.`);
+            type = Type.STRING;
+    }
+    const result = { type, description: param.description };
+    if (type === Type.ARRAY) {
+        if (param.items) {
+            if (param.items.type === 'object') {
+                result.items = convertParameterToGeminiFormat(param.items);
+            }
+            else {
+                result.items = {
+                    type: param.items.type === 'string' ? Type.STRING :
+                        param.items.type === 'number' ? Type.NUMBER :
+                            param.items.type === 'boolean' ? Type.BOOLEAN : Type.STRING
+                };
+                if (param.items.enum) {
+                    result.items.enum = param.items.enum;
+                }
+            }
+        }
+        else {
+            result.items = { type: Type.STRING };
+        }
+        if (param.enum) {
+            result.items.enum = param.enum;
+        }
+    }
+    else if (type === Type.OBJECT) {
+        if (param.properties && typeof param.properties === 'object') {
+            result.properties = {};
+            for (const [propName, propSchema] of Object.entries(param.properties)) {
+                result.properties[propName] = convertParameterToGeminiFormat(propSchema);
+            }
+        }
+        else {
+            result.properties = {};
+        }
+    }
+    else if (param.enum) {
+        result.format = 'enum';
+        result.enum = param.enum;
+    }
+    return result;
+}
+function convertToGeminiFunctionDeclarations(tools) {
+    return tools
+        .map(tool => {
+        if (tool.definition.function.name === 'google_web_search') {
+            console.log('[Gemini] Enabling Google Search grounding');
+            return null;
+        }
+        const properties = {};
+        const toolParams = tool.definition?.function?.parameters?.properties;
+        if (toolParams) {
+            for (const [name, param] of Object.entries(toolParams)) {
+                properties[name] = convertParameterToGeminiFormat(param);
+            }
+        }
+        else {
+            console.warn(`Tool ${tool.definition?.function?.name || 'Unnamed Tool'} has missing or invalid parameters definition.`);
+        }
+        return {
+            name: tool.definition.function.name,
+            description: tool.definition.function.description,
+            parameters: {
+                type: Type.OBJECT,
+                properties,
+                required: Array.isArray(tool.definition?.function?.parameters?.required)
+                    ? tool.definition.function.parameters.required
+                    : [],
+            },
+        };
+    })
+        .filter(Boolean);
+}
+function getImageMimeType(imageData) {
+    if (imageData.includes('data:image/jpeg'))
+        return 'image/jpeg';
+    if (imageData.includes('data:image/png'))
+        return 'image/png';
+    if (imageData.includes('data:image/gif'))
+        return 'image/gif';
+    if (imageData.includes('data:image/webp'))
+        return 'image/webp';
+    return 'image/jpeg';
+}
+function cleanBase64Data(imageData) {
+    return imageData.replace(/^data:image\/[a-z]+;base64,/, '');
+}
+function formatGroundingChunks(chunks) {
+    return chunks
+        .filter(c => c?.web?.uri)
+        .map((c, i) => `${i + 1}. ${c.web.title || 'Untitled'} – ${c.web.uri}`)
+        .join('\n');
+}
+async function convertToGeminiContents(messages) {
+    const contents = [];
+    for (const msg of messages) {
+        if (msg.type === 'function_call') {
+            let args = {};
+            try {
+                const parsedArgs = JSON.parse(msg.arguments || '{}');
+                args =
+                    typeof parsedArgs === 'object' && parsedArgs !== null
+                        ? parsedArgs
+                        : { value: parsedArgs };
+            }
+            catch (e) {
+                console.error(`Failed to parse function call arguments for ${msg.name}:`, msg.arguments, e);
+                args = {
+                    error: 'Invalid JSON arguments provided',
+                    raw_args: msg.arguments,
+                };
+            }
+            contents.push({
+                role: 'model',
+                parts: [
+                    {
+                        functionCall: {
+                            name: msg.name,
+                            args,
+                        },
+                    },
+                ],
+            });
+        }
+        else if (msg.type === 'function_call_output') {
+            if (typeof msg.output === 'string') {
+                const extracted = extractBase64Image(msg.output);
+                if (extracted.found && extracted.image_id !== null) {
+                    const image_id = extracted.image_id;
+                    const originalImageData = extracted.images[image_id];
+                    const processedImageData = await resizeAndTruncateForGemini(originalImageData);
+                    const mimeType = getImageMimeType(processedImageData);
+                    const cleanedImageData = cleanBase64Data(processedImageData);
+                    const parts = [];
+                    parts.push({
+                        functionResponse: {
+                            name: msg.name,
+                            response: {
+                                content: extracted.replaceContent.trim() ||
+                                    `[image #${image_id}]`,
+                            },
+                        },
+                    });
+                    parts.push({
+                        inlineData: {
+                            mimeType: mimeType,
+                            data: cleanedImageData,
+                        },
+                    });
+                    contents.push({
+                        role: 'user',
+                        parts: parts,
+                    });
+                }
+                else {
+                    contents.push({
+                        role: 'user',
+                        parts: [
+                            {
+                                functionResponse: {
+                                    name: msg.name,
+                                    response: { content: msg.output || '' },
+                                },
+                            },
+                        ],
+                    });
+                }
+            }
+            else {
+                contents.push({
+                    role: 'user',
+                    parts: [
+                        {
+                            functionResponse: {
+                                name: msg.name,
+                                response: { content: msg.output || '' },
+                            },
+                        },
+                    ],
+                });
+            }
+        }
+        else {
+            const role = msg.role === 'assistant' ? 'model' : 'user';
+            let textContent = '';
+            if (typeof msg.content === 'string') {
+                const extracted = extractBase64Image(msg.content);
+                if (extracted.found && extracted.image_id !== null) {
+                    const image_id = extracted.image_id;
+                    const originalImageData = extracted.images[image_id];
+                    const processedImageData = await resizeAndTruncateForGemini(originalImageData);
+                    const mimeType = getImageMimeType(processedImageData);
+                    const cleanedImageData = cleanBase64Data(processedImageData);
+                    const parts = [];
+                    if (extracted.replaceContent.trim()) {
+                        parts.push({
+                            text: extracted.replaceContent.trim(),
+                        });
+                    }
+                    parts.push({
+                        inlineData: {
+                            mimeType: mimeType,
+                            data: cleanedImageData,
+                        },
+                    });
+                    contents.push({
+                        role,
+                        parts: parts,
+                    });
+                    continue;
+                }
+                else {
+                    textContent = msg.content;
+                }
+            }
+            else if (msg.content &&
+                typeof msg.content === 'object' &&
+                'text' in msg.content) {
+                textContent = msg.content.text;
+            }
+            if (textContent && textContent.trim() !== '') {
+                if (msg.type === 'thinking') {
+                    textContent = 'Thinking: ' + textContent;
+                }
+                contents.push({
+                    role,
+                    parts: [{ text: textContent.trim() }],
+                });
+            }
+        }
+    }
+    return contents;
+}
+const THINKING_BUDGET_CONFIGS = {
+    '-low': 0,
+    '-medium': 2048,
+    '-high': 12288,
+    '-max': 24576,
+};
+export class GeminiProvider {
+    client;
+    constructor(apiKey) {
+        const key = apiKey || process.env.GOOGLE_API_KEY;
+        if (!key) {
+            throw new Error('Failed to initialize Gemini client. GOOGLE_API_KEY is missing or not provided.');
+        }
+        this.client = new GoogleGenAI({
+            apiKey: key,
+            vertexai: false,
+        });
+    }
+    async createEmbedding(modelId, input, opts) {
+        try {
+            let actualModelId = modelId.startsWith('gemini/')
+                ? modelId.substring(7)
+                : modelId;
+            let thinkingConfig = null;
+            for (const [suffix, budget] of Object.entries(THINKING_BUDGET_CONFIGS)) {
+                if (actualModelId.endsWith(suffix)) {
+                    thinkingConfig = { thinkingBudget: budget };
+                    actualModelId = actualModelId.slice(0, -suffix.length);
+                    break;
+                }
+            }
+            console.log(`[Gemini] Generating embedding with model ${actualModelId}`);
+            const payload = {
+                model: actualModelId,
+                contents: input,
+                config: {
+                    taskType: opts?.taskType ?? 'SEMANTIC_SIMILARITY',
+                },
+            };
+            if (thinkingConfig) {
+                payload.config.thinkingConfig = thinkingConfig;
+            }
+            const response = await this.client.models.embedContent(payload);
+            console.log('[Gemini] Embedding response structure:', JSON.stringify(response, (key, value) => key === 'values' &&
+                Array.isArray(value) &&
+                value.length > 10
+                ? `[${value.length} items]`
+                : value, 2));
+            if (!response.embeddings || !Array.isArray(response.embeddings)) {
+                console.error('[Gemini] Unexpected embedding response structure:', response);
+                throw new Error('Invalid embedding response structure from Gemini API');
+            }
+            const estimatedTokens = typeof input === 'string'
+                ? Math.ceil(input.length / 4)
+                : input.reduce((sum, text) => sum + Math.ceil(text.length / 4), 0);
+            let extractedValues = [];
+            let dimensions = 0;
+            if (response.embeddings.length > 0) {
+                if (response.embeddings[0].values) {
+                    extractedValues = response.embeddings.map(e => e.values);
+                    dimensions = extractedValues[0].length;
+                }
+                else {
+                    console.warn('[Gemini] Could not find expected "values" property in embeddings response');
+                    extractedValues =
+                        response.embeddings;
+                    dimensions = Array.isArray(extractedValues[0])
+                        ? extractedValues[0].length
+                        : 0;
+                }
+            }
+            costTracker.addUsage({
+                model: modelId,
+                input_tokens: estimatedTokens,
+                output_tokens: 0,
+                metadata: {
+                    dimensions,
+                },
+            });
+            if (Array.isArray(input) && input.length > 1) {
+                return extractedValues;
+            }
+            else {
+                let result;
+                if (Array.isArray(extractedValues) &&
+                    extractedValues.length >= 1) {
+                    const firstValue = extractedValues[0];
+                    if (Array.isArray(firstValue)) {
+                        result = firstValue;
+                    }
+                    else {
+                        console.error('[Gemini] Unexpected format in embedding result:', firstValue);
+                        result = [];
+                    }
+                }
+                else {
+                    result = [];
+                }
+                let adjustedResult = result;
+                if (result.length !== 3072) {
+                    console.warn(`Gemini embedding returned ${result.length} dimensions, adjusting to 3072...`);
+                    if (result.length > 3072) {
+                        adjustedResult = result.slice(0, 3072);
+                    }
+                    else {
+                        adjustedResult = [
+                            ...result,
+                            ...Array(3072 - result.length).fill(0),
+                        ];
+                    }
+                }
+                return adjustedResult;
+            }
+        }
+        catch (error) {
+            console.error('[Gemini] Error generating embedding:', error);
+            throw error;
+        }
+    }
+    async *retryStreamOnIncompleteJson(requestFn, maxRetries = 2) {
+        let attempts = 0;
+        while (attempts <= maxRetries) {
+            try {
+                const stream = await requestFn();
+                for await (const chunk of stream) {
+                    yield chunk;
+                }
+                return;
+            }
+            catch (error) {
+                attempts++;
+                const errorMsg = error instanceof Error ? error.message : String(error);
+                if (errorMsg.includes('Incomplete JSON segment') &&
+                    attempts <= maxRetries) {
+                    console.warn(`[Gemini] Incomplete JSON segment error, retrying (${attempts}/${maxRetries})...`);
+                    await new Promise(resolve => setTimeout(resolve, 1000 * attempts));
+                    continue;
+                }
+                throw error;
+            }
+        }
+    }
+    async *createResponseStream(model, messages, agent) {
+        const tools = agent
+            ? await agent.getTools()
+            : [];
+        const settings = agent?.modelSettings;
+        let contentBuffer = '';
+        const messageId = uuidv4();
+        let eventOrder = 0;
+        const deltaBuffers = new Map();
+        let hasYieldedToolStart = false;
+        const shownGrounding = new Set();
+        let requestId = undefined;
+        const chunks = [];
+        try {
+            const contents = await convertToGeminiContents(messages);
+            if (contents.length === 0) {
+                throw new Error('No valid content found in messages after conversion.');
+            }
+            const lastContent = contents[contents.length - 1];
+            if (lastContent.role !== 'user') {
+                console.warn("Last message in history is not from 'user'. Gemini might not respond as expected.");
+            }
+            let thinkingBudget = null;
+            for (const [suffix, budget] of Object.entries(THINKING_BUDGET_CONFIGS)) {
+                if (model.endsWith(suffix)) {
+                    thinkingBudget = budget;
+                    model = model.slice(0, -suffix.length);
+                    break;
+                }
+            }
+            const config = {
+                thinkingConfig: {
+                    includeThoughts: true,
+                },
+            };
+            if (thinkingBudget) {
+                config.thinkingConfig.thinkingBudget = thinkingBudget;
+            }
+            if (settings?.stop_sequence) {
+                config.stopSequences = [settings.stop_sequence];
+            }
+            if (settings?.temperature) {
+                config.temperature = settings.temperature;
+            }
+            if (settings?.max_tokens) {
+                config.maxOutputTokens = settings.max_tokens;
+            }
+            if (settings?.top_p) {
+                config.topP = settings.top_p;
+            }
+            if (settings?.top_k) {
+                config.topK = settings.top_k;
+            }
+            if (settings?.json_schema) {
+                config.responseMimeType = 'application/json';
+                config.responseSchema = settings.json_schema.schema;
+                if (config.responseSchema) {
+                    const removeAdditionalProperties = (obj) => {
+                        if (!obj || typeof obj !== 'object') {
+                            return;
+                        }
+                        if ('additionalProperties' in obj) {
+                            delete obj.additionalProperties;
+                        }
+                        if (obj.properties &&
+                            typeof obj.properties === 'object') {
+                            Object.values(obj.properties).forEach(prop => {
+                                removeAdditionalProperties(prop);
+                            });
+                        }
+                        if (obj.items) {
+                            removeAdditionalProperties(obj.items);
+                        }
+                        ['oneOf', 'anyOf', 'allOf'].forEach(key => {
+                            if (obj[key] && Array.isArray(obj[key])) {
+                                obj[key].forEach((subSchema) => {
+                                    removeAdditionalProperties(subSchema);
+                                });
+                            }
+                        });
+                    };
+                    removeAdditionalProperties(config.responseSchema);
+                }
+            }
+            let hasGoogleWebSearch = false;
+            if (tools && tools.length > 0) {
+                hasGoogleWebSearch = tools.some(tool => tool.definition.function.name === 'google_web_search');
+                const functionDeclarations = convertToGeminiFunctionDeclarations(tools);
+                let allowedFunctionNames = [];
+                if (functionDeclarations.length > 0) {
+                    config.tools = [{ functionDeclarations }];
+                    if (settings?.tool_choice) {
+                        let toolChoice;
+                        if (typeof settings.tool_choice === 'object' &&
+                            settings.tool_choice?.type === 'function' &&
+                            settings.tool_choice?.function?.name) {
+                            toolChoice = FunctionCallingConfigMode.ANY;
+                            allowedFunctionNames = [
+                                settings.tool_choice.function.name,
+                            ];
+                        }
+                        else if (settings.tool_choice === 'required') {
+                            toolChoice = FunctionCallingConfigMode.ANY;
+                        }
+                        else if (settings.tool_choice === 'auto') {
+                            toolChoice = FunctionCallingConfigMode.AUTO;
+                        }
+                        else if (settings.tool_choice === 'none') {
+                            toolChoice = FunctionCallingConfigMode.NONE;
+                        }
+                        if (toolChoice) {
+                            config.toolConfig = {
+                                functionCallingConfig: {
+                                    mode: toolChoice,
+                                },
+                            };
+                            if (allowedFunctionNames.length > 0) {
+                                config.toolConfig.functionCallingConfig.allowedFunctionNames =
+                                    allowedFunctionNames;
+                            }
+                        }
+                    }
+                }
+                else if (!hasGoogleWebSearch) {
+                    console.warn('Tools were provided but resulted in empty declarations after conversion.');
+                }
+            }
+            if (hasGoogleWebSearch) {
+                console.log('[Gemini] Enabling Google Search grounding');
+                config.tools = [{ googleSearch: {} }];
+                config.toolConfig = {
+                    functionCallingConfig: {
+                        mode: FunctionCallingConfigMode.ANY,
+                        allowedFunctionNames: ['googleSearch'],
+                    },
+                };
+            }
+            const requestParams = {
+                model,
+                contents,
+                config,
+            };
+            requestId = log_llm_request(agent.agent_id, 'google', model, requestParams);
+            const getStreamFn = () => this.client.models.generateContentStream(requestParams);
+            const response = this.retryStreamOnIncompleteJson(getStreamFn);
+            let usageMetadata;
+            for await (const chunk of response) {
+                chunks.push(chunk);
+                if (isPaused()) {
+                    console.log(`[Gemini] System paused during stream for model ${model}. Aborting processing.`);
+                    yield {
+                        type: 'message_delta',
+                        content: '\n⏸️ Stream paused by user.',
+                        message_id: messageId,
+                        order: 999,
+                    };
+                    break;
+                }
+                if (chunk.functionCalls && chunk.functionCalls.length > 0) {
+                    const toolCallsToEmit = [];
+                    for (const fc of chunk.functionCalls) {
+                        if (fc && fc.name) {
+                            const callId = `call_${uuidv4()}`;
+                            toolCallsToEmit.push({
+                                id: callId,
+                                type: 'function',
+                                function: {
+                                    name: fc.name,
+                                    arguments: JSON.stringify(fc.args || {}),
+                                },
+                            });
+                        }
+                    }
+                    if (toolCallsToEmit.length > 0 && !hasYieldedToolStart) {
+                        yield {
+                            type: 'tool_start',
+                            tool_calls: toolCallsToEmit,
+                        };
+                        hasYieldedToolStart = true;
+                        continue;
+                    }
+                }
+                if (chunk.text) {
+                    contentBuffer += chunk.text;
+                    for (const ev of bufferDelta(deltaBuffers, messageId, chunk.text, content => ({
+                        type: 'message_delta',
+                        content,
+                        message_id: messageId,
+                        order: eventOrder++,
+                    }))) {
+                        yield ev;
+                    }
+                }
+                const gChunks = chunk.candidates?.[0]?.groundingMetadata?.groundingChunks;
+                if (Array.isArray(gChunks)) {
+                    const newChunks = gChunks.filter(c => c?.web?.uri && !shownGrounding.has(c.web.uri));
+                    if (newChunks.length) {
+                        newChunks.forEach(c => shownGrounding.add(c.web.uri));
+                        const formatted = formatGroundingChunks(newChunks);
+                        yield {
+                            type: 'message_delta',
+                            content: '\n\nSearch Results:\n' + formatted + '\n',
+                            message_id: messageId,
+                            order: eventOrder++,
+                        };
+                        contentBuffer +=
+                            '\n\nSearch Results:\n' + formatted + '\n';
+                    }
+                }
+                if (chunk.candidates?.[0]?.content?.parts) {
+                    for (const part of chunk.candidates[0].content.parts) {
+                        if (part.inlineData?.data) {
+                            yield {
+                                type: 'file_complete',
+                                data_format: 'base64',
+                                data: part.inlineData.data,
+                                mime_type: part.inlineData.mimeType || 'image/png',
+                                message_id: uuidv4(),
+                                order: eventOrder++,
+                            };
+                        }
+                    }
+                }
+                if (chunk.usageMetadata) {
+                    usageMetadata = chunk.usageMetadata;
+                }
+            }
+            for (const ev of flushBufferedDeltas(deltaBuffers, (_id, content) => ({
+                type: 'message_delta',
+                content,
+                message_id: messageId,
+                order: eventOrder++,
+            }))) {
+                yield ev;
+            }
+            if (usageMetadata) {
+                costTracker.addUsage({
+                    model,
+                    input_tokens: usageMetadata.promptTokenCount || 0,
+                    output_tokens: usageMetadata.candidatesTokenCount || 0,
+                    cached_tokens: usageMetadata.cachedContentTokenCount || 0,
+                    metadata: {
+                        total_tokens: usageMetadata.totalTokenCount || 0,
+                        reasoning_tokens: usageMetadata.thoughtsTokenCount || 0,
+                        tool_tokens: usageMetadata.toolUsePromptTokenCount || 0,
+                    },
+                });
+            }
+            else {
+                console.error('No usage metadata found in the response. This may affect token tracking.');
+                costTracker.addUsage({
+                    model,
+                    input_tokens: 0,
+                    output_tokens: 0,
+                    cached_tokens: 0,
+                    metadata: {
+                        total_tokens: 0,
+                        source: 'estimated',
+                    },
+                });
+            }
+            if (!hasYieldedToolStart && contentBuffer) {
+                yield {
+                    type: 'message_complete',
+                    content: contentBuffer,
+                    message_id: messageId,
+                };
+            }
+        }
+        catch (error) {
+            log_llm_error(requestId, error);
+            const errorMessage = error instanceof Error
+                ? error.stack || error.message
+                : String(error);
+            if (errorMessage.includes('Incomplete JSON segment')) {
+                console.error('[Gemini] Stream terminated with incomplete JSON. This may indicate network issues or timeouts.');
+            }
+            console.error('\n=== Gemini error ===');
+            console.dir(error, { depth: null });
+            console.error('\n=== JSON dump of error ===');
+            console.error(JSON.stringify(error, Object.getOwnPropertyNames(error), 2));
+            console.error('\n=== Manual property walk ===');
+            for (const key of Reflect.ownKeys(error)) {
+                console.error(`${String(key)}:`, error[key]);
+            }
+            yield {
+                type: 'error',
+                error: `Gemini error ${model}: ${errorMessage}`,
+            };
+            if (!hasYieldedToolStart && contentBuffer) {
+                yield {
+                    type: 'message_complete',
+                    content: contentBuffer,
+                    message_id: messageId,
+                };
+            }
+        }
+        finally {
+            log_llm_response(requestId, chunks);
+        }
+    }
+}
+export const geminiProvider = new GeminiProvider();
+//# sourceMappingURL=gemini.js.map