feat: Add reflection

evals/highlights.ts

@@ -1,5 +1,5 @@
 import { type Example, Run } from "langsmith";
-import { graph } from "../src/agent/index";
+import { graph } from "../src/agent/open-canvas/index";
 import { evaluate, EvaluationResult } from "langsmith/evaluation";
 import "dotenv/config";
 

langgraph.json

@@ -3,7 +3,8 @@
   "dockerfile_lines": [],
   "dependencies": ["."],
   "graphs": {
-    "agent": "./src/agent/index.ts:graph"
+    "agent": "./src/agent/open-canvas/index.ts:graph",
+    "reflection": "./src/agent/reflection/index.ts:graph"
   },
   "env": ".env"
 }

package.json

@@ -19,7 +19,7 @@
     "@codemirror/lang-javascript": "^6.2.2",
     "@codemirror/lang-php": "^6.0.1",
     "@codemirror/lang-python": "^6.1.6",
-    "@langchain/anthropic": "^0.3.1",
+    "@langchain/anthropic": "^0.3.3",
     "@langchain/core": "^0.3.3",
     "@langchain/langgraph": "^0.2.10",
     "@langchain/langgraph-sdk": "^0.0.14",

src/agent/index.ts → src/agent/open-canvas/index.ts

@@ -1,5 +1,5 @@
 import { END, Send, START, StateGraph } from "@langchain/langgraph";
-import { GraphAnnotation } from "./state";
+import { OpenCanvasGraphAnnotation } from "./state";
 import { generatePath } from "./nodes/generatePath";
 import { generateFollowup } from "./nodes/generateFollowup";
 import { generateArtifact } from "./nodes/generateArtifact";
@@ -10,7 +10,7 @@ import { respondToQuery } from "./nodes/respondToQuery";
 import { rewriteCodeArtifactTheme } from "./nodes/rewriteCodeArtifactTheme";
 
 const defaultInputs: Omit<
-  typeof GraphAnnotation.State,
+  typeof OpenCanvasGraphAnnotation.State,
   "messages" | "artifacts"
 > = {
   selectedArtifactId: undefined,
@@ -26,7 +26,7 @@ const defaultInputs: Omit<
   portLanguage: undefined,
 };
 
-const routeNode = (state: typeof GraphAnnotation.State) => {
+const routeNode = (state: typeof OpenCanvasGraphAnnotation.State) => {
   if (!state.next) {
     throw new Error("'next' state field not set.");
   }
@@ -36,13 +36,13 @@ const routeNode = (state: typeof GraphAnnotation.State) => {
   });
 };
 
