130 lines
3.9 KiB
TypeScript
130 lines
3.9 KiB
TypeScript
import { createGoogleGenerativeAI } from '@ai-sdk/google';
|
|
import { generateObject, generateText as aiGenerateText, embed, stepCountIs } from 'ai';
|
|
import { z } from 'zod';
|
|
import { AIProvider, TagSuggestion, TitleSuggestion, ToolUseOptions, ToolCallResult } from '../types';
|
|
|
|
export class GoogleProvider implements AIProvider {
|
|
private model: any;
|
|
private embeddingModel: any;
|
|
|
|
constructor(apiKey: string, modelName: string = 'gemini-1.5-flash', embeddingModelName: string = 'text-embedding-004') {
|
|
const google = createGoogleGenerativeAI({
|
|
apiKey: apiKey,
|
|
});
|
|
|
|
this.model = google(modelName);
|
|
this.embeddingModel = google.textEmbeddingModel(embeddingModelName);
|
|
}
|
|
|
|
async generateTags(content: string): Promise<TagSuggestion[]> {
|
|
try {
|
|
const { object } = await generateObject({
|
|
model: this.model,
|
|
schema: z.object({
|
|
tags: z.array(z.object({
|
|
tag: z.string().describe('Short tag name in lowercase'),
|
|
confidence: z.number().min(0).max(1).describe('Confidence level between 0 and 1')
|
|
}))
|
|
}),
|
|
prompt: `Analyze the following note and suggest 1 to 5 relevant tags.
|
|
Note content: "${content}"`,
|
|
});
|
|
|
|
return object.tags;
|
|
} catch (e) {
|
|
console.error('Error generating tags (Google):', e);
|
|
return [];
|
|
}
|
|
}
|
|
|
|
async getEmbeddings(text: string): Promise<number[]> {
|
|
try {
|
|
const { embedding } = await embed({
|
|
model: this.embeddingModel,
|
|
value: text,
|
|
});
|
|
return embedding;
|
|
} catch (e) {
|
|
console.error('Error generating embeddings (Google):', e);
|
|
throw e;
|
|
}
|
|
}
|
|
|
|
async generateTitles(prompt: string): Promise<TitleSuggestion[]> {
|
|
try {
|
|
const { object } = await generateObject({
|
|
model: this.model,
|
|
schema: z.object({
|
|
titles: z.array(z.object({
|
|
title: z.string().describe('Suggested title'),
|
|
confidence: z.number().min(0).max(1).describe('Confidence level between 0 and 1')
|
|
}))
|
|
}),
|
|
prompt: prompt,
|
|
});
|
|
|
|
return object.titles;
|
|
} catch (e) {
|
|
console.error('Error generating titles (Google):', e);
|
|
return [];
|
|
}
|
|
}
|
|
|
|
async generateText(prompt: string): Promise<string> {
|
|
try {
|
|
const { text } = await aiGenerateText({
|
|
model: this.model,
|
|
prompt: prompt,
|
|
});
|
|
|
|
return text.trim();
|
|
} catch (e) {
|
|
console.error('Error generating text (Google):', e);
|
|
throw e;
|
|
}
|
|
}
|
|
|
|
async chat(messages: any[], systemPrompt?: string): Promise<any> {
|
|
try {
|
|
const { text } = await aiGenerateText({
|
|
model: this.model,
|
|
system: systemPrompt,
|
|
messages: messages,
|
|
});
|
|
|
|
return { text: text.trim() };
|
|
} catch (e) {
|
|
console.error('Error in chat (Google):', e);
|
|
throw e;
|
|
}
|
|
}
|
|
|
|
async generateWithTools(options: ToolUseOptions): Promise<ToolCallResult> {
|
|
const { tools, maxSteps = 10, systemPrompt, messages, prompt } = options
|
|
const opts: Record<string, any> = {
|
|
model: this.model,
|
|
tools,
|
|
stopWhen: stepCountIs(maxSteps),
|
|
}
|
|
if (systemPrompt) opts.system = systemPrompt
|
|
if (messages) opts.messages = messages
|
|
else if (prompt) opts.prompt = prompt
|
|
|
|
const result = await aiGenerateText(opts as any)
|
|
return {
|
|
toolCalls: result.toolCalls?.map((tc: any) => ({ toolName: tc.toolName, input: tc.input })) || [],
|
|
toolResults: result.toolResults?.map((tr: any) => ({ toolName: tr.toolName, input: tr.input, output: tr.output })) || [],
|
|
text: result.text,
|
|
steps: result.steps?.map((step: any) => ({
|
|
text: step.text,
|
|
toolCalls: step.toolCalls?.map((tc: any) => ({ toolName: tc.toolName, input: tc.input })) || [],
|
|
toolResults: step.toolResults?.map((tr: any) => ({ toolName: tr.toolName, input: tr.input, output: tr.output })) || []
|
|
})) || []
|
|
}
|
|
}
|
|
|
|
getModel() {
|
|
return this.model;
|
|
}
|
|
}
|