79 lines
2.4 KiB
TypeScript
79 lines
2.4 KiB
TypeScript
import { GoogleGenAI } from '@google/genai';
|
|
import { IPromptAgent, PromptEnhancementOptions, AgentResult } from '../types';
|
|
import { detectLanguage, detectEnhancements } from '../utils';
|
|
|
|
export abstract class BaseAgent implements IPromptAgent {
|
|
protected ai: GoogleGenAI;
|
|
protected model = 'gemini-2.5-flash';
|
|
|
|
abstract readonly templateType: string;
|
|
|
|
constructor(apiKey: string) {
|
|
if (!apiKey) {
|
|
throw new Error('Gemini API key is required');
|
|
}
|
|
this.ai = new GoogleGenAI({ apiKey });
|
|
}
|
|
|
|
protected abstract getSystemPrompt(): string;
|
|
protected abstract getTemplate(): string;
|
|
|
|
async enhance(rawPrompt: string, _options: PromptEnhancementOptions): Promise<AgentResult> {
|
|
const timestamp = new Date().toISOString();
|
|
|
|
console.log(
|
|
`[${timestamp}] [${this.templateType}Agent] Enhancing prompt: "${rawPrompt.substring(0, 50)}..."`,
|
|
);
|
|
|
|
try {
|
|
const systemPrompt = this.getSystemPrompt();
|
|
const userPrompt = this.buildUserPrompt(rawPrompt);
|
|
|
|
const response = await this.ai.models.generateContent({
|
|
model: this.model,
|
|
config: { responseModalities: ['TEXT'] },
|
|
contents: [
|
|
{
|
|
role: 'user' as const,
|
|
parts: [{ text: `${systemPrompt}\n\n${userPrompt}` }],
|
|
},
|
|
],
|
|
});
|
|
|
|
if (response.candidates && response.candidates[0] && response.candidates[0].content) {
|
|
const content = response.candidates[0].content;
|
|
const enhancedPrompt = content.parts?.[0]?.text?.trim() || '';
|
|
|
|
console.log(`[${timestamp}] [${this.templateType}Agent] Enhancement successful`);
|
|
|
|
return {
|
|
success: true,
|
|
enhancedPrompt,
|
|
detectedLanguage: detectLanguage(rawPrompt),
|
|
appliedTemplate: this.templateType,
|
|
enhancements: detectEnhancements(rawPrompt, enhancedPrompt),
|
|
};
|
|
}
|
|
|
|
return {
|
|
success: false,
|
|
error: 'No enhanced prompt received from API',
|
|
enhancements: [],
|
|
};
|
|
} catch (error) {
|
|
console.error(`[${timestamp}] [${this.templateType}Agent] Enhancement failed:`, error);
|
|
return {
|
|
success: false,
|
|
error: error instanceof Error ? error.message : 'Enhancement failed',
|
|
enhancements: [],
|
|
};
|
|
}
|
|
}
|
|
|
|
protected buildUserPrompt(rawPrompt: string): string {
|
|
return `Transform this prompt into a professional image generation prompt: "${rawPrompt}"
|
|
|
|
Target template/style: ${this.templateType}`;
|
|
}
|
|
}
|