move to genai

.env.example

@@ -7,10 +7,8 @@ LANGCHAIN_API_KEY=
 ANTHROPIC_API_KEY=
 # OpenAI used for content generation
 OPENAI_API_KEY=
-
-# In Google Cloud Enable GeminiAPI, Create a service account and download the key
-# Set the path to the key file
-GOOGLE_APPLICATION_CREDENTIALS=
+# Optional, only required if using `Gemini 1.5 Flash` as the model.
+GOOGLE_API_KEY=
 
 # LangGraph Deployment, or local development server via LangGraph Studio.
 # If running locally, this URL should be set in the `constants.ts` file.

package.json

@@ -33,7 +33,7 @@
     "@langchain/anthropic": "^0.3.5",
     "@langchain/community": "^0.3.9",
     "@langchain/core": "^0.3.14",
-    "@langchain/google-vertexai": "^0.1.0",
+    "@langchain/google-genai": "^0.1.0",
     "@langchain/langgraph": "^0.2.18",
     "@langchain/langgraph-sdk": "^0.0.17",
     "@langchain/openai": "^0.3.11",

src/agent/open-canvas/nodes/customAction.ts

@@ -13,7 +13,7 @@ import {
 import {
   ensureStoreInConfig,
   formatReflections,
-  getModelNameFromConfig,
+  getModelNameAndProviderFromConfig,
 } from "../../utils";
 import {
   CUSTOM_QUICK_ACTION_ARTIFACT_CONTENT_PROMPT,
@@ -39,9 +39,11 @@ export const customAction = async (
     throw new Error("No custom quick action ID found.");
   }
 
-  const modelName = getModelNameFromConfig(config);
+  const { modelName, modelProvider } =
+    getModelNameAndProviderFromConfig(config);
   const smallModel = await initChatModel(modelName, {
     temperature: 0.5,
+    modelProvider,
   });
 
   const store = ensureStoreInConfig(config);

src/agent/open-canvas/nodes/generateFollowup.ts

@@ -6,7 +6,7 @@ import { Reflections } from "../../../types";
 import {
   ensureStoreInConfig,
   formatReflections,
-  getModelNameFromConfig,
+  getModelNameAndProviderFromConfig,
 } from "../../utils";
 import { FOLLOWUP_ARTIFACT_PROMPT } from "../prompts";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
@@ -18,10 +18,12 @@ export const generateFollowup = async (
   state: typeof OpenCanvasGraphAnnotation.State,
   config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
-  const modelName = getModelNameFromConfig(config);
+  const { modelName, modelProvider } =
+    getModelNameAndProviderFromConfig(config);
   const smallModel = await initChatModel(modelName, {
     temperature: 0.5,
     maxTokens: 250,
+    modelProvider,
   });
 
   const store = ensureStoreInConfig(config);

src/agent/open-canvas/nodes/respondToQuery.ts

@@ -6,7 +6,7 @@ import {
   ensureStoreInConfig,
   formatArtifactContentWithTemplate,
   formatReflections,
-  getModelNameFromConfig,
+  getModelNameAndProviderFromConfig,
 } from "../../utils";
 import { CURRENT_ARTIFACT_PROMPT, NO_ARTIFACT_PROMPT } from "../prompts";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
@@ -18,9 +18,11 @@ export const respondToQuery = async (
   state: typeof OpenCanvasGraphAnnotation.State,
   config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
-  const modelName = getModelNameFromConfig(config);
+  const { modelName, modelProvider } =
+    getModelNameAndProviderFromConfig(config);
   const smallModel = await initChatModel(modelName, {
     temperature: 0.5,
+    modelProvider,
   });
 
   const prompt = `You are an AI assistant tasked with responding to the users question.

src/agent/open-canvas/nodes/rewriteArtifact.ts

@@ -9,7 +9,7 @@ import {
   ensureStoreInConfig,
   formatArtifactContent,
   formatReflections,
-  getModelNameFromConfig,
+  getModelNameAndProviderFromConfig,
 } from "../../utils";
 import {
   ArtifactCodeV3,
@@ -67,9 +67,11 @@ export const rewriteArtifact = async (
     ])
     .withConfig({ runName: "optionally_update_artifact_meta" });
 
-  const modelName = getModelNameFromConfig(config);
+  const { modelName, modelProvider } =
+    getModelNameAndProviderFromConfig(config);
   const smallModel = await initChatModel(modelName, {
     temperature: 0.5,
+    modelProvider,
   });
 
   const store = ensureStoreInConfig(config);

src/agent/open-canvas/nodes/rewriteArtifactTheme.ts

@@ -6,7 +6,7 @@ import { ArtifactV3, Reflections } from "../../../types";
 import {
   ensureStoreInConfig,
   formatReflections,
-  getModelNameFromConfig,
+  getModelNameAndProviderFromConfig,
 } from "../../utils";
 import {
   ADD_EMOJIS_TO_ARTIFACT_PROMPT,
@@ -21,9 +21,11 @@ export const rewriteArtifactTheme = async (
   state: typeof OpenCanvasGraphAnnotation.State,
   config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
-  const modelName = getModelNameFromConfig(config);
+  const { modelName, modelProvider } =
+    getModelNameAndProviderFromConfig(config);
   const smallModel = await initChatModel(modelName, {
     temperature: 0.5,
+    modelProvider,
   });
 
   const store = ensureStoreInConfig(config);

src/agent/open-canvas/nodes/rewriteCodeArtifactTheme.ts

@@ -1,4 +1,4 @@
-import { getModelNameFromConfig } from "@/agent/utils";
+import { getModelNameAndProviderFromConfig } from "@/agent/utils";
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
 import { initChatModel } from "langchain/chat_models/universal";
 import { getArtifactContent } from "../../../hooks/use-graph/utils";
@@ -16,9 +16,11 @@ export const rewriteCodeArtifactTheme = async (
   state: typeof OpenCanvasGraphAnnotation.State,
   config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
-  const modelName = getModelNameFromConfig(config);
+  const { modelName, modelProvider } =
+    getModelNameAndProviderFromConfig(config);
   const smallModel = await initChatModel(modelName, {
     temperature: 0.5,
+    modelProvider,
   });
 
   const currentArtifactContent = state.artifact

src/agent/utils.ts

@@ -80,6 +80,41 @@ export const formatArtifactContentWithTemplate = (
   );
 };
 
+export const getModelNameAndProviderFromConfig = (
+  config: LangGraphRunnableConfig
+): { modelName: string; modelProvider: string } => {
+  const customModelName = config.configurable?.customModelName as string;
+  if (!customModelName) {
+    throw new Error("Model name is missing in config.");
+  }
+  if (customModelName.includes("gpt-")) {
+    return {
+      modelName: customModelName,
+      modelProvider: "openai",
+    };
+  }
+  if (customModelName.includes("claude-")) {
+    return {
+      modelName: customModelName,
+      modelProvider: "anthropic",
+    };
+  }
+  if (customModelName.includes("fireworks/")) {
+    return {
+      modelName: customModelName,
+      modelProvider: "fireworks",
+    };
+  }
+  if (customModelName.includes("gemini-")) {
+    return {
+      modelName: customModelName,
+      modelProvider: "google-genai",
+    };
+  }
+
+  throw new Error("Unknown model provider");
+};
+
 export const getModelNameFromConfig = (
   config: LangGraphRunnableConfig
 ): string => {

src/components/ModelSelector.tsx

@@ -14,10 +14,15 @@ import {
   ANTHROPIC_MODELS,
   OPENAI_MODELS,
   FIREWORKS_MODELS,
-  GEMINI_MODELS
+  GEMINI_MODELS,
 } from "@/constants";
 
-const allModels = [...ANTHROPIC_MODELS, ...OPENAI_MODELS, ...FIREWORKS_MODELS, ...GEMINI_MODELS];
+const allModels = [
+  ...ANTHROPIC_MODELS,
+  ...OPENAI_MODELS,
+  ...FIREWORKS_MODELS,
+  ...GEMINI_MODELS,
+];
 
 const modelNameToLabel = (modelName: ALL_MODEL_NAMES) => {
   const model = allModels.find((m) => m.name === modelName);

src/constants.ts

@@ -1,5 +1,5 @@
 export const LANGGRAPH_API_URL =
-  process.env.LANGGRAPH_API_URL ?? "http://localhost:59811";
+  process.env.LANGGRAPH_API_URL ?? "http://localhost:54790";
 // v2 is tied to the 'open-canvas-prod' deployment.
 export const ASSISTANT_ID_COOKIE = "oc_assistant_id_v2";
 // export const ASSISTANT_ID_COOKIE = "oc_assistant_id";
@@ -48,7 +48,7 @@ export const GEMINI_MODELS = [
   {
     name: "gemini-1.5-flash",
     label: "Gemini 1.5 Flash",
-  }
+  },
 ];
 export const DEFAULT_MODEL_NAME: ALL_MODEL_NAMES = "claude-3-haiku-20240307";
 export type OPENAI_MODEL_NAMES = (typeof OPENAI_MODELS)[number]["name"];