chore(api): config sqlite history

nicobytes · Oct 27, 2024 · 6cc67ea · 6cc67ea
1 parent 74f8105
commit 6cc67ea
Show file tree

Hide file tree

Showing 8 changed files with 273 additions and 8 deletions.
diff --git a/apps/api/package.json b/apps/api/package.json
@@ -10,6 +10,7 @@
   "dependencies": {
     "@langchain/core": "^0.3.15",
     "@langchain/langgraph": "^0.2.18",
+    "@langchain/langgraph-checkpoint-sqlite": "^0.1.2",
     "@langchain/mistralai": "^0.1.1",
     "@langchain/openai": "^0.3.11",
     "hono": "^4.6.7",

diff --git a/apps/api/src/bindings.ts b/apps/api/src/bindings.ts
@@ -1,5 +1,4 @@
 export type Bindings = {
-	DB_STORE: D1Database;
 	JWT_SECRET: string;
 	CLOUDFLARE_ACCOUNT_ID: string;
 	CLOUDFLARE_API_TOKEN: string;
@@ -11,4 +10,6 @@ export type Bindings = {
 	R2_URL: string;
 	FB_TOKEN: string;
 	FB_VERIFY_TOKEN: string;
+	HISTORY_KV: KVNamespace;
+	HISTORY_DB: D1Database;
 };
diff --git a/apps/api/src/chatbot/graph.state.ts b/apps/api/src/chatbot/graph.state.ts
@@ -4,6 +4,7 @@ import { MyNodes } from "./nodes";
 export const GraphState = Annotation.Root({
 	...MessagesAnnotation.spec,
 	lastAgent: Annotation<MyNodes>,
+	isReadyToBook: Annotation<boolean>,
 });
 
 export type GraphState = typeof GraphState.State;
diff --git a/apps/api/src/chatbot/graph.ts b/apps/api/src/chatbot/graph.ts
@@ -1,5 +1,5 @@
 import { END, START, StateGraph } from "@langchain/langgraph";
-import { MemorySaver } from "@langchain/langgraph";
+import { SqliteSaver } from "@langchain/langgraph-checkpoint-sqlite";
 
 import { GraphState } from "./graph.state";
 
@@ -20,8 +20,8 @@ interface Props {
 	mistralKey: string;
 }
 
-export const createGraph = (data: Props) => {
-	const memory = new MemorySaver();
+export const createGraph = (data: Props, db: D1Database) => {
+	const memory = new SqliteSaver(db);
 
 	const llmGpt4 = models.gpt4(data.openAIKey);
 	const llmMistral = models.mistral(data.mistralKey);

diff --git a/apps/api/src/chatbot/index.ts b/apps/api/src/chatbot/index.ts
@@ -24,6 +24,7 @@ app.post("/webhook", async (c) => {
 	const fbToken = c.env.FB_TOKEN;
 	const openAIKey = c.env.OPENAI_API_KEY;
 	const mistralKey = c.env.MISTRAL_API_KEY;
+	const d1 = c.env.HISTORY_DB;
 
 	const body = await c.req.json();
 	const message = body.entry?.[0]?.changes[0]?.value?.messages?.[0];
@@ -38,7 +39,7 @@ app.post("/webhook", async (c) => {
 		headers.append("Authorization", `Bearer ${fbToken}`);
 		headers.append("Content-Type", "application/json");
 
-		const agent = createGraph({ openAIKey, mistralKey });
+		const agent = createGraph({ openAIKey, mistralKey }, d1);
 		const userMessage = message.text.body;
 		const threadId = `thread_${message.from}`;
 		const agentResponse = await agent.invoke(

diff --git a/apps/api/src/chatbot/nodes/booking.node.ts b/apps/api/src/chatbot/nodes/booking.node.ts
@@ -38,6 +38,7 @@ export const bookingNode = (llm: ChatOpenAI) => {
 		return {
 			messages: response,
 			lastAgent: MyNodes.BOOKING,
+			isReadyToBook: true,
 		};
 	};
 };
diff --git a/apps/api/wrangler.toml b/apps/api/wrangler.toml
@@ -11,4 +11,13 @@ binding = "AI"
 
 [observability]
 enabled = true
-head_sampling_rate = 1
+head_sampling_rate = 1
+
+[[kv_namespaces]]
+binding = "HISTORY_KV"
+id = "efd83beca3174ee2ba9bdeaf99f73b4d"
+
+[[d1_databases]]
+binding = "HISTORY_DB"
+database_name = "prod-d1-history"
+database_id = "482cf989-efd8-45ae-a3cd-5e5891a0e573"