use reflections throughout graph

src/agent/open-canvas/index.ts

@@ -8,6 +8,7 @@ import { rewriteArtifactTheme } from "./nodes/rewriteArtifactTheme";
 import { updateArtifact } from "./nodes/updateArtifact";
 import { respondToQuery } from "./nodes/respondToQuery";
 import { rewriteCodeArtifactTheme } from "./nodes/rewriteCodeArtifactTheme";
+import { reflect } from "../reflection";
 
 const defaultInputs: Omit<
   typeof OpenCanvasGraphAnnotation.State,
@@ -57,6 +58,7 @@ const builder = new StateGraph(OpenCanvasGraphAnnotation)
   .addNode("generateArtifact", generateArtifact)
   .addNode("generateFollowup", generateFollowup)
   .addNode("cleanState", cleanState)
+  .addNode("reflect", reflect)
   // Edges
   .addEdge("generateArtifact", "generateFollowup")
   .addEdge("updateArtifact", "generateFollowup")
@@ -65,7 +67,9 @@ const builder = new StateGraph(OpenCanvasGraphAnnotation)
   .addEdge("rewriteCodeArtifactTheme", "generateFollowup")
   // End edges
   .addEdge("respondToQuery", "cleanState")
-  .addEdge("generateFollowup", "cleanState")
+  // Only reflect if an artifact was generated/updated.
+  .addEdge("generateFollowup", "reflect")
+  .addEdge("reflect", "cleanState")
   .addEdge("cleanState", END);
 
 export const graph = builder.compile();

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

@@ -1,21 +1,36 @@
 import { ChatOpenAI } from "@langchain/openai";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { NEW_ARTIFACT_PROMPT } from "../prompts";
-import { Artifact } from "../../../types";
+import { Artifact, Reflections } from "../../../types";
 import { z } from "zod";
 import { v4 as uuidv4 } from "uuid";
+import { ensureStoreInConfig, formatReflections } from "@/agent/utils";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
 
 /**
  * Generate a new artifact based on the user's query.
  */
 export const generateArtifact = async (
-  state: typeof OpenCanvasGraphAnnotation.State
+  state: typeof OpenCanvasGraphAnnotation.State,
+  config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,
   });
 
