Skip to content

Commit deed8e5

Browse files
hoshinotsuyoshihoshinotsuyoshi
committed
poc(chat): introduce LangGraph-based pipeline for schema build mode
* Replaced the previous agent-based approach with a LangGraph workflow (`runChat`) for the `'build'` mode * Refined the system prompt to consistently guide the LLM toward proper RFC 6902 JSON Patch output, wrapped in \`\`\`json fences * Improved the prompt structure to allow informative commentary outside the JSON Patch block * Added fallback mechanism for non-schema-editing messages (i.e., omit JSON Patch) * Introduced unit tests for LangGraph flow to ensure format correctness and schema alignment * Updated dependencies to include `@langchain/langgraph`
1 parent b34315b commit deed8e5

File tree

7 files changed

+373
-29
lines changed

7 files changed

+373
-29
lines changed

‎frontend/apps/app/app/api/chat/route.ts‎

Lines changed: 11 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import { convertSchemaToText } from '@/app/lib/schema/convertSchemaToText'
22
import { isSchemaUpdated } from '@/app/lib/vectorstore/supabaseVectorStore'
33
import { syncSchemaVectorStore } from '@/app/lib/vectorstore/syncSchemaVectorStore'
4-
import { mastra } from '@/lib/mastra'
4+
import { runChat } from '@/lib/chat/langGraph'
55
import * as Sentry from '@sentry/nextjs'
66
import { NextResponse } from 'next/server'
77

88
export async function POST(request: Request) {
9-
const { message, schemaData, history, mode, projectId } = await request.json()
9+
const { message, schemaData, mode, projectId, history } = await request.json()
1010

1111
if (!message || typeof message !== 'string' || !message.trim()) {
1212
return NextResponse.json({ error: 'Message is required' }, { status: 400 })
@@ -19,9 +19,6 @@ export async function POST(request: Request) {
1919
)
2020
}
2121

22-
// Determine which agent to use based on the mode
23-
const agentName =
24-
mode === 'build' ? 'databaseSchemaBuildAgent' : 'databaseSchemaAskAgent'
2522
try {
2623
// Check if schema has been updated
2724
const schemaUpdated = await isSchemaUpdated(schemaData)
@@ -53,31 +50,17 @@ export async function POST(request: Request) {
5350
// Convert schema to text
5451
const schemaText = convertSchemaToText(schemaData)
5552

56-
// Get the agent from Mastra
57-
const agent = mastra.getAgent(agentName)
58-
if (!agent) {
59-
throw new Error(`${agentName} not found in Mastra instance`)
53+
// Use LangGraph pipeline for build mode, fallback to original for ask mode
54+
let responseText: string | undefined
55+
if (mode === 'build') {
56+
responseText = await runChat(message, schemaText, formattedChatHistory)
57+
} else {
58+
// For ask mode, we'll keep the original implementation for now
59+
// This can be refactored later to use a separate LangGraph pipeline
60+
throw new Error('Ask mode not yet implemented with LangGraph')
6061
}
6162

62-
// Create a response using the agent
63-
const response = await agent.generate([
64-
{
65-
role: 'system',
66-
content: `
67-
Complete Schema Information:
68-
${schemaText}
69-
70-
Previous conversation:
71-
${formattedChatHistory}
72-
`,
73-
},
74-
{
75-
role: 'user',
76-
content: message,
77-
},
78-
])
79-
80-
return new Response(response.text, {
63+
return new Response(responseText, {
8164
headers: {
8265
'Content-Type': 'text/plain; charset=utf-8',
8366
},

‎frontend/apps/app/components/ChatbotButton/components/ChatbotDialog/ChatbotDialog.tsx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ export const ChatbotDialog: FC<ChatbotDialogProps> = ({
9090
tableGroups,
9191
history,
9292
projectId,
93+
mode: 'build',
9394
}),
9495
})
9596

‎frontend/apps/app/lib/chat/__tests__/langGraph.test.ts‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
import { beforeEach, describe, expect, it, vi } from 'vitest'
2+
import { runChat } from '../langGraph'
3+
4+
// Mock the dependencies
5+
vi.mock('@/lib/mastra', () => ({
6+
mastra: {
7+
getAgent: vi.fn(() => ({
8+
generate: vi.fn(() =>
9+
Promise.resolve({
10+
text: `Added a summary column to track design session outcomes!
11+
12+
\`\`\`json
13+
[
14+
{
15+
"op": "add",
16+
"path": "/tables/design_sessions/columns/summary",
17+
"value": {
18+
"name": "summary",
19+
"type": "text",
20+
"not_null": false
21+
}
22+
}
23+
]
24+
\`\`\`
25+
26+
This will help you quickly understand what happened in each session.`,
27+
}),
28+
),
29+
})),
30+
},
31+
}))
32+
33+
vi.mock('@langchain/openai', () => ({
34+
ChatOpenAI: vi.fn(() => ({
35+
invoke: vi.fn(() =>
36+
Promise.resolve({
37+
content: `\`\`\`json
38+
[
39+
{
40+
"op": "add",
41+
"path": "/tables/design_sessions/columns/summary",
42+
"value": {
43+
"name": "summary",
44+
"type": "text",
45+
"not_null": false
46+
}
47+
}
48+
]
49+
\`\`\``,
50+
}),
51+
),
52+
})),
53+
}))
54+
55+
describe('LangGraph Chat Pipeline', () => {
56+
beforeEach(() => {
57+
vi.clearAllMocks()
58+
})
59+
60+
it('should return a response with JSON Patch fence for schema changes', async () => {
61+
const userMsg = 'Add summary column'
62+
const schemaText = 'CREATE TABLE design_sessions (id SERIAL PRIMARY KEY);'
63+
const chatHistory = 'No previous conversation.'
64+
65+
const result = await runChat(userMsg, schemaText, chatHistory)
66+
67+
expect(result).toContain('```json')
68+
expect(result).toContain('summary')
69+
expect(result).toContain('design_sessions')
70+
})
71+
72+
it('should parse JSON Patch correctly', async () => {
73+
const userMsg = 'Add summary column'
74+
const schemaText = 'CREATE TABLE design_sessions (id SERIAL PRIMARY KEY);'
75+
const chatHistory = 'No previous conversation.'
76+
77+
const result = await runChat(userMsg, schemaText, chatHistory)
78+
79+
// Extract JSON from the response
80+
const jsonMatch = result?.match(/```json\s+([\s\S]+?)\s*```/i)
81+
expect(jsonMatch).toBeTruthy()
82+
83+
if (jsonMatch) {
84+
const patch = JSON.parse(jsonMatch[1]) as unknown[]
85+
expect(Array.isArray(patch)).toBe(true)
86+
expect(patch[0]).toHaveProperty('op', 'add')
87+
expect(patch[0]).toHaveProperty('path')
88+
expect(patch[0]).toHaveProperty('value')
89+
}
90+
})
91+
92+
it('should handle chat history in the prompt', async () => {
93+
const userMsg = 'Add summary column'
94+
const schemaText = 'CREATE TABLE design_sessions (id SERIAL PRIMARY KEY);'
95+
const chatHistory = 'User: Hello\nAssistant: Hi there!'
96+
97+
const result = await runChat(userMsg, schemaText, chatHistory)
98+
99+
expect(result).toBeTruthy()
100+
expect(typeof result).toBe('string')
101+
})
102+
103+
it('should handle empty chat history', async () => {
104+
const userMsg = 'Add summary column'
105+
const schemaText = 'CREATE TABLE design_sessions (id SERIAL PRIMARY KEY);'
106+
const chatHistory = 'No previous conversation.'
107+
108+
const result = await runChat(userMsg, schemaText, chatHistory)
109+
110+
expect(result).toBeTruthy()
111+
expect(typeof result).toBe('string')
112+
})
113+
})

