Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes #149: Add Temperature and Max Tokens Configuration #176

Merged
merged 10 commits into from
Jan 28, 2025
3 changes: 3 additions & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down
7 changes: 5 additions & 2 deletions src/agent/open-canvas/nodes/customAction.ts
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -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,
Expand Down
14 changes: 7 additions & 7 deletions src/agent/open-canvas/nodes/generate-artifact/index.ts
Original file line number Diff line number Diff line change
@@ -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";

/**
Expand Down
12 changes: 7 additions & 5 deletions src/agent/open-canvas/nodes/generatePath.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion src/agent/open-canvas/nodes/replyToGeneralInput.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand Down
7 changes: 5 additions & 2 deletions src/agent/open-canvas/nodes/rewriteArtifactTheme.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
2 changes: 1 addition & 1 deletion src/agent/open-canvas/nodes/rewriteCodeArtifactTheme.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
12 changes: 6 additions & 6 deletions src/agent/open-canvas/nodes/updateArtifact.ts
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
12 changes: 6 additions & 6 deletions src/agent/open-canvas/nodes/updateHighlightedText.ts
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
39 changes: 28 additions & 11 deletions src/agent/utils.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -139,6 +140,7 @@ export const getModelConfig = (
): {
modelName: string;
modelProvider: string;
modelConfig?: CustomModelConfig;
azureConfig?: {
azureOpenAIApiKey: string;
azureOpenAIApiInstanceName: string;
Expand All @@ -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/", "");
Expand All @@ -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,
};
Expand Down Expand Up @@ -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,
Expand Down
36 changes: 27 additions & 9 deletions src/components/canvas/canvas.tsx
Original file line number Diff line number Diff line change
@@ -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);
Expand Down Expand Up @@ -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);
}
Expand Down
13 changes: 13 additions & 0 deletions src/components/chat-interface/model-selector/constants.ts
Original file line number Diff line number Diff line change
@@ -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,
];
Loading