+  const store = ensureStoreInConfig(config);
+  const assistantId = config.configurable?.assistant_id;
+  if (!assistantId) {
+    throw new Error("`assistant_id` not found in configurable");
+  }
+  const memoryNamespace = ["memories", assistantId];
+  const memoryKey = "reflection";
+  const memories = await store.get(memoryNamespace, memoryKey);
+  const memoriesAsString = memories?.value
+    ? formatReflections(memories.value as Reflections)
+    : "No reflections found.";
+
   const modelWithArtifactTool = smallModel.bindTools(
     [
       {
@@ -46,8 +61,16 @@ export const generateArtifact = async (
     { tool_choice: "generate_artifact" }
   );
 
+  const formattedNewArtifactPrompt = NEW_ARTIFACT_PROMPT.replace(
+    "{reflections}",
+    memoriesAsString
+  );
+
   const response = await modelWithArtifactTool.invoke(
-    [{ role: "system", content: NEW_ARTIFACT_PROMPT }, ...state.messages],
+    [
+      { role: "system", content: formattedNewArtifactPrompt },
+      ...state.messages,
+    ],
     { runName: "generate_artifact" }
   );
   const newArtifact: Artifact = {

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

@@ -1,24 +1,41 @@
 import { ChatOpenAI } from "@langchain/openai";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { FOLLOWUP_ARTIFACT_PROMPT } from "../prompts";
+import { ensureStoreInConfig, formatReflections } from "@/agent/utils";
+import { Reflections } from "../../../types";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
 
 /**
  * Generate a followup message after generating or updating an artifact.
  */
 export const generateFollowup = async (
-  state: typeof OpenCanvasGraphAnnotation.State
+  state: typeof OpenCanvasGraphAnnotation.State,
+  config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o-mini",
     temperature: 0.5,
     maxTokens: 250,
   });
 
+  const store = ensureStoreInConfig(config);
+  const assistantId = config.configurable?.assistant_id;
+  if (!assistantId) {
+    throw new Error("`assistant_id` not found in configurable");
+  }
+  const memoryNamespace = ["memories", assistantId];
+  const memoryKey = "reflection";
+  const memories = await store.get(memoryNamespace, memoryKey);
+  const memoriesAsString = memories?.value
+    ? formatReflections(memories.value as Reflections)
+    : "No reflections found.";
+
   const recentArtifact = state.artifacts[state.artifacts.length - 1];
   const formattedPrompt = FOLLOWUP_ARTIFACT_PROMPT.replace(
     "{artifactContent}",
     recentArtifact.content
-  );
+  ).replace("{reflections}", memoriesAsString);
+
   const response = await smallModel.invoke([
     { role: "user", content: formattedPrompt },
   ]);

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

@@ -0,0 +1,35 @@
+import { Client } from "@langchain/langgraph-sdk";
+import { OpenCanvasGraphAnnotation } from "../state";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+
+export const reflect = async (
+  state: typeof OpenCanvasGraphAnnotation.State,
+  config: LangGraphRunnableConfig
+) => {
+  const langGraphClient = new Client();
+
+  const selectedArtifact = state.selectedArtifactId
+    ? state.artifacts.find((art) => art.id === state.selectedArtifactId)
+    : state.artifacts[state.artifacts.length - 1];
+  const reflectionInput = {
+    messages: state.messages,
+    artifact: selectedArtifact,
+  };
+  const reflectionConfig = {
+    configurable: {
+      // Ensure we pass in the current graph's assistant ID as this is
+      // how we fetch & store the memories.
+      assistant_id: config.configurable?.assistant_id,
+    },
+  };
+
+  const newThread = await langGraphClient.threads.create();
+  // Create a new reflection run, but do not `wait` for it to finish.
+  // Intended to be a background run.
+  await langGraphClient.runs.create(newThread.thread_id, "reflection", {
+    input: reflectionInput,
+    config: reflectionConfig,
+  });
+
+  return {};
+};

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

@@ -1,12 +1,16 @@
 import { ChatOpenAI } from "@langchain/openai";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { formatArtifacts } from "../utils";
+import { ensureStoreInConfig, formatReflections } from "@/agent/utils";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+import { Reflections } from "../../../types";
 
 /**
  * Generate responses to questions. Does not generate artifacts.
  */
 export const respondToQuery = async (
-  state: typeof OpenCanvasGraphAnnotation.State
+  state: typeof OpenCanvasGraphAnnotation.State,
+  config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
@@ -17,14 +21,30 @@ export const respondToQuery = async (
   
 The user has generated artifacts in the past. Use the following artifacts as context when responding to the users question.
 
+You also have the following reflections on style guidelines and general memories/facts about the user to use when generating your response.
+<reflections>
+{reflections}
+</reflections>
+
 <artifacts>
 {artifacts}
 </artifacts>`;
 
-  const formattedPrompt = prompt.replace(
-    "{artifacts}",
-    formatArtifacts(state.artifacts)
-  );
+  const store = ensureStoreInConfig(config);
+  const assistantId = config.configurable?.assistant_id;
+  if (!assistantId) {
+    throw new Error("`assistant_id` not found in configurable");
+  }
+  const memoryNamespace = ["memories", assistantId];
+  const memoryKey = "reflection";
+  const memories = await store.get(memoryNamespace, memoryKey);
+  const memoriesAsString = memories?.value
+    ? formatReflections(memories.value as Reflections)
+    : "No reflections found.";
+
+  const formattedPrompt = prompt
+    .replace("{artifacts}", formatArtifacts(state.artifacts))
+    .replace("{reflections}", memoriesAsString);
 
   const response = await smallModel.invoke([
     { role: "system", content: formattedPrompt },

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

@@ -1,15 +1,31 @@
 import { ChatOpenAI } from "@langchain/openai";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { UPDATE_ENTIRE_ARTIFACT_PROMPT } from "../prompts";
+import { ensureStoreInConfig, formatReflections } from "@/agent/utils";
+import { Reflections } from "../../../types";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
 
 export const rewriteArtifact = async (
-  state: typeof OpenCanvasGraphAnnotation.State
+  state: typeof OpenCanvasGraphAnnotation.State,
+  config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,
   });
 
+  const store = ensureStoreInConfig(config);
+  const assistantId = config.configurable?.assistant_id;
+  if (!assistantId) {
+    throw new Error("`assistant_id` not found in configurable");
+  }
+  const memoryNamespace = ["memories", assistantId];
+  const memoryKey = "reflection";
+  const memories = await store.get(memoryNamespace, memoryKey);
+  const memoriesAsString = memories?.value
+    ? formatReflections(memories.value as Reflections)
+    : "No reflections found.";
+
   const selectedArtifact = state.artifacts.find(
     (artifact) => artifact.id === state.selectedArtifactId
   );
@@ -20,7 +36,7 @@ export const rewriteArtifact = async (
   const formattedPrompt = UPDATE_ENTIRE_ARTIFACT_PROMPT.replace(
     "{artifactContent}",
     selectedArtifact.content
-  );
+  ).replace("{reflections}", memoriesAsString);
 
   const recentHumanMessage = state.messages.findLast(
     (message) => message._getType() === "human"

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

@@ -7,15 +7,31 @@ import {
   CHANGE_ARTIFACT_READING_LEVEL_PROMPT,
   CHANGE_ARTIFACT_TO_PIRATE_PROMPT,
 } from "../prompts";
+import { ensureStoreInConfig, formatReflections } from "@/agent/utils";
+import { Reflections } from "../../../types";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
 
 export const rewriteArtifactTheme = async (
-  state: typeof OpenCanvasGraphAnnotation.State
+  state: typeof OpenCanvasGraphAnnotation.State,
+  config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,
   });
 
+  const store = ensureStoreInConfig(config);
+  const assistantId = config.configurable?.assistant_id;
+  if (!assistantId) {
+    throw new Error("`assistant_id` not found in configurable");
+  }
+  const memoryNamespace = ["memories", assistantId];
+  const memoryKey = "reflection";
+  const memories = await store.get(memoryNamespace, memoryKey);
+  const memoriesAsString = memories?.value
+    ? formatReflections(memories.value as Reflections)
+    : "No reflections found.";
+
   const selectedArtifact = state.artifacts.find(
     (artifact) => artifact.id === state.selectedArtifactId
   );
@@ -83,6 +99,8 @@ export const rewriteArtifactTheme = async (
     throw new Error("No theme selected");
   }
 
+  formattedPrompt = formattedPrompt.replace("{reflections}", memoriesAsString);
+
   const newArtifactValues = await smallModel.invoke([
     { role: "user", content: formattedPrompt },
   ]);

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

@@ -1,18 +1,34 @@
 import { ChatOpenAI } from "@langchain/openai";
 import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { UPDATE_HIGHLIGHTED_ARTIFACT_PROMPT } from "../prompts";
+import { ensureStoreInConfig, formatReflections } from "@/agent/utils";
+import { Reflections } from "../../../types";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
 
 /**
  * Update an existing artifact based on the user's query.
  */
 export const updateArtifact = async (
-  state: typeof OpenCanvasGraphAnnotation.State
+  state: typeof OpenCanvasGraphAnnotation.State,
+  config: LangGraphRunnableConfig
 ): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,
   });
 
+  const store = ensureStoreInConfig(config);
+  const assistantId = config.configurable?.assistant_id;
+  if (!assistantId) {
+    throw new Error("`assistant_id` not found in configurable");
+  }
+  const memoryNamespace = ["memories", assistantId];
+  const memoryKey = "reflection";
+  const memories = await store.get(memoryNamespace, memoryKey);
+  const memoriesAsString = memories?.value
+    ? formatReflections(memories.value as Reflections)
+    : "No reflections found.";
+
   const selectedArtifact = state.artifacts.find(
     (artifact) => artifact.id === state.selectedArtifactId
   );
@@ -51,7 +67,8 @@ export const updateArtifact = async (
     highlightedText
   )
     .replace("{beforeHighlight}", beforeHighlight)
-    .replace("{afterHighlight}", afterHighlight);
+    .replace("{afterHighlight}", afterHighlight)
+    .replace("{reflections}", memoriesAsString);
 
   const recentHumanMessage = state.messages.findLast(
     (message) => message._getType() === "human"