‎frontend/apps/app/lib/chat/langGraph.ts‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
import { mastra } from '@/lib/mastra'
2+
import { Annotation, END, START, StateGraph } from '@langchain/langgraph'
3+
import { ChatOpenAI } from '@langchain/openai'
4+
5+
////////////////////////////////////////////////////////////////
6+
// 1. Type definitions for the StateGraph
7+
////////////////////////////////////////////////////////////////
8+
interface ChatState {
9+
userMsg: string
10+
schemaText: string
11+
chatHistory: string
12+
sysPrompt?: string
13+
14+
draft?: string
15+
patch?: unknown[]
16+
valid?: boolean
17+
retryCount?: number
18+
}
19+
20+
// define the annotations for the StateGraph
21+
const ChatStateAnnotation = Annotation.Root({
22+
userMsg: Annotation<string>,
23+
schemaText: Annotation<string>,
24+
chatHistory: Annotation<string>,
25+
sysPrompt: Annotation<string>,
26+
draft: Annotation<string>,
27+
patch: Annotation<unknown[]>,
28+
valid: Annotation<boolean>,
29+
retryCount: Annotation<number>,
30+
})
31+
32+
////////////////////////////////////////////////////////////////
33+
// 2. Implementation of the StateGraph nodes
34+
////////////////////////////////////////////////////////////////
35+
36+
const buildPrompt = async (s: ChatState): Promise<Partial<ChatState>> => {
37+
const sysPrompt = `
38+
You are Build Agent, an energetic and innovative system designer who builds and edits ERDs with lightning speed.
39+
Your role is to execute user instructions immediately and offer smart suggestions for schema improvements.
40+
You speak in a lively, action-oriented tone, showing momentum and confidence.
41+
42+
Your personality is bold, constructive, and enthusiastic — like a master architect in a hardhat, ready to build.
43+
You say things like "Done!", "You can now...", and "Shall we move to the next step?".
44+
45+
Your communication should feel fast, fresh, and forward-moving, like a green plant constantly growing.
46+
47+
Do:
48+
- Confirm execution quickly: "Added!", "Created!", "Linked!"
49+
- Propose the next steps: "Would you like to add an index?", "Let's relate this to the User table too!"
50+
- Emphasize benefits: "This makes tracking updates easier."
51+
52+
Don't:
53+
- Hesitate ("Maybe", "We'll have to check...")
54+
- Use long, uncertain explanations
55+
- Get stuck in abstract talk — focus on action and outcomes
56+
57+
When in doubt, prioritize momentum, simplicity, and clear results.
58+
59+
<SCHEMA>
60+
${s.schemaText}
61+
</SCHEMA>
62+
63+
Previous conversation:
64+
${s.chatHistory}
65+
66+
#### REQUIRED OUTPUT FORMAT
67+
1. **Always** wrap your RFC 6902 JSON Patch in a **\`\`\`json … \`\`\`** code fence.
68+
2. Any text *other than* the JSON Patch (explanations, suggestions, etc.) may appear **before or after** the fence.
69+
**Do not** add filler phrases such as "Here is the patch" or "See below."
70+
Instead, include only meaningful comments—design rationale, next steps, trade-offs, and so on.
71+
3. If the user's question **does not** involve a schema change, **omit** the JSON Patch fence entirely.
72+
`
73+
return { sysPrompt }
74+
}
75+
76+
const draft = async (s: ChatState): Promise<Partial<ChatState>> => {
77+
const agent = mastra.getAgent('databaseSchemaBuildAgent')
78+
if (!agent) {
79+
throw new Error('databaseSchemaBuildAgent not found in Mastra instance')
80+
}
81+
if (!s.sysPrompt) {
82+
throw new Error('System prompt not built')
83+
}
84+
const res = await agent.generate([
85+
{ role: 'system', content: s.sysPrompt },
86+
{ role: 'user', content: s.userMsg },
87+
])
88+
return { draft: res.text }
89+
}
90+
91+
const check = async (s: ChatState): Promise<Partial<ChatState>> => {
92+
const m = s.draft?.match(/```json\s+([\s\S]+?)\s*```/i)
93+
if (!m) return { valid: false }
94+
try {
95+
return { valid: true, patch: JSON.parse(m[1]) }
96+
} catch {
97+
return { valid: false }
98+
}
99+
}
100+
101+
const remind = async (s: ChatState): Promise<Partial<ChatState>> => {
102+
const llm = new ChatOpenAI({ model: 'gpt-4o-mini' })
103+
const res = await llm.invoke([
104+
{
105+
role: 'system',
106+
content:
107+
'Return ONLY the ```json code fence with the RFC 6902 patch. No intro text.',
108+
},
109+
{ role: 'user', content: s.userMsg },
110+
])
111+
return { draft: res.content as string, retryCount: (s.retryCount ?? 0) + 1 }
112+
}
113+
114+
////////////////////////////////////////////////////////////////
115+
// 3. build StateGraph
116+
////////////////////////////////////////////////////////////////
117+
export const runChat = async (
118+
userMsg: string,
119+
schemaText: string,
120+
chatHistory: string,
121+
) => {
122+
try {
123+
const graph = new StateGraph(ChatStateAnnotation)
124+
125+
graph
126+
.addNode('buildPrompt', buildPrompt)
127+
.addNode('drafted', draft)
128+
.addNode('check', check)
129+
.addNode('remind', remind)
130+
.addEdge(START, 'buildPrompt')
131+
.addEdge('buildPrompt', 'drafted')
132+
.addEdge('remind', 'check')
133+
134+
// conditional edges
135+
.addConditionalEdges('check', (s: ChatState) => {
136+
if (s.valid) return END
137+
if ((s.retryCount ?? 0) >= 3) return END // give up
138+
return 'remind'
139+
})
140+
141+
// execution
142+
const compiled = graph.compile()
143+
const result = await compiled.invoke(
144+
{
145+
userMsg,
146+
schemaText,
147+
chatHistory,
148+
retryCount: 0,
149+
},
150+
{
151+
recursionLimit: 4, // for avoid deep recursion
152+
},
153+
)
154+
155+
return result.draft ?? 'No response generated'
156+
} catch (error) {
157+
console.error(
158+
'StateGraph execution failed, falling back to manual execution:',
159+
error,
160+
)
161+
// some fallback logic
162+
}
163+
}

