Merge pull request #34 from langchain-ai/brace/new-store-api

feat(agent): Implement the new store API
This commit is contained in:
Brace Sproul
2024-10-01 15:17:36 -07:00
committed by GitHub
18 changed files with 491 additions and 469 deletions
+3 -2
View File
@@ -16,8 +16,8 @@
"@assistant-ui/react-syntax-highlighter": "^0.0.11",
"@langchain/anthropic": "^0.3.1",
"@langchain/core": "^0.3.3",
"@langchain/langgraph": "^0.2.8",
"@langchain/langgraph-sdk": "^0.0.11",
"@langchain/langgraph": "^0.2.10",
"@langchain/langgraph-sdk": "^0.0.14",
"@radix-ui/react-avatar": "^1.1.0",
"@radix-ui/react-dialog": "^1.1.1",
"@radix-ui/react-dropdown-menu": "^2.1.1",
@@ -27,6 +27,7 @@
"@radix-ui/react-slot": "^1.1.0",
"@radix-ui/react-toast": "^1.2.1",
"@radix-ui/react-tooltip": "^1.1.2",
"@supabase/supabase-js": "^2.45.4",
"@types/react-syntax-highlighter": "^15.5.13",
"@vercel/kv": "^2.0.0",
"class-variance-authority": "^0.7.0",
+34 -76
View File
@@ -1,9 +1,9 @@
import { VercelMemoryStore } from "@/stores/vercel";
import {
Annotation,
BaseStore,
END,
LangGraphRunnableConfig,
MessagesAnnotation,
SharedValue,
START,
StateGraph,
} from "@langchain/langgraph";
@@ -12,28 +12,17 @@ import { BaseMessage } from "@langchain/core/messages";
import { RunnableConfig } from "@langchain/core/runnables";
import { ChatAnthropic } from "@langchain/anthropic";
import { DEFAULT_SYSTEM_RULES } from "../constants";
import { UserRules } from "@/hooks/useGraph";
import { getRulesFromStore, putRulesInStore } from "./utils";
const DEFAULT_SYSTEM_RULES_STRING = `- ${DEFAULT_SYSTEM_RULES.join("\n- ")}`;
const DEFAULT_RULES_STRING = "*no rules have been set yet*";
const GraphAnnotation = Annotation.Root({
...MessagesAnnotation.spec,
/**
* Shared user rules on how to generate text.
* Use `assistant_id` so it matches the assistant,
* and can be shared across users.
*/
userRules: SharedValue.on("assistant_id"),
/**
* Whether or not writing content was generated in the conversation.
*/
contentGenerated: Annotation<boolean>(),
/**
* The user rules defined in the shared value.
* @TODO remove once api for fetching shared values is available.
*/
rules: Annotation<UserRules>(),
contentGenerated: Annotation<boolean>,
});
const GraphConfig = Annotation.Root({
@@ -46,12 +35,7 @@ const GraphConfig = Annotation.Root({
* Whether or not the user has accepted the text generated by the AI.
* If this is true, the graph will route to a node which generates rules.
*/
hasAcceptedText: Annotation<boolean>(),
/**
* Whether or not to only get the rules.
* @TODO remove once api for fetching shared values is available.
*/
onlyGetRules: Annotation<boolean>(),
hasAcceptedText: Annotation<boolean>,
});
const RULES_PROMPT = `The user has defined two sets of rules. The first set is for style guidelines, and the second set is for content guidelines.
@@ -76,35 +60,31 @@ System rules:
const callModel = async (
state: typeof GraphAnnotation.State,
config?: RunnableConfig
config: LangGraphRunnableConfig
) => {
const model = new ChatAnthropic({
model: "claude-3-5-sonnet-20240620",
temperature: 0,
});
let styleRules: string | undefined;
let contentRules: string | undefined;
if (state.userRules) {
if (state.userRules.styleRules?.length) {
styleRules = `- ${state.userRules.styleRules.join("\n - ")}`;
}
if (state.userRules.contentRules?.length) {
contentRules = `- ${state.userRules.contentRules.join("\n - ")}`;
}
}
const { styleRules, contentRules } = await getRulesFromStore(config);
const styleRulesString = styleRules ? `- ${styleRules.join("\n - ")}` : null;
const contentRulesString = contentRules
? `- ${contentRules.join("\n - ")}`
: null;
let systemPrompt = SYSTEM_PROMPT.replace(
"{systemRules}",
config?.configurable?.systemRules ?? DEFAULT_SYSTEM_RULES_STRING
);
if (styleRules || contentRules) {
if (styleRulesString || contentRulesString) {
systemPrompt = systemPrompt
.replace("{rulesPrompt}", RULES_PROMPT)
.replace("{styleRules}", styleRules || DEFAULT_RULES_STRING)
.replace("{contentRules}", contentRules || DEFAULT_RULES_STRING);
.replace("{styleRules}", styleRulesString || DEFAULT_RULES_STRING)
.replace("{contentRules}", contentRulesString || DEFAULT_RULES_STRING);
} else {
systemPrompt.replace("{rulesPrompt}", "");
systemPrompt = systemPrompt.replace("{rulesPrompt}", "");
}
const response = await model.invoke(
@@ -142,7 +122,7 @@ const _prepareConversation = (messages: BaseMessage[]): string => {
*/
const generateInsights = async (
state: typeof GraphAnnotation.State,
config?: RunnableConfig
config: LangGraphRunnableConfig
) => {
const systemPrompt = `This conversation contains back and fourth between an AI assistant, and a user who is using the assistant to generate text.
@@ -192,24 +172,22 @@ And here are the default system rules:
Respond with updated rules to keep in mind for future conversations. Try to keep the rules you list high signal-to-noise - don't include unnecessary ones, but make sure the ones you do add are descriptive. Combine ones that seem similar and/or contradictory`;
let styleRules = DEFAULT_RULES_STRING;
let contentRules = DEFAULT_RULES_STRING;
if (state.userRules) {
if (state.userRules.styleRules?.length) {
styleRules = `- ${state.userRules.styleRules.join("\n - ")}`;
}
if (state.userRules.contentRules?.length) {
contentRules = `- ${state.userRules.contentRules.join("\n - ")}`;
}
}
const { styleRules, contentRules } = await getRulesFromStore(config);
const styleRulesString = styleRules
? `- ${styleRules.join("\n - ")}`
: DEFAULT_RULES_STRING;
const contentRulesString = contentRules
? `- ${contentRules.join("\n - ")}`
: DEFAULT_RULES_STRING;
const prompt = systemPrompt
.replace(
"{systemRules}",
config?.configurable?.systemRules ?? DEFAULT_SYSTEM_RULES_STRING
config.configurable?.systemRules ?? DEFAULT_SYSTEM_RULES_STRING
)
.replace("{styleRules}", styleRules)
.replace("{contentRules}", contentRules)
.replace("{styleRules}", styleRulesString)
.replace("{contentRules}", contentRulesString)
.replace("{conversation}", _prepareConversation(state.messages));
const userRulesSchema = z.object({
@@ -240,10 +218,9 @@ Respond with updated rules to keep in mind for future conversations. Try to keep
config
);
await putRulesInStore(config, result);
return {
userRules: {
...result,
},
userAcceptedText: false,
};
};
@@ -303,12 +280,6 @@ If the assistant generated content, set 'contentGenerated' to true.
]);
};
const getRules = (state: typeof GraphAnnotation.State) => {
return {
rules: state.userRules,
};
};
/**
* Conditional edge which is always called first. This edge
* determines whether or not revisions have been made, and if so,
@@ -319,29 +290,19 @@ const shouldGenerateInsights = (
_state: typeof GraphAnnotation.State,
config?: RunnableConfig
) => {
const { hasAcceptedText, onlyGetRules } = {
hasAcceptedText: false,
onlyGetRules: false,
...(config?.configurable || {}),
};
const { hasAcceptedText = false } = config?.configurable ?? {};
if (onlyGetRules) {
return "getRules";
}
if (hasAcceptedText) {
return "generateInsights";
}
return "callModel";
};
export function buildGraph(store?: VercelMemoryStore) {
export function buildGraph() {
const workflow = new StateGraph(GraphAnnotation, GraphConfig)
.addNode("callModel", callModel)
.addNode("generateInsights", generateInsights)
.addNode("wasContentGenerated", wasContentGenerated)
// At this time there isn't a good way to fetch values from the store
// so instead we have a node which can return them.
.addNode("getRules", getRules)
// Always start by checking whether or not to generate insights
.addConditionalEdges(START, shouldGenerateInsights)
// Always check if content was generated after calling the model
@@ -349,10 +310,7 @@ export function buildGraph(store?: VercelMemoryStore) {
// No further action by the graph is necessary after either
// generating a response via `callModel`, or rules via `generateInsights`.
.addEdge("generateInsights", END)
.addEdge("wasContentGenerated", END)
.addEdge("getRules", END);
.addEdge("wasContentGenerated", END);
return workflow.compile({
store,
});
return workflow.compile();
}
+44
View File
@@ -0,0 +1,44 @@
import { createNamespace, USER_RULES_STORE_KEY } from "../lib/store";
import { UserRules } from "../types";
import { BaseStore, LangGraphRunnableConfig } from "@langchain/langgraph";
const validateStore = (config: LangGraphRunnableConfig): BaseStore => {
if (!config.store) {
throw new Error("Store not found in config.");
}
return config.store;
};
export const getRulesFromStore = async (
config: LangGraphRunnableConfig
): Promise<UserRules> => {
const store = validateStore(config);
const assistantId = config.configurable?.assistant_id;
if (!assistantId) {
throw new Error("Assistant ID not found in config.");
}
const namespace = createNamespace(assistantId);
const rules = await store.get(namespace, USER_RULES_STORE_KEY);
return {
styleRules: rules?.value?.styleRules ?? null,
contentRules: rules?.value?.contentRules ?? null,
};
};
export const putRulesInStore = async (
config: LangGraphRunnableConfig,
rules: UserRules
): Promise<void> => {
const store = validateStore(config);
const assistantId = config.configurable?.assistant_id;
if (!assistantId) {
throw new Error("Assistant ID not found in config.");
}
const namespace = createNamespace(assistantId);
await store.put(namespace, USER_RULES_STORE_KEY, rules);
};
-2
View File
@@ -1,7 +1,5 @@
import { NextRequest, NextResponse } from "next/server";
export const runtime = "edge";
function getCorsHeaders() {
return {
"Access-Control-Allow-Origin": "*",
-67
View File
@@ -1,67 +0,0 @@
// ---------------------------------------
// WARNING
// This API endpoint is no longer used. Instead, use the LangGraph Cloud endpoint.
// ---------------------------------------
import { buildGraph } from "@/agent";
import { VercelMemoryStore } from "@/stores/vercel";
import { createClient } from "@vercel/kv";
import { NextRequest, NextResponse } from "next/server";
const vercelKvClient = () => {
if (!process.env.KV_REST_API_TOKEN || !process.env.KV_REST_API_URL) {
throw new Error("Missing Vercel token or URL environment");
}
return createClient({
token: process.env.KV_REST_API_TOKEN,
url: process.env.KV_REST_API_URL,
});
};
export async function POST(req: NextRequest) {
const reqJson = await req.json();
const {
messages,
assistantId,
hasAcceptedText,
contentGenerated,
systemRules,
} = reqJson;
// Unlike in the studio, we need to pass a store here since it's not set by default.
const store = new VercelMemoryStore({
client: vercelKvClient(),
});
const graph = buildGraph(store);
const config = {
configurable: { assistant_id: assistantId },
version: "v2" as const,
};
const stream = graph.streamEvents(
{ messages, hasAcceptedText, contentGenerated, systemRules },
config
);
const encoder = new TextEncoder();
const readableStream = new ReadableStream({
async start(controller) {
try {
for await (const event of stream) {
controller.enqueue(encoder.encode(JSON.stringify(event) + "\n"));
}
} finally {
controller.close();
}
},
});
return new NextResponse(readableStream, {
headers: {
"Content-Type": "application/json",
"Transfer-Encoding": "chunked",
},
});
}
-72
View File
@@ -1,72 +0,0 @@
import { buildGraph } from "@/agent";
import { VercelMemoryStore } from "@/stores/vercel";
import {
Annotation,
END,
SharedValue,
START,
StateGraph,
} from "@langchain/langgraph";
import { createClient } from "@vercel/kv";
import { NextRequest, NextResponse } from "next/server";
const vercelKvClient = () => {
if (!process.env.KV_REST_API_TOKEN || !process.env.KV_REST_API_URL) {
throw new Error("Missing Vercel token or URL environment");
}
return createClient({
token: process.env.KV_REST_API_TOKEN,
url: process.env.KV_REST_API_URL,
});
};
const buildGetRulesGraph = (store: VercelMemoryStore) => {
const GraphAnnotation = Annotation.Root({
userRules: SharedValue.on("assistant_id"),
styleRules: Annotation<string[]>(),
contentRules: Annotation<string[]>(),
});
const getRules = (
state: typeof GraphAnnotation.State
): Partial<typeof GraphAnnotation.State> => {
if (!state.userRules) {
return {
contentRules: [],
styleRules: [],
};
}
return {
contentRules: (state.userRules.contentRules as string[]) || [],
styleRules: (state.userRules.styleRules as string[]) || [],
};
};
const workflow = new StateGraph(GraphAnnotation)
.addNode("getRules", getRules)
.addEdge(START, "getRules")
.addEdge("getRules", END);
return workflow.compile({ store });
};
export async function POST(req: NextRequest) {
const reqJson = await req.json();
const { assistantId } = reqJson;
// Unlike in the studio, we need to pass a store here since it's not set by default.
const store = new VercelMemoryStore({
client: vercelKvClient(),
});
const graph = buildGetRulesGraph(store);
const config = { configurable: { assistant_id: assistantId } };
const result = await graph.invoke({}, config);
return new NextResponse(JSON.stringify(result), {
headers: {
"Content-Type": "application/json",
},
});
}
+67
View File
@@ -0,0 +1,67 @@
import { Client } from "@langchain/langgraph-sdk";
import { NextRequest, NextResponse } from "next/server";
export async function GET(req: NextRequest) {
if (!process.env.LANGGRAPH_API_URL || !process.env.LANGCHAIN_API_KEY) {
return new NextResponse(
JSON.stringify({
error: "LANGGRAPH_API_URL and LANGCHAIN_API_KEY must be set",
}),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
const searchParams = req.nextUrl.searchParams;
const namespaceParam = searchParams.get("namespace");
const key = searchParams.get("key");
if (!namespaceParam || !key) {
return new NextResponse(
JSON.stringify({ error: "Missing namespace or key" }),
{
status: 400,
headers: { "Content-Type": "application/json" },
}
);
}
// Parse the namespace from URL-encoded string to an array of strings
const namespace: string = decodeURIComponent(namespaceParam);
const namespaceArr: string[] = namespace.split(".");
if (!Array.isArray(namespaceArr)) {
return new NextResponse(
JSON.stringify({ error: "Invalid namespace format" }),
{
status: 400,
headers: { "Content-Type": "application/json" },
}
);
}
const lgClient = new Client({
apiKey: process.env.LANGCHAIN_API_KEY,
apiUrl: process.env.LANGGRAPH_API_URL,
});
try {
const result = await lgClient.store.getItem(namespaceArr, key);
return new NextResponse(JSON.stringify(result ?? {}), {
headers: { "Content-Type": "application/json" },
});
} catch (e) {
console.error("Err fetching store");
console.error(e);
return new NextResponse(
JSON.stringify({ error: "Failed to get item from store" }),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
}
+85
View File
@@ -0,0 +1,85 @@
import { createClient } from "@supabase/supabase-js";
import { NextRequest, NextResponse } from "next/server";
export async function GET(req: NextRequest) {
if (
!process.env.SUPABASE_SERVICE_ROLE_KEY ||
!process.env.NEXT_PUBLIC_SUPABASE_URL
) {
return new NextResponse(
JSON.stringify({
error:
"SUPABASE_SERVICE_ROLE_KEY and NEXT_PUBLIC_SUPABASE_URL must be set",
}),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
// Initialize Supabase client
const supabase = createClient(
process.env.NEXT_PUBLIC_SUPABASE_URL,
process.env.SUPABASE_SERVICE_ROLE_KEY
);
const searchParams = req.nextUrl.searchParams;
const userId = searchParams.get("userId");
const assistantId = searchParams.get("assistantId");
if (!userId || !assistantId) {
return new NextResponse(
JSON.stringify({ error: "Missing userId or assistantId" }),
{
status: 400,
headers: { "Content-Type": "application/json" },
}
);
}
try {
// Fetch the latest system rules
const { data, error } = await supabase
.from("user_rules")
.select("system_rules")
.eq("user_id", userId)
.eq("assistant_id", assistantId)
.limit(1)
.single();
if (error) {
console.error("Error getting system rules:", {
error,
});
return new NextResponse(
JSON.stringify({ error: "Failed to get system rules." }),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
if (!data) {
return new NextResponse(JSON.stringify({ error: "No rules found" }), {
status: 404,
headers: { "Content-Type": "application/json" },
});
}
return new NextResponse(JSON.stringify(data), {
status: 200,
headers: { "Content-Type": "application/json" },
});
} catch (error) {
console.error("Error fetching system rules:", error);
return new NextResponse(
JSON.stringify({ error: "Failed to fetch system rules" }),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
}
+82
View File
@@ -0,0 +1,82 @@
import { NextRequest, NextResponse } from "next/server";
import { createClient } from "@supabase/supabase-js";
export async function POST(req: NextRequest) {
if (
!process.env.SUPABASE_SERVICE_ROLE_KEY ||
!process.env.NEXT_PUBLIC_SUPABASE_URL
) {
return new NextResponse(
JSON.stringify({
error:
"SUPABASE_SERVICE_ROLE_KEY and NEXT_PUBLIC_SUPABASE_URL must be set",
}),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
// Initialize Supabase client
const supabase = createClient(
process.env.NEXT_PUBLIC_SUPABASE_URL,
process.env.SUPABASE_SERVICE_ROLE_KEY
);
const { assistantId, userId, systemRules } = await req.json();
if (!userId || !assistantId || !systemRules) {
return new NextResponse(
JSON.stringify({
error: "Missing userId, assistantId, or an array of systemRules.",
}),
{
status: 400,
headers: { "Content-Type": "application/json" },
}
);
}
try {
// Insert new row into user_rules table
const { data, error } = await supabase
.from("user_rules")
.upsert(
{
user_id: userId,
assistant_id: assistantId,
system_rules: systemRules,
},
{ onConflict: "user_id,assistant_id", ignoreDuplicates: false }
)
.select();
if (error) {
console.error("Error inserting system rules:", {
error,
});
return new NextResponse(
JSON.stringify({ error: "Failed to insert system rules." }),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
return new NextResponse(JSON.stringify(data), {
status: 200,
headers: { "Content-Type": "application/json" },
});
} catch (error) {
console.error("Error inserting system rules:", error);
return new NextResponse(
JSON.stringify({ error: "Failed to insert system rules." }),
{
status: 500,
headers: { "Content-Type": "application/json" },
}
);
}
}
-88
View File
@@ -1,88 +0,0 @@
import { initVercelStore, VercelMemoryStore } from "@/stores/vercel";
import { NextRequest, NextResponse } from "next/server";
export const runtime = "edge";
function getCorsHeaders() {
return {
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Methods": "GET, POST, PUT, PATCH, DELETE, OPTIONS",
"Access-Control-Allow-Headers": "*",
};
}
const NAMESPACE = "system_rules";
async function handleSetRules(
store: VercelMemoryStore,
fields: { systemRules: string; assistantId: string }
): Promise<void> {
const key = fields.assistantId;
const values = {
systemRules: fields.systemRules,
};
const input: Array<[string, string, Record<string, any> | null]> = [
[NAMESPACE, key, values],
];
await store.put(input);
}
async function handleGetRules(
store: VercelMemoryStore,
fields: { assistantId: string }
): Promise<{ systemRules: string | null }> {
const results = await store.list([NAMESPACE]);
if (results && results[NAMESPACE] && results[NAMESPACE][fields.assistantId]) {
return {
systemRules: results[NAMESPACE][fields.assistantId].systemRules,
};
}
return {
systemRules: null,
};
}
async function handleRequest(req: NextRequest, method: string) {
try {
const store = initVercelStore();
if (method === "POST") {
const { systemRules, assistantId } = await req.json();
await handleSetRules(store, { systemRules, assistantId });
return NextResponse.json({ success: true }, { status: 200 });
} else if (method === "GET") {
// GET reqs can not have a body, so we need to parse the query string to get the assistantId
const url = new URL(req.url);
const assistantId = url.searchParams.get("assistantId");
if (!assistantId) {
return NextResponse.json(
{ error: "assistantId is required" },
{ status: 400 }
);
}
const systemRules = await handleGetRules(store, { assistantId });
return NextResponse.json(systemRules, {
status: 200,
});
} else {
return new NextResponse(null, { status: 405 });
}
} catch (e: any) {
return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
}
}
export const GET = (req: NextRequest) => handleRequest(req, "GET");
export const POST = (req: NextRequest) => handleRequest(req, "POST");
// Add a new OPTIONS handler
export const OPTIONS = () => {
return new NextResponse(null, {
status: 204,
headers: {
...getCorsHeaders(),
},
});
};
+5 -4
View File
@@ -23,15 +23,16 @@ export default function Home() {
isGetAssistantsLoading,
getAssistantsByUserId,
updateAssistantMetadata,
userRules,
isLoadingUserRules,
} = useGraph({ userId, refreshAssistants });
const {
userRules,
isLoadingUserRules,
setSystemRules,
systemRules,
setSystemRulesAndSave,
isLoadingSystemRules,
} = useRules(assistantId);
getUserRules,
} = useRules({ assistantId, userId });
return (
<main className="h-screen">
@@ -59,7 +60,7 @@ export default function Home() {
<ContentComposerChatInterface
createAssistant={createAssistant}
systemRules={systemRules}
sendMessage={sendMessage}
sendMessage={async (params) => sendMessage(params, getUserRules)}
streamMessage={streamMessage}
userId={userId}
/>
+1 -1
View File
@@ -8,8 +8,8 @@ import {
DialogTrigger,
} from "./ui/dialog";
import { Button } from "./ui/button";
import { UserRules } from "@/hooks/useGraph";
import { Loader } from "lucide-react";
import { UserRules } from "@/types";
export interface GeneratedRulesProps {
isLoadingUserRules: boolean;
+5 -73
View File
@@ -11,11 +11,6 @@ export interface GraphInput {
systemRules: string | undefined;
}
export interface UserRules {
styleRules?: string[];
contentRules?: string[];
}
export interface UseGraphInput {
userId: string | undefined;
refreshAssistants: () => Promise<void>;
@@ -26,8 +21,6 @@ export function useGraph(input: UseGraphInput) {
const [threadId, setThreadId] = useState<string>();
const [assistantId, setAssistantId] = useState<string>();
const [isGetAssistantsLoading, setIsGetAssistantsLoading] = useState(false);
const [isLoadingUserRules, setIsLoadingUserRules] = useState(false);
const [userRules, setUserRules] = useState<UserRules | undefined>();
useEffect(() => {
if (typeof window === "undefined") return;
@@ -55,25 +48,6 @@ export function useGraph(input: UseGraphInput) {
}
}, [input.userId]);
// TODO: remove after a couple days when all existing users have been updated.
useEffect(() => {
if (typeof window === "undefined") return;
if (!input.userId) return;
void ensureAssistantIsTiedToUser(input.userId);
}, [assistantId, input.userId]);
useEffect(() => {
if (!assistantId) return;
const fetchRules = async () => {
if (!userRules) {
await getUserRules();
}
};
void fetchRules();
}, [assistantId]);
const createAssistant = async (
graphId: string,
userId: string,
@@ -126,7 +100,10 @@ export function useGraph(input: UseGraphInput) {
});
};
const sendMessage = async (params: GraphInput) => {
const sendMessage = async (
params: GraphInput,
getRulesCallback: () => Promise<void>
) => {
const { messages, hasAcceptedText, contentGenerated, systemRules } = params;
if (!assistantId) {
throw new Error("Assistant ID is required");
@@ -149,7 +126,7 @@ export function useGraph(input: UseGraphInput) {
if (hasAcceptedText) {
// Do not await so it is not blocking
getUserRules().catch((_) => {
getRulesCallback().catch((_) => {
toast({
title: "Failed to re-fetch user rules.",
description: "Please refresh the page to see the updated rules.",
@@ -192,49 +169,6 @@ export function useGraph(input: UseGraphInput) {
return updatedAssistant;
};
const ensureAssistantIsTiedToUser = async (userId: string) => {
if (!assistantId || getCookie(USER_TIED_TO_ASSISTANT) === "true") return;
const client = createClient();
const currentAssistant = await client.assistants.get(assistantId);
if (
currentAssistant.metadata &&
"userId" in currentAssistant.metadata &&
currentAssistant.metadata.userId === userId
) {
setCookie(USER_TIED_TO_ASSISTANT, "true");
return;
}
// Update assistant metadata to include userId
await updateAssistantMetadata(assistantId, {
metadata: {
...currentAssistant.metadata,
userId,
},
});
setCookie(USER_TIED_TO_ASSISTANT, "true");
};
const getUserRules = async () => {
if (!assistantId || assistantId === "") return;
setIsLoadingUserRules(true);
const client = createClient();
try {
const response = await client.runs.wait(null, assistantId, {
input: {},
config: { configurable: { onlyGetRules: true } },
});
const { rules } = response as Record<string, any>;
if (rules?.styleRules?.length || rules?.contentRules?.length) {
setUserRules(rules);
}
} finally {
setIsLoadingUserRules(false);
}
};
return {
assistantId,
setAssistantId: updateAssistant,
@@ -244,7 +178,5 @@ export function useGraph(input: UseGraphInput) {
isGetAssistantsLoading,
getAssistantsByUserId,
updateAssistantMetadata,
userRules,
isLoadingUserRules,
};
}
+68 -9
View File
@@ -1,12 +1,23 @@
import { DEFAULT_SYSTEM_RULES } from "@/constants";
import { createNamespace, USER_RULES_STORE_KEY } from "@/lib/store";
import { UserRules } from "@/types";
import { useState, useEffect } from "react";
import { useToast } from "./use-toast";
const DEFAULT_SYSTEM_RULES_STRING = `- ${DEFAULT_SYSTEM_RULES.join("\n- ")}`;
export function useRules(assistantId: string | undefined) {
export interface UseRulesInput {
assistantId: string | undefined;
userId: string | undefined;
}
export function useRules({ assistantId, userId }: UseRulesInput) {
const { toast } = useToast();
const [systemRules, setSystemRules] = useState<string>();
const [isLoadingSystemRules, setIsLoadingSystemRules] = useState(false);
const [isSavingSystemRules, setIsSavingSystemRules] = useState(false);
const [isLoadingUserRules, setIsLoadingUserRules] = useState(false);
const [userRules, setUserRules] = useState<UserRules | undefined>();
useEffect(() => {
if (!assistantId) return;
@@ -15,18 +26,59 @@ export function useRules(assistantId: string | undefined) {
if (!systemRules) {
await getSystemRules();
}
if (!userRules) {
await getUserRules();
}
};
void fetchRules();
}, [assistantId]);
const getUserRules = async (): Promise<void> => {
if (!assistantId) return;
setIsLoadingUserRules(true);
try {
const namespace = encodeURIComponent(
createNamespace(assistantId).join(".")
);
const queryParams = new URLSearchParams({
namespace,
key: USER_RULES_STORE_KEY,
});
const fullUrl = `/api/store/get?${queryParams.toString()}`;
const response = await fetch(fullUrl);
if (!response.ok) {
toast({
title: "An error occurred fetching user rules",
});
return;
}
const rules = await response.json();
if (!rules || !rules.value) {
// Successfully hit API, no rules yet stored.
// no-op
return;
}
setUserRules(rules.value);
} catch (e) {
toast({
title: "An error occurred fetching user rules",
});
} finally {
setIsLoadingUserRules(false);
}
};
const getSystemRules = async () => {
if (!assistantId || assistantId === "") return;
if (!assistantId || assistantId === "" || !userId || userId === "") return;
setIsLoadingSystemRules(true);
try {
const queryParams = new URLSearchParams({ assistantId });
const fullUrl = `/api/system_rules?${queryParams.toString()}`;
const queryParams = new URLSearchParams({ assistantId, userId });
const fullUrl = `/api/system_rules/get?${queryParams.toString()}`;
const response = await fetch(fullUrl);
if (!response.ok) {
@@ -34,8 +86,8 @@ export function useRules(assistantId: string | undefined) {
}
const data = await response.json();
if (data?.systemRules) {
setSystemRules(data.systemRules);
if (data?.system_rules) {
setSystemRules(data.system_rules);
} else {
setSystemRules(DEFAULT_SYSTEM_RULES_STRING);
}
@@ -45,17 +97,21 @@ export function useRules(assistantId: string | undefined) {
};
const setSystemRulesAndSave = async (newSystemRules: string) => {
if (!assistantId || assistantId === "") return;
if (!assistantId || assistantId === "" || !userId || userId === "") return;
setIsSavingSystemRules(true);
try {
setSystemRules(newSystemRules);
await fetch("/api/system_rules", {
await fetch("/api/system_rules/put", {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({ assistantId, systemRules: newSystemRules }),
body: JSON.stringify({
assistantId,
userId,
systemRules: newSystemRules,
}),
});
} finally {
setIsSavingSystemRules(false);
@@ -69,5 +125,8 @@ export function useRules(assistantId: string | undefined) {
systemRules,
isLoadingSystemRules,
isSavingSystemRules,
userRules,
isLoadingUserRules,
getUserRules,
};
}
+5
View File
@@ -0,0 +1,5 @@
export const USER_RULES_STORE_KEY = "rules";
export const createNamespace = (assistantId: string) => {
return ["assistant_id", assistantId, "userRules"];
};
-62
View File
@@ -1,62 +0,0 @@
import { BaseStore, type Values } from "@langchain/langgraph";
import { createClient, kv, type VercelKV } from "@vercel/kv";
export class VercelMemoryStore extends BaseStore {
protected client: VercelKV;
constructor(fields?: { client?: VercelKV }) {
super();
this.client = fields?.client || kv;
}
async list(
prefixes: string[]
): Promise<Record<string, Record<string, Values>>> {
const result: Record<string, Record<string, Values>> = {};
for (const prefix of prefixes) {
const keys = await this.client.keys(`${prefix}:*`);
result[prefix] = {};
for (const fullKey of keys) {
const value = await this.client.get<string>(fullKey);
if (value !== null) {
// Get the last part of the key. This represents the key of the shared value object the user has set.
const items = fullKey.split(":");
const key = items[items.length - 1];
result[prefix][key] =
typeof value === "string" ? JSON.parse(value) : value;
}
}
}
return result;
}
async put(writes: Array<[string, string, Values | null]>): Promise<void> {
const pipeline = this.client.pipeline();
for (const [namespace, key, value] of writes) {
const fullKey = `${namespace}:${key}`;
if (value === null) {
pipeline.del(fullKey);
} else {
pipeline.set(fullKey, JSON.stringify(value));
}
}
await pipeline.exec();
}
}
export function initVercelStore() {
if (!process.env.KV_REST_API_TOKEN || !process.env.KV_REST_API_URL) {
throw new Error("Missing Vercel token or URL environment");
}
return new VercelMemoryStore({
client: createClient({
token: process.env.KV_REST_API_TOKEN,
url: process.env.KV_REST_API_URL,
}),
});
}
+5
View File
@@ -14,3 +14,8 @@ export interface ToolCall {
}
export type Model = "gpt-4o-mini" | string; // Add other model options as needed
export type UserRules = {
styleRules: string[];
contentRules: string[];
};
+87 -13
View File
@@ -233,29 +233,29 @@
zod "^3.22.4"
zod-to-json-schema "^3.22.3"
"@langchain/langgraph-checkpoint@~0.0.6":
version "0.0.6"
resolved "https://registry.yarnpkg.com/@langchain/langgraph-checkpoint/-/langgraph-checkpoint-0.0.6.tgz#69f0c5c9aeefd48dcf0fa1ffa0744d8139a9f27d"
integrity sha512-hQsznlUMFKyOCaN9VtqNSSemfKATujNy5ePM6NX7lruk/Mmi2t7R9SsBnf9G2Yts+IaIwv3vJJaAFYEHfqbc5g==
"@langchain/langgraph-checkpoint@~0.0.9":
version "0.0.9"
resolved "https://registry.yarnpkg.com/@langchain/langgraph-checkpoint/-/langgraph-checkpoint-0.0.9.tgz#fdeb7654b112831161093d2867323e0450706ad8"
integrity sha512-9KrTxnKqTCRDxYOsvQ4UOuM878S1Sp4ZUejfGBdZc9yaWGzRGV4aEYJGt8GDSBwBUYd7gz2gNi+q4xtxvwIZig==
dependencies:
uuid "^10.0.0"
"@langchain/langgraph-sdk@^0.0.11":
version "0.0.11"
resolved "https://registry.yarnpkg.com/@langchain/langgraph-sdk/-/langgraph-sdk-0.0.11.tgz#d51ca65489f2d208bf32493bf33f89e7c3df5498"
integrity sha512-Bos40RcYoXP5eONBxti/IKp4t/6egQjq84mxU4hN0gvuqIfPakguBPcQZNnVhb9KSHbGa7bpHbxXRYVW8iDwPg==
"@langchain/langgraph-sdk@^0.0.14":
version "0.0.14"
resolved "https://registry.yarnpkg.com/@langchain/langgraph-sdk/-/langgraph-sdk-0.0.14.tgz#aae3495208f6bcc2438f7cd6616b21a0dfa91e6f"
integrity sha512-hDu5Q92px6M3frZbKPOg2jWb8cCxU83oEt+GtfOY0MzID60+XocjsHdwSv5EEj32X9yzINGq6jHlHg1EHqjZyA==
dependencies:
"@types/json-schema" "^7.0.15"
p-queue "^6.6.2"
p-retry "4"
uuid "^9.0.0"
"@langchain/langgraph@^0.2.8":
version "0.2.8"
resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.2.8.tgz#9606982686ee857064a217dc5599ebdbc9aaf2fe"
integrity sha512-sQ3NqwZzdvILeiYQQCDCBFj+FLd3oBfg2sxMo3e5g7vd5+zd/hpK5+JRTHbsMZte0PTAlTbQ5YbfCC2D6K9AVw==
"@langchain/langgraph@^0.2.10":
version "0.2.10"
resolved "https://registry.yarnpkg.com/@langchain/langgraph/-/langgraph-0.2.10.tgz#f026b132c5e5e01bd91803bd0124c94bdad6076a"
integrity sha512-xkWkcpngcpz32nglzQyX2SyS00I+h2Ao2XjTxtONCmPJR6N4HCXJzgpEekpvUeCN1XtggWGBMWNBb75Z8CgGfQ==
dependencies:
"@langchain/langgraph-checkpoint" "~0.0.6"
"@langchain/langgraph-checkpoint" "~0.0.9"
double-ended-queue "^2.1.0-0"
uuid "^10.0.0"
zod "^3.23.8"
@@ -712,6 +712,63 @@
resolved "https://registry.yarnpkg.com/@rushstack/eslint-patch/-/eslint-patch-1.10.4.tgz#427d5549943a9c6fce808e39ea64dbe60d4047f1"
integrity sha512-WJgX9nzTqknM393q1QJDJmoW28kUfEnybeTfVNcNAPnIx210RXm2DiXiHzfNPJNIUUb1tJnz/l4QGtJ30PgWmA==
"@supabase/auth-js@2.65.0":
version "2.65.0"
resolved "https://registry.yarnpkg.com/@supabase/auth-js/-/auth-js-2.65.0.tgz#e345c492f8cbc31cd6289968eae0e349ff0f39e9"
integrity sha512-+wboHfZufAE2Y612OsKeVP4rVOeGZzzMLD/Ac3HrTQkkY4qXNjI6Af9gtmxwccE5nFvTiF114FEbIQ1hRq5uUw==
dependencies:
"@supabase/node-fetch" "^2.6.14"
"@supabase/functions-js@2.4.1":
version "2.4.1"
resolved "https://registry.yarnpkg.com/@supabase/functions-js/-/functions-js-2.4.1.tgz#373e75f8d3453bacd71fb64f88d7a341d7b53ad7"
integrity sha512-8sZ2ibwHlf+WkHDUZJUXqqmPvWQ3UHN0W30behOJngVh/qHHekhJLCFbh0AjkE9/FqqXtf9eoVvmYgfCLk5tNA==
dependencies:
"@supabase/node-fetch" "^2.6.14"
"@supabase/node-fetch@2.6.15", "@supabase/node-fetch@^2.6.14":
version "2.6.15"
resolved "https://registry.yarnpkg.com/@supabase/node-fetch/-/node-fetch-2.6.15.tgz#731271430e276983191930816303c44159e7226c"
integrity sha512-1ibVeYUacxWYi9i0cf5efil6adJ9WRyZBLivgjs+AUpewx1F3xPi7gLgaASI2SmIQxPoCEjAsLAzKPgMJVgOUQ==
dependencies:
whatwg-url "^5.0.0"
"@supabase/postgrest-js@1.16.1":
version "1.16.1"
resolved "https://registry.yarnpkg.com/@supabase/postgrest-js/-/postgrest-js-1.16.1.tgz#68dfa0581d8ae4296378cb8815bbde3f4602aef5"
integrity sha512-EOSEZFm5pPuCPGCmLF1VOCS78DfkSz600PBuvBND/IZmMciJ1pmsS3ss6TkB6UkuvTybYiBh7gKOYyxoEO3USA==
dependencies:
"@supabase/node-fetch" "^2.6.14"
"@supabase/realtime-js@2.10.2":
version "2.10.2"
resolved "https://registry.yarnpkg.com/@supabase/realtime-js/-/realtime-js-2.10.2.tgz#c2b42d17d723d2d2a9146cfad61dc3df1ce3127e"
integrity sha512-qyCQaNg90HmJstsvr2aJNxK2zgoKh9ZZA8oqb7UT2LCh3mj9zpa3Iwu167AuyNxsxrUE8eEJ2yH6wLCij4EApA==
dependencies:
"@supabase/node-fetch" "^2.6.14"
"@types/phoenix" "^1.5.4"
"@types/ws" "^8.5.10"
ws "^8.14.2"
"@supabase/storage-js@2.7.0":
version "2.7.0"
resolved "https://registry.yarnpkg.com/@supabase/storage-js/-/storage-js-2.7.0.tgz#9ff322d2c3b141087aa34115cf14205e4980ce75"
integrity sha512-iZenEdO6Mx9iTR6T7wC7sk6KKsoDPLq8rdu5VRy7+JiT1i8fnqfcOr6mfF2Eaqky9VQzhP8zZKQYjzozB65Rig==
dependencies:
"@supabase/node-fetch" "^2.6.14"
"@supabase/supabase-js@^2.45.4":
version "2.45.4"
resolved "https://registry.yarnpkg.com/@supabase/supabase-js/-/supabase-js-2.45.4.tgz#0bcf8722f1732dfe3e4c5190d23e3938dcc689c3"
integrity sha512-E5p8/zOLaQ3a462MZnmnz03CrduA5ySH9hZyL03Y+QZLIOO4/Gs8Rdy4ZCKDHsN7x0xdanVEWWFN3pJFQr9/hg==
dependencies:
"@supabase/auth-js" "2.65.0"
"@supabase/functions-js" "2.4.1"
"@supabase/node-fetch" "2.6.15"
"@supabase/postgrest-js" "1.16.1"
"@supabase/realtime-js" "2.10.2"
"@supabase/storage-js" "2.7.0"
"@swc/counter@^0.1.3":
version "0.1.3"
resolved "https://registry.yarnpkg.com/@swc/counter/-/counter-0.1.3.tgz#cc7463bd02949611c6329596fccd2b0ec782b0e9"
@@ -819,6 +876,11 @@
dependencies:
undici-types "~6.19.2"
"@types/phoenix@^1.5.4":
version "1.6.5"
resolved "https://registry.yarnpkg.com/@types/phoenix/-/phoenix-1.6.5.tgz#5654e14ec7ad25334a157a20015996b6d7d2075e"
integrity sha512-xegpDuR+z0UqG9fwHqNoy3rI7JDlvaPh2TY47Fl80oq6g+hXT+c/LEuE43X48clZ6lOfANl5WrPur9fYO1RJ/w==
"@types/prop-types@*":
version "15.7.12"
resolved "https://registry.yarnpkg.com/@types/prop-types/-/prop-types-15.7.12.tgz#12bb1e2be27293c1406acb6af1c3f3a1481d98c6"
@@ -866,6 +928,13 @@
resolved "https://registry.yarnpkg.com/@types/uuid/-/uuid-10.0.0.tgz#e9c07fe50da0f53dc24970cca94d619ff03f6f6d"
integrity sha512-7gqG38EyHgyP1S+7+xomFtL+ZNHcKv6DwNaCZmJmo1vgMugyF3TCnXVg4t1uk89mLNwnLtnY3TpOpCOyp1/xHQ==
"@types/ws@^8.5.10":
version "8.5.12"
resolved "https://registry.yarnpkg.com/@types/ws/-/ws-8.5.12.tgz#619475fe98f35ccca2a2f6c137702d85ec247b7e"
integrity sha512-3tPRkv1EtkDpzlgyKyI8pGsGZAGPEaXeu0DOj5DI25Ja91bdAYddYHbADRYVrZMRbfW+1l5YwXVDKohDJNQxkQ==
dependencies:
"@types/node" "*"
"@typescript-eslint/parser@^5.4.2 || ^6.0.0 || 7.0.0 - 7.2.0":
version "7.2.0"
resolved "https://registry.yarnpkg.com/@typescript-eslint/parser/-/parser-7.2.0.tgz#44356312aea8852a3a82deebdacd52ba614ec07a"
@@ -5003,6 +5072,11 @@ wrappy@1:
resolved "https://registry.yarnpkg.com/wrappy/-/wrappy-1.0.2.tgz#b5243d8f3ec1aa35f1364605bc0d1036e30ab69f"
integrity sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==
ws@^8.14.2:
version "8.18.0"
resolved "https://registry.yarnpkg.com/ws/-/ws-8.18.0.tgz#0d7505a6eafe2b0e712d232b42279f53bc289bbc"
integrity sha512-8VbfWfHLbbwu3+N6OKsOMpBdT4kXPDDB9cJk2bJ6mh9ucxdlnNvH1e+roYkKmN9Nxw2yjz7VzeO9oOz2zJ04Pw==
xtend@^4.0.0:
version "4.0.2"
resolved "https://registry.yarnpkg.com/xtend/-/xtend-4.0.2.tgz#bb72779f5fa465186b1f438f674fa347fdb5db54"