Fixes #149: Add Temperature and Max Tokens Configuration

package.json

@@ -45,6 +45,8 @@
     "@radix-ui/react-hover-card": "^1.1.2",
     "@radix-ui/react-icons": "^1.3.0",
     "@radix-ui/react-label": "^2.1.0",
+    "@radix-ui/react-navigation-menu": "^1.2.1",
+    "@radix-ui/react-popover": "^1.1.2",
     "@radix-ui/react-progress": "^1.1.0",
     "@radix-ui/react-select": "^2.1.1",
     "@radix-ui/react-slider": "^1.2.1",
@@ -60,6 +62,7 @@
     "@vercel/kv": "^2.0.0",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.1",
+    "cmdk": "1.0.0",
     "date-fns": "^4.1.0",
     "dotenv": "^16.4.5",
     "eslint-plugin-unused-imports": "^4.1.4",

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

@@ -1,6 +1,5 @@
 import { BaseMessage } from "@langchain/core/messages";
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { getModelFromConfig } from "../../utils";
 import { getArtifactContent } from "../../../contexts/utils";
 import { isArtifactMarkdownContent } from "../../../lib/artifact_content_types";
 import {
@@ -10,7 +9,11 @@ import {
   CustomQuickAction,
   Reflections,
 } from "../../../types";
-import { ensureStoreInConfig, formatReflections } from "../../utils";
+import {
+  ensureStoreInConfig,
+  formatReflections,
+  getModelFromConfig,
+} from "../../utils";
 import {
   CUSTOM_QUICK_ACTION_ARTIFACT_CONTENT_PROMPT,
   CUSTOM_QUICK_ACTION_ARTIFACT_PROMPT_PREFIX,

src/agent/open-canvas/nodes/generate-artifact/index.ts

@@ -1,16 +1,16 @@
-import {
-  OpenCanvasGraphAnnotation,
-  OpenCanvasGraphReturnType,
-} from "../../state";
-import { LangGraphRunnableConfig } from "@langchain/langgraph";
 import {
   getFormattedReflections,
-  getModelFromConfig,
   getModelConfig,
+  getModelFromConfig,
   optionallyGetSystemPromptFromConfig,
 } from "@/agent/utils";
-import { ARTIFACT_TOOL_SCHEMA } from "./schemas";
 import { ArtifactV3 } from "@/types";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+import {
+  OpenCanvasGraphAnnotation,
+  OpenCanvasGraphReturnType,
+} from "../../state";
+import { ARTIFACT_TOOL_SCHEMA } from "./schemas";
 import { createArtifactContent, formatNewArtifactPrompt } from "./utils";
 
 /**

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

@@ -1,3 +1,10 @@
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+import { z } from "zod";
+import { getArtifactContent } from "../../../contexts/utils";
+import {
+  formatArtifactContentWithTemplate,
+  getModelFromConfig,
+} from "../../utils";
 import {
   CURRENT_ARTIFACT_PROMPT,
   NO_ARTIFACT_PROMPT,
@@ -6,11 +13,6 @@ import {
   ROUTE_QUERY_PROMPT,
 } from "../prompts";
 import { OpenCanvasGraphAnnotation } from "../state";
-import { z } from "zod";
-import { formatArtifactContentWithTemplate } from "../../utils";
-import { getArtifactContent } from "../../../contexts/utils";
-import { getModelFromConfig } from "../../utils";
-import { LangGraphRunnableConfig } from "@langchain/langgraph";
 
 /**
  * Routes to the proper node in the graph based on the user's query.
@@ -19,7 +21,6 @@ export const generatePath = async (
   state: typeof OpenCanvasGraphAnnotation.State,
   config: LangGraphRunnableConfig
 ) => {
-  console.log("config.configurable!!", config.configurable);
   if (state.highlightedCode) {
     return {
       next: "updateArtifact",

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

@@ -1,11 +1,11 @@
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { getModelFromConfig } from "../../utils";
 import { getArtifactContent } from "../../../contexts/utils";
 import { Reflections } from "../../../types";
 import {
   ensureStoreInConfig,
   formatArtifactContentWithTemplate,
   formatReflections,
+  getModelFromConfig,
 } from "../../utils";
 import { CURRENT_ARTIFACT_PROMPT, NO_ARTIFACT_PROMPT } from "../prompts";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";

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

@@ -1,9 +1,12 @@
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { getModelFromConfig } from "../../utils";
 import { getArtifactContent } from "../../../contexts/utils";
 import { isArtifactMarkdownContent } from "../../../lib/artifact_content_types";
 import { ArtifactV3, Reflections } from "../../../types";
-import { ensureStoreInConfig, formatReflections } from "../../utils";
+import {
+  ensureStoreInConfig,
+  formatReflections,
+  getModelFromConfig,
+} from "../../utils";
 import {
   ADD_EMOJIS_TO_ARTIFACT_PROMPT,
   CHANGE_ARTIFACT_LANGUAGE_PROMPT,

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

@@ -1,8 +1,8 @@
 import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { getModelFromConfig } from "../../utils";
 import { getArtifactContent } from "../../../contexts/utils";
 import { isArtifactCodeContent } from "../../../lib/artifact_content_types";
 import { ArtifactCodeV3, ArtifactV3 } from "../../../types";
+import { getModelFromConfig } from "../../utils";
 import {
   ADD_COMMENTS_TO_CODE_ARTIFACT_PROMPT,
   ADD_LOGS_TO_CODE_ARTIFACT_PROMPT,

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

@@ -1,15 +1,15 @@
-import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
-import { UPDATE_HIGHLIGHTED_ARTIFACT_PROMPT } from "../prompts";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+import { getArtifactContent } from "../../../contexts/utils";
+import { isArtifactCodeContent } from "../../../lib/artifact_content_types";
+import { ArtifactCodeV3, ArtifactV3, Reflections } from "../../../types";
 import {
   ensureStoreInConfig,
   formatReflections,
   getModelConfig,
   getModelFromConfig,
 } from "../../utils";
-import { ArtifactCodeV3, ArtifactV3, Reflections } from "../../../types";
-import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { getArtifactContent } from "../../../contexts/utils";
-import { isArtifactCodeContent } from "../../../lib/artifact_content_types";
+import { UPDATE_HIGHLIGHTED_ARTIFACT_PROMPT } from "../prompts";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 
 /**
  * Update an existing artifact based on the user's query.

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

@@ -1,13 +1,13 @@
-import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
-import { ArtifactMarkdownV3 } from "../../../types";
-import { getArtifactContent } from "../../../contexts/utils";
-import { isArtifactMarkdownContent } from "../../../lib/artifact_content_types";
 import { getModelConfig, getModelFromConfig } from "@/agent/utils";
-import { LangGraphRunnableConfig } from "@langchain/langgraph";
-import { RunnableBinding } from "@langchain/core/runnables";
 import { BaseLanguageModelInput } from "@langchain/core/language_models/base";
 import { AIMessageChunk } from "@langchain/core/messages";
+import { RunnableBinding } from "@langchain/core/runnables";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
 import { ConfigurableChatModelCallOptions } from "langchain/chat_models/universal";
+import { getArtifactContent } from "../../../contexts/utils";
+import { isArtifactMarkdownContent } from "../../../lib/artifact_content_types";
+import { ArtifactMarkdownV3 } from "../../../types";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 
 const PROMPT = `You are an expert AI writing assistant, tasked with rewriting some text a user has selected. The selected text is nested inside a larger 'block'. You should always respond with ONLY the updated text block in accordance with the user's request.
 You should always respond with the full markdown text block, as it will simply replace the existing block in the artifact.

src/agent/utils.ts

@@ -1,7 +1,8 @@
 import { isArtifactCodeContent } from "@/lib/artifact_content_types";
+import { CustomModelConfig } from "@/types";
 import { BaseStore, LangGraphRunnableConfig } from "@langchain/langgraph";
-import { ArtifactCodeV3, ArtifactMarkdownV3, Reflections } from "../types";
 import { initChatModel } from "langchain/chat_models/universal";
+import { ArtifactCodeV3, ArtifactMarkdownV3, Reflections } from "../types";
 
 export const formatReflections = (
   reflections: Reflections,
@@ -136,6 +137,7 @@ export const getModelConfig = (
 ): {
   modelName: string;
   modelProvider: string;
+  modelConfig?: CustomModelConfig;
   azureConfig?: {
     azureOpenAIApiKey: string;
     azureOpenAIApiInstanceName: string;
@@ -145,9 +147,9 @@ export const getModelConfig = (
   };
 } => {
   const customModelName = config.configurable?.customModelName as string;
-  if (!customModelName) {
-    throw new Error("Model name is missing in config.");
-  }
+  if (!customModelName) throw new Error("Model name is missing in config.");
+
+  const modelConfig = config.configurable?.modelConfig as CustomModelConfig;
 
   if (customModelName.startsWith("azure/")) {
     const actualModelName = customModelName.replace("azure/", "");
@@ -167,27 +169,32 @@ export const getModelConfig = (
     };
   }
 
+  const providerConfig = {
+    modelName: customModelName,
+    modelConfig,
+  };
+
   if (customModelName.includes("gpt-")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "openai",
     };
   }
   if (customModelName.includes("claude-")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "anthropic",
     };
   }
   if (customModelName.includes("fireworks/")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "fireworks",
     };
   }
   if (customModelName.includes("gemini-")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "google-genai",
     };
   }
@@ -208,12 +215,14 @@ export async function getModelFromConfig(
     maxTokens?: number;
   }
 ) {
-  const { temperature = 0.5, maxTokens } = extra || {};
-  const { modelName, modelProvider, azureConfig } = getModelConfig(config);
+  const { modelName, modelProvider, azureConfig, modelConfig } =
+    getModelConfig(config);
+  const { temperature, maxTokens } = extra || {};
+
   return await initChatModel(modelName, {
     modelProvider,
-    temperature,
-    maxTokens,
+    temperature: temperature ?? modelConfig?.temperatureRange.current ?? 0.5,
+    maxTokens: maxTokens ?? modelConfig?.maxTokens.current,
     ...(azureConfig != null
       ? {
           azureOpenAIApiKey: azureConfig.azureOpenAIApiKey,

src/components/canvas/canvas.tsx

@@ -1,25 +1,30 @@
 "use client";
 
 import { ArtifactRenderer } from "@/components/artifacts/ArtifactRenderer";
-import { ContentComposerChatInterface } from "./content-composer";
-import { ALL_MODEL_NAMES } from "@/constants";
+import {
+  ALL_MODEL_NAMES,
+  DEFAULT_MODEL_CONFIG,
+  DEFAULT_MODEL_NAME,
+} from "@/constants";
+import { useGraphContext } from "@/contexts/GraphContext";
 import { useToast } from "@/hooks/use-toast";
 import { getLanguageTemplate } from "@/lib/get_language_template";
 import { cn } from "@/lib/utils";
 import {
   ArtifactCodeV3,
   ArtifactMarkdownV3,
   ArtifactV3,
+  CustomModelConfig,
   ProgrammingLanguageOptions,
 } from "@/types";
-import { useEffect, useState } from "react";
-import { useGraphContext } from "@/contexts/GraphContext";
-import React from "react";
+import React, { useEffect, useState } from "react";
+import { ContentComposerChatInterface } from "./content-composer";
 
 export function CanvasComponent() {
   const { threadData, graphData, userData } = useGraphContext();
   const { user } = userData;
-  const { threadId, clearThreadsWithNoValues, setModelName } = threadData;
+  const { threadId, clearThreadsWithNoValues, setModelName, setModelConfig } =
+    threadData;
   const { setArtifact } = graphData;
   const { toast } = useToast();
   const [chatStarted, setChatStarted] = useState(false);
@@ -88,9 +93,22 @@ export function CanvasComponent() {
             // Chat should only be "started" if there are messages present
             if ((thread.values as Record<string, any>)?.messages?.length) {
               setChatStarted(true);
-              setModelName(
-                thread?.metadata?.customModelName as ALL_MODEL_NAMES
-              );
+              if (thread?.metadata?.customModelName) {
+                setModelName(
+                  thread.metadata.customModelName as ALL_MODEL_NAMES
+                );
+              } else {
+                setModelName(DEFAULT_MODEL_NAME);
+              }
+
+              if (thread?.metadata?.modelConfig) {
+                setModelConfig(
+                  thread?.metadata?.customModelName as ALL_MODEL_NAMES,
+                  thread.metadata?.modelConfig as CustomModelConfig
+                );
+              } else {
+                setModelConfig(DEFAULT_MODEL_NAME, DEFAULT_MODEL_CONFIG);
+              }
             } else {
               setChatStarted(false);
             }

src/components/chat-interface/model-selector/constants.ts

@@ -0,0 +1,13 @@
+import {
+  ANTHROPIC_MODELS,
+  OPENAI_MODELS,
+  FIREWORKS_MODELS,
+  GEMINI_MODELS,
+} from "@/constants";
+
+export const ALL_MODELS = [
+  ...ANTHROPIC_MODELS,
+  ...OPENAI_MODELS,
+  ...FIREWORKS_MODELS,
+  ...GEMINI_MODELS,
+];