‎frontend/apps/app/lib/mastra/agents/databaseSchemaBuildAgent.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,33 @@ Don't:
3030
3131
When in doubt, prioritize momentum, simplicity, and clear results.
3232
33+
---
34+
35+
#### REQUIRED OUTPUT FORMAT
36+
1. **Always** wrap your RFC 6902 JSON Patch in a **\`\`\`json … \`\`\`** code fence.
37+
2. Any text *other than* the JSON Patch (explanations, suggestions, etc.) may appear **before or after** the fence.
38+
**Do not** add filler phrases such as "Here is the patch" or "See below."
39+
Instead, include only meaningful comments—design rationale, next steps, trade-offs, and so on.
40+
3. Example:
41+
42+
\`\`\`markdown
43+
### Why we need \`summary\`
44+
45+
Adding a nullable \`summary\` helps …
46+
\`summary\` will be displayed on …
47+
48+
\`\`\`json
49+
[
50+
{ "op": "add",
51+
"path": "/tables/design_sessions/columns/summary",
52+
"value": { "name": "summary", "type": "text", "not_null": false } }
53+
]
54+
\`\`\`
55+
56+
Next, we might add an index …
57+
\`\`\`
58+
59+
4. If the user’s question **does not** involve a schema change, **omit** the JSON Patch fence entirely.
3360
`,
3461
model: openai('o4-mini-2025-04-16'),
3562
})

‎frontend/apps/app/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
"@codemirror/view": "6.36.8",
1515
"@langchain/community": "0.3.43",
1616
"@langchain/core": "0.3.55",
17+
"@langchain/langgraph": "0.2.73",
1718
"@langchain/openai": "0.5.10",
1819
"@lezer/highlight": "1.2.1",
1920
"@liam-hq/db": "workspace:*",

0 commit comments

Comments
 (0)