Skip to content

Commit 5fdcf87

Browse files
vkarpov15vkarpov15
authored
Merge pull request #156 from mongoosejs/codex/add-ai-input-for-document-creation
Add AI draft streaming for create-document flow
2 parents 3b30898 + 14e36f6 commit 5fdcf87

File tree

7 files changed

+253
-2
lines changed

7 files changed

+253
-2
lines changed

‎backend/actions/Model/index.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,5 +12,6 @@ exports.getDocumentsStream = require('./getDocumentsStream');
1212
exports.getCollectionInfo = require('./getCollectionInfo');
1313
exports.getIndexes = require('./getIndexes');
1414
exports.listModels = require('./listModels');
15+
exports.streamChatMessage = require('./streamChatMessage');
1516
exports.updateDocument = require('./updateDocument');
1617
exports.updateDocuments = require('./updateDocuments');

‎backend/actions/Model/streamChatMessage.js‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
'use strict';
2+
3+
const Archetype = require('archetype');
4+
const authorize = require('../../authorize');
5+
const streamLLM = require('../../integrations/streamLLM');
6+
const getModelDescriptions = require('../../helpers/getModelDescriptions');
7+
8+
const StreamChatMessageParams = new Archetype({
9+
model: {
10+
$type: 'string',
11+
$required: true
12+
},
13+
content: {
14+
$type: 'string',
15+
$required: true
16+
},
17+
documentData: {
18+
$type: 'string'
19+
},
20+
roles: {
21+
$type: ['string']
22+
}
23+
}).compile('StreamChatMessageParams');
24+
25+
module.exports = ({ db, options }) => async function* streamChatMessage(params) {
26+
const { model, content, documentData, roles } = new StreamChatMessageParams(params);
27+
28+
await authorize('Model.streamChatMessage', roles);
29+
30+
const Model = db.models[model];
31+
if (Model == null) {
32+
throw new Error(`Model ${model} not found`);
33+
}
34+
35+
const modelDescriptions = getModelDescriptions({ models: { [Model.modelName]: Model } });
36+
const context = [
37+
modelDescriptions,
38+
'Current draft document:\n' + (documentData || '')
39+
].join('\n\n');
40+
const system = systemPrompt + '\n\n' + context + (options?.context ? '\n\n' + options.context : '');
41+
42+
const llmMessages = [{ role: 'user', content: [{ type: 'text', text: content }] }];
43+
const textStream = streamLLM(llmMessages, system, options);
44+
45+
for await (const textPart of textStream) {
46+
yield { textPart };
47+
}
48+
49+
return {};
50+
};
51+
52+
const systemPrompt = `
53+
You are a helpful assistant that drafts MongoDB documents for the user.
54+
55+
Use the model description and the current draft document to refine the user's intent.
56+
57+
Return only the updated document body as a JavaScript object literal. Do not use Markdown or code fences.
58+
`.trim();

‎backend/authorize.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ const actionsToRequiredRoles = {
2323
'Model.getDocumentsStream': ['owner', 'admin', 'member', 'readonly'],
2424
'Model.getIndexes': ['owner', 'admin', 'member', 'readonly'],
2525
'Model.listModels': ['owner', 'admin', 'member', 'readonly'],
26+
'Model.streamChatMessage': ['owner', 'admin', 'member', 'readonly'],
2627
'Model.updateDocuments': ['owner', 'admin', 'member']
2728
};
2829

‎backend/helpers/getModelDescriptions.js‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,16 @@ const formatSchemaTypeInstance = schemaType => {
1717
return schemaType.instance;
1818
};
1919

20+
const formatEnum = schemaType => {
21+
if (!schemaType.options?.enum) {
22+
return '';
23+
}
24+
return ` (enum: ${JSON.stringify(schemaType.options.enum)})`;
25+
};
26+
2027
const formatSchemaPath = (path, schemaType) => `- ${path}: ${formatSchemaTypeInstance(schemaType)}` +
2128
formatRef(schemaType) +
29+
formatEnum(schemaType) +
2230
(schemaType.schema ? formatNestedSchema(schemaType) : '');
2331

