pit-manager
Version:
Centralized prompt management system for Human Behavior AI agents
249 lines (248 loc) • 11.4 kB
JavaScript
;
/**
* Google provider implementation for TypeScript.
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.GoogleProvider = void 0;
const generative_ai_1 = require("@google/generative-ai");
const exceptions_1 = require("../core/exceptions");
const base_1 = require("./base");
class GoogleProvider extends base_1.BaseProvider {
constructor(apiKey, config = {}) {
super(apiKey, config);
if (!this.apiKey) {
throw new exceptions_1.APIKeyError('Google API key is required');
}
this.genAI = new generative_ai_1.GoogleGenerativeAI(this.apiKey);
}
get name() {
return 'google';
}
get supportedModels() {
return Object.keys(GoogleProvider.PRICING);
}
async chatCompletion(model, messages, options = {}) {
if (!this.supportsModel(model)) {
const available = this.supportedModels.join(', ');
throw new exceptions_1.ProviderError(`Model ${model} not supported by Google. Available: ${available}`);
}
try {
// Initialize generation config
const generationConfig = {};
if (options.temperature !== undefined) {
generationConfig.temperature = options.temperature;
}
if (options.maxTokens !== undefined) {
generationConfig.maxOutputTokens = options.maxTokens;
}
// Safety settings - moderate default
const safetySettings = [
{
category: generative_ai_1.HarmCategory.HARM_CATEGORY_HARASSMENT,
threshold: generative_ai_1.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
},
{
category: generative_ai_1.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
threshold: generative_ai_1.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
},
{
category: generative_ai_1.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
threshold: generative_ai_1.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
},
{
category: generative_ai_1.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
threshold: generative_ai_1.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
},
];
// Handle structured output configuration
let modelConfig = {
model,
generationConfig,
safetySettings
};
// Add structured output support
if (options.responseMimeType && options.responseSchema) {
modelConfig = {
...modelConfig,
generationConfig: {
...generationConfig,
responseMimeType: options.responseMimeType,
responseSchema: options.responseSchema
}
};
}
// Extract system instruction from options if present
if (options.systemPrompt) {
modelConfig.systemInstruction = options.systemPrompt;
}
const geminiModel = this.genAI.getGenerativeModel(modelConfig);
// Convert PIT messages to Gemini format
// Gemini chat expects alternating user/model messages
const geminiHistory = [];
let systemInstruction;
// Extract system message if present
const currentMessages = [];
for (const msg of messages) {
if (msg.role === 'system') {
systemInstruction = typeof msg.content === 'string' ? msg.content : msg.content.text || '';
}
else {
currentMessages.push(msg);
}
}
// Build conversation history (all but last message)
if (currentMessages.length > 1) {
for (let i = 0; i < currentMessages.length - 1; i++) {
const msg = currentMessages[i];
const role = msg.role === 'user' ? 'user' : 'model';
// Handle multimodal content
if (typeof msg.content === 'string') {
geminiHistory.push({
role,
parts: [{ text: msg.content }]
});
}
else if (msg.content.type === 'multimodal' && msg.content.parts) {
const parts = msg.content.parts.map(part => {
if (part.type === 'text') {
return { text: part.text || '' };
}
else if (part.type === 'image' && part.image) {
return {
inlineData: {
mimeType: part.image.mimeType,
data: part.image.data
}
};
}
return null;
}).filter(part => part !== null);
geminiHistory.push({ role, parts });
}
}
}
// Get the final message to send
let finalParts = [];
if (currentMessages.length > 0) {
const lastMsg = currentMessages[currentMessages.length - 1];
if (typeof lastMsg.content === 'string') {
// Prepend system instruction if present
const text = systemInstruction ? `${systemInstruction}\n\n${lastMsg.content}` : lastMsg.content;
finalParts = [{ text }];
}
else if (lastMsg.content.type === 'multimodal' && lastMsg.content.parts) {
// Handle multimodal last message
finalParts = lastMsg.content.parts.map((part, index) => {
if (part.type === 'text') {
// Prepend system instruction to first text part if present
const text = (index === 0 && systemInstruction)
? `${systemInstruction}\n\n${part.text || ''}`
: (part.text || '');
return { text };
}
else if (part.type === 'image' && part.image) {
return {
inlineData: {
mimeType: part.image.mimeType,
data: part.image.data
}
};
}
return null;
}).filter(part => part !== null);
// If no text parts exist but we have system instruction, add it
if (systemInstruction && !finalParts.some(p => p.text)) {
finalParts.unshift({ text: systemInstruction });
}
}
}
// Start chat session
const chat = geminiModel.startChat({
history: geminiHistory
});
// Make API call
const startTime = new Date();
const result = await chat.sendMessage(finalParts);
const endTime = new Date();
// Calculate latency
const latency = endTime.getTime() - startTime.getTime();
// Extract response data
const response = await result.response;
const content = response.text();
// Extract usage (Gemini provides token counts)
const usageMetadata = response.usageMetadata;
const tokenUsage = {
input: usageMetadata?.promptTokenCount || 0,
output: usageMetadata?.candidatesTokenCount || 0,
total: usageMetadata?.totalTokenCount || 0
};
// Calculate cost
const cost = this.calculateCost(model, tokenUsage);
// Determine finish reason
let finishReason = 'stop';
if (response.candidates && response.candidates[0]?.finishReason) {
finishReason = response.candidates[0].finishReason.toLowerCase();
}
return {
content,
tokens: tokenUsage,
model,
provider: this.name,
latency,
cost,
timestamp: startTime,
finishReason,
rawResponse: {
text: content,
usageMetadata: {
promptTokenCount: usageMetadata?.promptTokenCount || 0,
candidatesTokenCount: usageMetadata?.candidatesTokenCount || 0,
totalTokenCount: usageMetadata?.totalTokenCount || 0
},
candidates: response.candidates?.map(candidate => ({
finishReason: candidate.finishReason || null,
safetyRatings: candidate.safetyRatings?.map(rating => ({
category: rating.category,
probability: rating.probability
})) || []
})) || []
}
};
}
catch (error) {
// Google API doesn't have specific rate limit errors, check message
if (error.message?.toLowerCase().includes('quota') || error.message?.toLowerCase().includes('rate')) {
throw new exceptions_1.RateLimitError(`Google rate limit exceeded: ${error.message}`);
}
throw new exceptions_1.ProviderError(`Unexpected error calling Google: ${error.message}`);
}
}
calculateCost(model, tokens) {
const pricing = GoogleProvider.PRICING[model];
if (!pricing) {
return 0.0;
}
const inputCost = (tokens.input / 1000) * pricing.input;
const outputCost = (tokens.output / 1000) * pricing.output;
return inputCost + outputCost;
}
}
exports.GoogleProvider = GoogleProvider;
// Google Gemini pricing per 1K tokens (updated with latest models)
GoogleProvider.PRICING = {
// Legacy models
'gemini-1.0-pro': { input: 0.0005, output: 0.0015 },
'gemini-pro': { input: 0.0005, output: 0.0015 }, // Alias for gemini-1.0-pro
'gemini-pro-vision': { input: 0.00025, output: 0.0005 },
// Gemini 1.5 series
'gemini-1.5-pro': { input: 0.00125, output: 0.005 },
'gemini-1.5-flash': { input: 0.000075, output: 0.0003 },
// Gemini 2.0 series (estimated pricing - adjust as needed)
'gemini-2.0-flash': { input: 0.0001, output: 0.0004 },
'gemini-2.0-flash-lite': { input: 0.00005, output: 0.0002 },
// Gemini 2.5 series (estimated pricing - adjust as needed)
'gemini-2.5-pro': { input: 0.002, output: 0.008 },
'gemini-2.5-flash': { input: 0.0001, output: 0.0004 },
'gemini-2.5-flash-lite': { input: 0.00005, output: 0.0002 },
'gemini-live-2.5-flash-preview': { input: 0.0001, output: 0.0004 }
};