-const cleanState = (_: typeof GraphAnnotation.State) => {
+const cleanState = (_: typeof OpenCanvasGraphAnnotation.State) => {
   return {
     ...defaultInputs,
   };
 };
 
-const builder = new StateGraph(GraphAnnotation)
+const builder = new StateGraph(OpenCanvasGraphAnnotation)
   // Start node & edge
   .addNode("generatePath", generatePath)
   .addEdge(START, "generatePath")

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

@@ -1,16 +1,16 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { GraphAnnotation, GraphReturnType } from "../state";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { NEW_ARTIFACT_PROMPT } from "../prompts";
-import { Artifact } from "../../types";
+import { Artifact } from "../../../types";
 import { z } from "zod";
 import { v4 as uuidv4 } from "uuid";
 
 /**
  * Generate a new artifact based on the user's query.
  */
 export const generateArtifact = async (
-  state: typeof GraphAnnotation.State
-): Promise<GraphReturnType> => {
+  state: typeof OpenCanvasGraphAnnotation.State
+): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,

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

@@ -1,13 +1,13 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { GraphAnnotation, GraphReturnType } from "../state";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { FOLLOWUP_ARTIFACT_PROMPT } from "../prompts";
 
 /**
  * Generate a followup message after generating or updating an artifact.
  */
 export const generateFollowup = async (
-  state: typeof GraphAnnotation.State
-): Promise<GraphReturnType> => {
+  state: typeof OpenCanvasGraphAnnotation.State
+): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o-mini",
     temperature: 0.5,

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

@@ -1,13 +1,13 @@
 import { ChatOpenAI } from "@langchain/openai";
 import { ROUTE_QUERY_PROMPT } from "../prompts";
-import { GraphAnnotation } from "../state";
+import { OpenCanvasGraphAnnotation } from "../state";
 import { formatArtifacts } from "../utils";
 import { z } from "zod";
 
 /**
  * Routes to the proper node in the graph based on the user's query.
  */
-export const generatePath = async (state: typeof GraphAnnotation.State) => {
+export const generatePath = async (state: typeof OpenCanvasGraphAnnotation.State) => {
   if (state.highlighted) {
     return {
       next: "updateArtifact",

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

@@ -1,13 +1,13 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { GraphAnnotation, GraphReturnType } from "../state";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { formatArtifacts } from "../utils";
 
 /**
  * Generate responses to questions. Does not generate artifacts.
  */
 export const respondToQuery = async (
-  state: typeof GraphAnnotation.State
-): Promise<GraphReturnType> => {
+  state: typeof OpenCanvasGraphAnnotation.State
+): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,

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

@@ -1,10 +1,10 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { GraphAnnotation, GraphReturnType } from "../state";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { UPDATE_ENTIRE_ARTIFACT_PROMPT } from "../prompts";
 
 export const rewriteArtifact = async (
-  state: typeof GraphAnnotation.State
-): Promise<GraphReturnType> => {
+  state: typeof OpenCanvasGraphAnnotation.State
+): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,

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

@@ -1,5 +1,5 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { GraphAnnotation, GraphReturnType } from "../state";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import {
   ADD_EMOJIS_TO_ARTIFACT_PROMPT,
   CHANGE_ARTIFACT_LANGUAGE_PROMPT,
@@ -9,8 +9,8 @@ import {
 } from "../prompts";
 
 export const rewriteArtifactTheme = async (
-  state: typeof GraphAnnotation.State
-): Promise<GraphReturnType> => {
+  state: typeof OpenCanvasGraphAnnotation.State
+): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,

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

@@ -1,5 +1,5 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { GraphAnnotation, GraphReturnType } from "../state";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import {
   ADD_COMMENTS_TO_CODE_ARTIFACT_PROMPT,
   ADD_LOGS_TO_CODE_ARTIFACT_PROMPT,
@@ -8,8 +8,8 @@ import {
 } from "../prompts";
 
 export const rewriteCodeArtifactTheme = async (
-  state: typeof GraphAnnotation.State
-): Promise<GraphReturnType> => {
+  state: typeof OpenCanvasGraphAnnotation.State
+): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,

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

@@ -1,13 +1,13 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { GraphAnnotation, GraphReturnType } from "../state";
+import { OpenCanvasGraphAnnotation, OpenCanvasGraphReturnType } from "../state";
 import { UPDATE_HIGHLIGHTED_ARTIFACT_PROMPT } from "../prompts";
 
 /**
  * Update an existing artifact based on the user's query.
  */
 export const updateArtifact = async (
-  state: typeof GraphAnnotation.State
-): Promise<GraphReturnType> => {
+  state: typeof OpenCanvasGraphAnnotation.State
+): Promise<OpenCanvasGraphReturnType> => {
   const smallModel = new ChatOpenAI({
     model: "gpt-4o",
     temperature: 0.5,

src/agent/state.ts → src/agent/open-canvas/state.ts

@@ -5,7 +5,7 @@ import {
   ProgrammingLanguageOptions,
   ReadingLevelOptions,
   Highlight,
-} from "../types";
+} from "../../types";
 import { Annotation, MessagesAnnotation } from "@langchain/langgraph";
 
 /**
@@ -23,7 +23,7 @@ const artifactsReducer = (
   return state.filter((a) => !updatedIds.has(a.id)).concat(update);
 };
 
-export const GraphAnnotation = Annotation.Root({
+export const OpenCanvasGraphAnnotation = Annotation.Root({
   ...MessagesAnnotation.spec,
   /**
    * The ID of the artifact to perform some action on.
@@ -79,4 +79,4 @@ export const GraphAnnotation = Annotation.Root({
   fixBugs: Annotation<boolean | undefined>,
 });
 
-export type GraphReturnType = Partial<typeof GraphAnnotation.State>;
+export type OpenCanvasGraphReturnType = Partial<typeof OpenCanvasGraphAnnotation.State>;

src/agent/utils.ts → src/agent/open-canvas/utils.ts

@@ -1,4 +1,4 @@
-import { Artifact } from "../types";
+import { Artifact } from "../../types";
 
 export const formatArtifacts = (
   messages: Artifact[],

src/agent/reflection/index.ts

@@ -0,0 +1,83 @@
+import { ChatAnthropic } from "@langchain/anthropic";
+import { type LangGraphRunnableConfig, type BaseStore, StateGraph, START } from "@langchain/langgraph";
+import { ReflectionGraphAnnotation, ReflectionGraphReturnType } from "./state";
+import { Memory } from "../../types";
+import { REFLECT_SYSTEM_PROMPT } from "./prompts";
+import { z } from "zod";
+
+const ensureStoreFromConfig = (config: LangGraphRunnableConfig): BaseStore => {
+  if (!config.store) {
+    throw new Error("`store` not found in config");
+  }
+
+  return config.store;
+};
+
+const formatMemories = (memories: Memory): string => {
+  const styleString = `The following is a list of style guidelines previously generated by you:
+<style-guidelines>
+- ${memories.styleRules.join("\n- ")}
+</style-guidelines>`;
+  const contentString = `The following is a list of memories/facts you previously generated about the user:
+<user-facts>
+- ${memories.content.join("\n- ")}
+</user-facts>`;
+
+
+  return styleString + "\n\n" + contentString;
+}
+
+export const reflect = async (
+  state: typeof ReflectionGraphAnnotation.State,
+  config: LangGraphRunnableConfig,
+): Promise<ReflectionGraphReturnType> => {
+  const store = ensureStoreFromConfig(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 ? formatMemories(memories.value as Memory) : "No memories found.";
+
+  const generateReflectionsSchema = z.object({
+    styleRules: z.array(z.string()).describe("The complete new list of style rules and guidelines."),
+    content: z.array(z.string()).describe("The complete new list of memories/facts about the user.")
+  });
+
+  const model = new ChatAnthropic({
+    model: "claude-3-5-sonnet-20240620",
+    temperature: 0,
+  }).withStructuredOutput(generateReflectionsSchema, {
+    name: "generate_reflections",
+  })
+
+  const formattedSystemPrompt = REFLECT_SYSTEM_PROMPT
+    .replace("{artifact}", state.artifact?.content ?? "No artifact found.")
+    .replace("{reflections}", memoriesAsString);
+
+  const result = await model.invoke([
+    {
+      role: "system",
+      content: formattedSystemPrompt
+    },
+    ...state.messages,
+  ]);
+
+  const newMemories = {
+    styleRules: result.styleRules,
+    content: result.content,
+  };
+
+  await store.put(memoryNamespace, memoryKey, newMemories);
+
+  return {};
+}
+
+const builder = new StateGraph(ReflectionGraphAnnotation)
+  .addNode("reflect", reflect)
+  .addEdge(START, "reflect");
+
+export const graph = builder.compile();

src/agent/reflection/prompts.ts

@@ -0,0 +1,30 @@
+export const REFLECT_SYSTEM_PROMPT = `You are an expert assistant, and writer. You are tasked with reflecting on the following conversation between a user and an AI assistant.
+You are also provided with an 'artifact' the user and assistant worked together on to write. Artifacts can be code, creative writing, emails, or any other form of written content.
+
+<artifact>
+{artifact}
+</artifact>
+
+You have also previously generated the following reflections about the user. Your reflections are broken down into two categories:
+1. Style Guidelines: These are the style guidelines you have generated for the user. Style guidelines can be anything from writing style, to code style, to design style.
+  They should be general, and apply to the all the users work, including the conversation and artifact generated.
+2. Content: These are general memories, facts, and insights you generate about the user. These can be anything from the users interests, to their goals, to their personality traits.
+  Ensure you think carefully about what goes in here, as the assistant will use these when generating future responses or artifacts for the user.
+  
+<reflections>
+{reflections}
+</reflections>
+
+Your job is to take all of the context and existing reflections and re-generate all. Use these guidelines when generating the new set of reflections:
+
+<system-guidelines>
+- Ensure your reflections are relevant to the conversation and artifact.
+- Remove duplicate reflections, or combine multiple reflections into one if they are duplicating content.
+- Do not remove reflections unless the conversation/artifact clearly demonstrates they should no longer be included.
+  This does NOT mean remove reflections if you see no evidence of them in the conversation/artifact, but instead remove them if the user indicates they are no longer relevant.
+- Think of the deeper meaning behind a users messages, or artifact to generate content reflections that will be relevant and useful for future interactions.
+- Keep the rules you list high signal-to-noise - don't include unnecessary reflections, but make sure the ones you do add are descriptive.
+- Your reflections should be very descriptive and detailed, ensuring they are clear and will not be misinterpreted.
+</system-guidelines>
+
+Finally, use the 'generate_reflections' tool to generate the new, full list of reflections.`

src/agent/reflection/state.ts

@@ -0,0 +1,15 @@
+import { Artifact } from "@/types";
+import { Annotation, MessagesAnnotation } from "@langchain/langgraph";
+
+export const ReflectionGraphAnnotation = Annotation.Root({
+  /**
+   * The chat history to reflect on.
+   */
+  ...MessagesAnnotation.spec,
+  /**
+   * The artifact to reflect on.
+   */
+  artifact: Annotation<Artifact | undefined>
+});
+
+export type ReflectionGraphReturnType = Partial<typeof ReflectionGraphAnnotation.State>;

src/types.ts

@@ -66,3 +66,15 @@ export type ReadingLevelOptions =
   | "teenager"
   | "college"
   | "phd";
+
+
+export interface Memory {
+  /**
+   * Style rules to follow for generating content.
+   */
+  styleRules: string[];
+  /**
+   * Key content to remember about the user when generating content.
+   */
+  content: string[];
+}

yarn.lock

@@ -487,10 +487,10 @@
     "@jridgewell/resolve-uri" "^3.1.0"
     "@jridgewell/sourcemap-codec" "^1.4.14"
 
-"@langchain/anthropic@^0.3.1":
-  version "0.3.1"
-  resolved "https://registry.yarnpkg.com/@langchain/anthropic/-/anthropic-0.3.1.tgz#7981e90e87566dedbd6cd5e4534a205cd618ec13"
-  integrity sha512-mE2zfv2IeogDgP53IYlfqZuXbmACkQDtBkg1Nriasrzvr1g+1Q85pfyUCnxc05p4S2IplbJhvoBtgdQUTGenqw==
+"@langchain/anthropic@^0.3.3":
+  version "0.3.3"
+  resolved "https://registry.yarnpkg.com/@langchain/anthropic/-/anthropic-0.3.3.tgz#68744efbebed581be45cd6fc85483de18511278c"
+  integrity sha512-OvnSV3Tjhb87n7CxWzIcJqcJEM4qoFDYYt6Rua7glQF/Ud5FBTurlzoMunLPTQeF5GdPiaOwP3nUw6I9gF7ppw==
   dependencies:
     "@anthropic-ai/sdk" "^0.27.3"
     fast-xml-parser "^4.4.1"