2432
const listModelPaths = Model => [

‎frontend/src/api.js‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,55 @@ if (window.MONGOOSE_STUDIO_CONFIG.isLambda) {
147147
updateDocument: function updateDocument(params) {
148148
return client.post('', { action: 'Model.updateDocument', ...params }).then(res => res.data);
149149
},
150+
streamChatMessage: async function* streamChatMessage(params) {
151+
const accessToken = window.localStorage.getItem('_mongooseStudioAccessToken') || null;
152+
const url = window.MONGOOSE_STUDIO_CONFIG.baseURL + '?' + new URLSearchParams({ action: 'Model.streamChatMessage', ...params }).toString();
153+
154+
const response = await fetch(url, {
155+
method: 'GET',
156+
headers: {
157+
Authorization: `${accessToken}`,
158+
Accept: 'text/event-stream'
159+
}
160+
});
161+
162+
if (!response.ok) {
163+
throw new Error(`HTTP error! Status: ${response.status}`);
164+
}
165+
166+
const reader = response.body.getReader();
167+
const decoder = new TextDecoder('utf-8');
168+
let buffer = '';
169+
170+
while (true) {
171+
const { done, value } = await reader.read();
172+
if (done) break;
173+
buffer += decoder.decode(value, { stream: true });
174+
175+
let eventEnd;
176+
while ((eventEnd = buffer.indexOf('\n\n')) !== -1) {
177+
const eventStr = buffer.slice(0, eventEnd);
178+
buffer = buffer.slice(eventEnd + 2);
179+
180+
// Parse SSE event
181+
const lines = eventStr.split('\n');
182+
let data = '';
183+
for (const line of lines) {
184+
if (line.startsWith('data:')) {
185+
data += line.slice(5).trim();
186+
}
187+
}
188+
if (data) {
189+
try {
190+
const res = JSON.parse(data);
191+
yield res;
192+
} catch (err) {
193+
yield data;
194+
}
195+
}
196+
}
197+
}
198+
},
150199
updateDocuments: function updateDocuments(params) {
151200
return client.post('', { action: 'Model.updateDocuments', ...params }).then(res => res.data);
152201
}
@@ -356,6 +405,55 @@ if (window.MONGOOSE_STUDIO_CONFIG.isLambda) {
356405
updateDocument: function updateDocument(params) {
357406
return client.post('/Model/updateDocument', params).then(res => res.data);
358407
},
408+
streamChatMessage: async function* streamChatMessage(params) {
409+
const accessToken = window.localStorage.getItem('_mongooseStudioAccessToken') || null;
410+
const url = window.MONGOOSE_STUDIO_CONFIG.baseURL + '/Model/streamChatMessage?' + new URLSearchParams(params).toString();
411+
412+
const response = await fetch(url, {
413+
method: 'GET',
414+
headers: {
415+
Authorization: `${accessToken}`,
416+
Accept: 'text/event-stream'
417+
}
418+
});
419+
420+
if (!response.ok) {
421+
throw new Error(`HTTP error! Status: ${response.status}`);
422+
}
423+
424+
const reader = response.body.getReader();
425+
const decoder = new TextDecoder('utf-8');
426+
let buffer = '';
427+
428+
while (true) {
429+
const { done, value } = await reader.read();
430+
if (done) break;
431+
buffer += decoder.decode(value, { stream: true });
432+
433+
let eventEnd;
434+
while ((eventEnd = buffer.indexOf('\n\n')) !== -1) {
435+
const eventStr = buffer.slice(0, eventEnd);
436+
buffer = buffer.slice(eventEnd + 2);
437+
438+
// Parse SSE event
439+
const lines = eventStr.split('\n');
440+
let data = '';
441+
for (const line of lines) {
442+
if (line.startsWith('data:')) {
443+
data += line.slice(5).trim();
444+
}
445+
}
446+
if (data) {
447+
try {
448+
yield JSON.parse(data);
449+
} catch (err) {
450+
// If not JSON, yield as string
451+
yield data;
452+
}
453+
}
454+
}
455+
}
456+
},
359457
updateDocuments: function updateDocument(params) {
360458
return client.post('/Model/updateDocuments', params).then(res => res.data);
361459
}

‎frontend/src/create-document/create-document.html‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,41 @@
11
<div>
2+
<div class="mb-4">
3+
<label class="block text-sm font-bold text-gray-900">AI Mode</label>
4+
<div class="mt-2 flex flex-col gap-2 sm:flex-row sm:items-center">
5+
<input
6+
v-model="aiPrompt"
7+
type="text"
8+
placeholder="Describe the document you'd like to create..."
9+
@keydown.enter.prevent="requestAiSuggestion()"
10+
class="w-full rounded-md border border-gray-300 px-3 py-2 text-sm shadow-sm focus:border-ultramarine-500 focus:outline-none focus:ring-1 focus:ring-ultramarine-500"
11+
/>
12+
<button
13+
@click="requestAiSuggestion()"
14+
:disabled="aiStreaming || !aiPrompt.trim()"
15+
class="inline-flex items-center justify-center rounded-md bg-ultramarine-600 px-3 py-2 text-sm font-semibold text-white shadow-sm hover:bg-ultramarine-500 disabled:cursor-not-allowed disabled:opacity-50"
16+
>
17+
{{ aiStreaming ? 'Generating...' : 'Generate' }}
18+
</button>
19+
</div>
20+
<p class="mt-2 text-xs text-gray-500">Use AI to draft the document. You can accept or reject the suggestion once it finishes.</p>
21+
<div v-if="aiSuggestionReady" class="mt-3 flex flex-wrap gap-2">
22+
<button
23+
@click="acceptAiSuggestion()"
24+
class="rounded-md bg-emerald-600 px-2.5 py-1.5 text-sm font-semibold text-white shadow-sm hover:bg-emerald-500"
25+
>
26+
Accept suggestion
27+
</button>
28+
<button
29+
@click="rejectAiSuggestion()"
30+
class="rounded-md bg-gray-100 px-2.5 py-1.5 text-sm font-semibold text-gray-700 shadow-sm hover:bg-gray-200"
31+
>
32+
Reject suggestion
33+
</button>
34+
</div>
35+
</div>
236
<div class="mb-2">
3-
<textarea class="border border-gray-200 p-2 h-[300px] w-full" ref="codeEditor"></textarea>
37+
<label class="block text-sm font-bold text-gray-900">Document to Create</label>
38+
<textarea class="border border-gray-200 p-2 h-[300px] w-full mt-2" ref="codeEditor"></textarea>
439
</div>
540
<button @click="createDocument()" class="rounded-md bg-ultramarine-600 px-2.5 py-1.5 text-sm font-semibold text-white shadow-sm hover:bg-ultramarine-500 focus-visible:outline focus-visible:outline-2 focus-visible:outline-offset-2 focus-visible:outline-teal-600">Submit</button>
641
<div v-if="errors.length > 0" class="rounded-md bg-red-50 p-4 mt-1">

‎frontend/src/create-document/create-document.js‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,60 @@ module.exports = app => app.component('create-document', {
2323
return {
2424
documentData: '',
2525
editor: null,
26-
errors: []
26+
errors: [],
27+
aiPrompt: '',
28+
aiSuggestion: '',
29+
aiOriginalDocument: '',
30+
aiStreaming: false,
31+
aiSuggestionReady: false
2732
};
2833
},
2934
methods: {
35+
async requestAiSuggestion() {
36+
if (this.aiStreaming) {
37+
return;
38+
}
39+
const prompt = this.aiPrompt.trim();
40+
if (!prompt) {
41+
return;
42+
}
43+
44+
this.aiOriginalDocument = this.editor.getValue();
45+
this.aiSuggestion = '';
46+
this.aiSuggestionReady = false;
47+
this.aiStreaming = true;
48+
49+
try {
50+
for await (const event of api.Model.streamChatMessage({
51+
model: this.currentModel,
52+
content: prompt,
53+
documentData: this.aiOriginalDocument
54+
})) {
55+
if (event?.textPart) {
56+
this.aiSuggestion += event.textPart;
57+
this.editor.setValue(this.aiSuggestion);
58+
}
59+
}
60+
this.aiSuggestionReady = true;
61+
} catch (err) {
62+
this.editor.setValue(this.aiOriginalDocument);
63+
this.$toast.error('Failed to generate a document suggestion.');
64+
throw err;
65+
} finally {
66+
this.aiStreaming = false;
67+
}
68+
},
69+
acceptAiSuggestion() {
70+
this.aiSuggestionReady = false;
71+
this.aiSuggestion = '';
72+
this.aiOriginalDocument = '';
73+
},
74+
rejectAiSuggestion() {
75+
this.editor.setValue(this.aiOriginalDocument);
76+
this.aiSuggestionReady = false;
77+
this.aiSuggestion = '';
78+
this.aiOriginalDocument = '';
79+
},
3080
async createDocument() {
3181
const data = EJSON.serialize(eval(`(${this.editor.getValue()})`));
3282
try {

0 commit comments

Comments
 (0)