Merge pull request #176 from ahmad2b/ahmad2b/model-config

Fixes #149: Add Temperature and Max Tokens Configuration

package.json

@@ -49,6 +49,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",
@@ -64,6 +66,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.

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,
@@ -139,6 +140,7 @@ export const getModelConfig = (
 ): {
   modelName: string;
   modelProvider: string;
+  modelConfig?: CustomModelConfig;
   azureConfig?: {
     azureOpenAIApiKey: string;
     azureOpenAIApiInstanceName: string;
@@ -150,9 +152,9 @@ export const getModelConfig = (
   baseUrl?: string;
 } => {
   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/", "");
@@ -172,30 +174,35 @@ export const getModelConfig = (
     };
   }
 
+  const providerConfig = {
+    modelName: customModelName,
+    modelConfig,
+  };
+
   if (customModelName.includes("gpt-")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "openai",
       apiKey: process.env.OPENAI_API_KEY,
     };
   }
   if (customModelName.includes("claude-")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "anthropic",
       apiKey: process.env.ANTHROPIC_API_KEY,
     };
   }
   if (customModelName.includes("fireworks/")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "fireworks",
       apiKey: process.env.FIREWORKS_API_KEY,
     };
   }
   if (customModelName.includes("gemini-")) {
     return {
-      modelName: customModelName,
+      ...providerConfig,
       modelProvider: "google-genai",
       apiKey: process.env.GOOGLE_API_KEY,
     };
@@ -225,9 +232,19 @@ export async function getModelFromConfig(
     maxTokens?: number;
   }
 ) {
-  const { temperature = 0.5, maxTokens } = extra || {};
-  const { modelName, modelProvider, azureConfig, apiKey, baseUrl } =
-    getModelConfig(config);
+  const {
+    modelName,
+    modelProvider,
+    azureConfig,
+    apiKey,
+    baseUrl,
+    modelConfig,
+  } = getModelConfig(config);
+  const { temperature = 0.5, maxTokens } = {
+    temperature: modelConfig?.temperatureRange.current,
+    maxTokens: modelConfig?.maxTokens.current,
+    ...extra,
+  };
 
   return await initChatModel(modelName, {
     modelProvider,

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,
+];