diff --git a/apps/web/app/api/user/settings/route.ts b/apps/web/app/api/user/settings/route.ts index 10edcdc0ca..c02156835b 100644 --- a/apps/web/app/api/user/settings/route.ts +++ b/apps/web/app/api/user/settings/route.ts @@ -28,6 +28,10 @@ async function saveAISettings(options: SaveSettingsBody) { } // use bedrock if no api key set return Model.CLAUDE_3_5_SONNET_BEDROCK; + case Provider.GOOGLE: + return options.aiModel || Model.GEMINI_1_5_PRO; + case Provider.OLLAMA: + return Model.OLLAMA; default: throw new Error("Invalid AI provider"); } diff --git a/apps/web/app/api/user/settings/validation.ts b/apps/web/app/api/user/settings/validation.ts index c0e53c6ce3..2760b59115 100644 --- a/apps/web/app/api/user/settings/validation.ts +++ b/apps/web/app/api/user/settings/validation.ts @@ -3,7 +3,12 @@ import { z } from "zod"; export const saveSettingsBody = z .object({ - aiProvider: z.enum([Provider.ANTHROPIC, Provider.OPEN_AI]), + aiProvider: z.enum([ + Provider.ANTHROPIC, + Provider.OPEN_AI, + Provider.GOOGLE, + ...(Provider.OLLAMA ? [Provider.OLLAMA] : []), + ]), aiModel: z.string(), aiApiKey: z.string().optional(), }) diff --git a/apps/web/package.json b/apps/web/package.json index e4627e6e37..f1e9e42f0e 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -15,6 +15,7 @@ "dependencies": { "@ai-sdk/amazon-bedrock": "^1.0.6", "@ai-sdk/anthropic": "^1.0.6", + "@ai-sdk/google": "^1.0.12", "@ai-sdk/openai": "^1.0.11", "@asteasolutions/zod-to-openapi": "^7.3.0", "@auth/core": "^0.37.4", diff --git a/apps/web/utils/ai/rule/generate-rules-prompt.ts b/apps/web/utils/ai/rule/generate-rules-prompt.ts index 9258f83879..42ccc03906 100644 --- a/apps/web/utils/ai/rule/generate-rules-prompt.ts +++ b/apps/web/utils/ai/rule/generate-rules-prompt.ts @@ -123,7 +123,7 @@ Your response should only include the list of general rules. Aim for 3-10 broadl const args = aiResponse.toolCalls[0].args; - logger.trace(args); + logger.trace("Args", { args }); return parseRulesResponse(args, hasSnippets); } diff --git a/apps/web/utils/llms/config.ts b/apps/web/utils/llms/config.ts index ddf5ef9cf7..49cfaf7601 100644 --- a/apps/web/utils/llms/config.ts +++ b/apps/web/utils/llms/config.ts @@ -5,6 +5,7 @@ const supportsOllama = env.NEXT_PUBLIC_OLLAMA_MODEL; export const Provider = { OPEN_AI: "openai", ANTHROPIC: "anthropic", + GOOGLE: "google", ...(supportsOllama ? { OLLAMA: "ollama" } : {}), }; @@ -15,12 +16,15 @@ export const Model = { // BEDROCK_ANTHROPIC_BACKUP_MODEL: // env.NEXT_PUBLIC_BEDROCK_ANTHROPIC_BACKUP_MODEL, CLAUDE_3_5_SONNET_ANTHROPIC: "claude-3-5-sonnet-20241022", + GEMINI_1_5_PRO: "gemini-1.5-pro-latest", + GEMINI_1_5_FLASH: "gemini-1.5-flash-latest", ...(supportsOllama ? { OLLAMA: env.NEXT_PUBLIC_OLLAMA_MODEL } : {}), }; export const providerOptions: { label: string; value: string }[] = [ { label: "OpenAI", value: Provider.OPEN_AI }, { label: "Anthropic", value: Provider.ANTHROPIC }, + { label: "Google", value: Provider.GOOGLE }, ...(supportsOllama && Provider.OLLAMA ? [{ label: "Ollama", value: Provider.OLLAMA }] : []), @@ -38,6 +42,16 @@ export const modelOptions: Record = value: "claude-3-5-sonnet", // used in ui only. can be either anthropic or bedrock }, ], + [Provider.GOOGLE]: [ + { + label: "Gemini 1.5 Pro", + value: Model.GEMINI_1_5_PRO, + }, + { + label: "Gemini 1.5 Flash", + value: Model.GEMINI_1_5_FLASH, + }, + ], ...(Provider.OLLAMA && Model.OLLAMA ? { [Provider.OLLAMA]: [{ label: "Ollama", value: Model.OLLAMA }], diff --git a/apps/web/utils/llms/index.ts b/apps/web/utils/llms/index.ts index a34ff5c99a..807bf399bb 100644 --- a/apps/web/utils/llms/index.ts +++ b/apps/web/utils/llms/index.ts @@ -10,6 +10,7 @@ import { import { createOpenAI } from "@ai-sdk/openai"; import { createAnthropic } from "@ai-sdk/anthropic"; import { createAmazonBedrock } from "@ai-sdk/amazon-bedrock"; +import { createGoogleGenerativeAI } from "@ai-sdk/google"; import { createOllama } from "ollama-ai-provider"; import { env } from "@/env"; import { saveAiUsage } from "@/utils/usage"; @@ -70,6 +71,17 @@ function getModel({ aiProvider, aiModel, aiApiKey }: UserAIFields) { }; } + if (provider === Provider.GOOGLE) { + if (!aiApiKey) throw new Error("Google API key is not set"); + + const model = aiModel || Model.GEMINI_1_5_PRO; + return { + provider: Provider.GOOGLE, + model, + llmModel: createGoogleGenerativeAI({ apiKey: aiApiKey })(model), + }; + } + if (provider === Provider.OLLAMA && env.NEXT_PUBLIC_OLLAMA_MODEL) { return { provider: Provider.OLLAMA, diff --git a/apps/web/utils/usage.ts b/apps/web/utils/usage.ts index 9e27215dfc..8ed0e645b6 100644 --- a/apps/web/utils/usage.ts +++ b/apps/web/utils/usage.ts @@ -91,6 +91,16 @@ const costs: Record< input: 0.8 / 1_000_000, output: 4 / 1_000_000, }, + // https://ai.google.dev/pricing#1_5pro + "gemini-1.5-pro-latest": { + input: 1.25 / 1_000_000, + output: 5 / 1_000_000, + }, + // https://ai.google.dev/pricing#1_5flash + "gemini-1.5-flash-latest": { + input: 0.075 / 1_000_000, + output: 0.3 / 1_000_000, + }, }; // returns cost in cents diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index c6386564c0..dbd3b74818 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -90,6 +90,9 @@ importers: '@ai-sdk/anthropic': specifier: ^1.0.6 version: 1.0.6(zod@3.24.1) + '@ai-sdk/google': + specifier: ^1.0.12 + version: 1.0.12(zod@3.24.1) '@ai-sdk/openai': specifier: ^1.0.11 version: 1.0.11(zod@3.24.1)