Skip to content

Commit da1e21e

Browse files
matt-aitkenmatt-aitken
committed
The AI supports editing queries now
1 parent d86d6cb commit da1e21e

File tree

5 files changed

+234
-45
lines changed

5 files changed

+234
-45
lines changed

‎apps/webapp/app/components/code/AIQueryInput.tsx‎

Lines changed: 78 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { PencilSquareIcon, SparklesIcon } from "@heroicons/react/20/solid";
12
import { AnimatePresence, motion } from "framer-motion";
23
import { Suspense, lazy, useCallback, useEffect, useRef, useState } from "react";
34
import { AISparkleIcon } from "~/assets/icons/AISparkleIcon";
@@ -17,22 +18,31 @@ import { Spinner } from "~/components/primitives/Spinner";
1718
import { useEnvironment } from "~/hooks/useEnvironment";
1819
import { useOrganization } from "~/hooks/useOrganizations";
1920
import { useProject } from "~/hooks/useProject";
21+
import { cn } from "~/utils/cn";
2022

2123
type StreamEventType =
2224
| { type: "thinking"; content: string }
2325
| { type: "tool_call"; tool: string; args: unknown }
24-
| { type: "tool_result"; tool: string; result: unknown }
2526
| { type: "result"; success: true; query: string }
2627
| { type: "result"; success: false; error: string };
2728

29+
export type AIQueryMode = "new" | "edit";
30+
2831
interface AIQueryInputProps {
2932
onQueryGenerated: (query: string) => void;
3033
/** Set this to a prompt to auto-populate and immediately submit */
3134
autoSubmitPrompt?: string;
35+
/** The current query in the editor (used for edit mode) */
36+
currentQuery?: string;
3237
}
3338

34-
export function AIQueryInput({ onQueryGenerated, autoSubmitPrompt }: AIQueryInputProps) {
39+
export function AIQueryInput({
40+
onQueryGenerated,
41+
autoSubmitPrompt,
42+
currentQuery,
43+
}: AIQueryInputProps) {
3544
const [prompt, setPrompt] = useState("");
45+
const [mode, setMode] = useState<AIQueryMode>("new");
3646
const [isLoading, setIsLoading] = useState(false);
3747
const [thinking, setThinking] = useState("");
3848
const [error, setError] = useState<string | null>(null);
@@ -48,9 +58,20 @@ export function AIQueryInput({ onQueryGenerated, autoSubmitPrompt }: AIQueryInpu
4858

4959
const resourcePath = `/resources/orgs/${organization.slug}/projects/${project.slug}/env/${environment.slug}/query/ai-generate`;
5060

61+
// Can only use edit mode if there's a current query
62+
const canEdit = Boolean(currentQuery?.trim());
63+
64+
// If mode is edit but there's no current query, switch to new
65+
useEffect(() => {
66+
if (mode === "edit" && !canEdit) {
67+
setMode("new");
68+
}
69+
}, [mode, canEdit]);
70+
5171
const submitQuery = useCallback(
52-
async (queryPrompt: string) => {
72+
async (queryPrompt: string, submitMode: AIQueryMode = mode) => {
5373
if (!queryPrompt.trim() || isLoading) return;
74+
if (submitMode === "edit" && !currentQuery?.trim()) return;
5475

5576
setIsLoading(true);
5677
setThinking("");
@@ -67,6 +88,10 @@ export function AIQueryInput({ onQueryGenerated, autoSubmitPrompt }: AIQueryInpu
6788
try {
6889
const formData = new FormData();
6990
formData.append("prompt", queryPrompt);
91+
formData.append("mode", submitMode);
92+
if (submitMode === "edit" && currentQuery) {
93+
formData.append("currentQuery", currentQuery);
94+
}
7095

7196
const response = await fetch(resourcePath, {
7297
method: "POST",
@@ -135,7 +160,7 @@ export function AIQueryInput({ onQueryGenerated, autoSubmitPrompt }: AIQueryInpu
135160
setIsLoading(false);
136161
}
137162
},
138-
[isLoading, resourcePath]
163+
[isLoading, resourcePath, mode, currentQuery]
139164
);
140165

141166
const processStreamEvent = useCallback(
@@ -147,9 +172,6 @@ export function AIQueryInput({ onQueryGenerated, autoSubmitPrompt }: AIQueryInpu
147172
case "tool_call":
148173
setThinking((prev) => prev + `\nValidating query...\n`);
149174
break;
150-
case "tool_result":
151-
// Optionally show validation result
152-
break;
153175
case "result":
154176
if (event.success) {
155177
onQueryGenerated(event.query);
@@ -217,7 +239,11 @@ export function AIQueryInput({ onQueryGenerated, autoSubmitPrompt }: AIQueryInpu
217239
<textarea
218240
ref={textareaRef}
219241
name="prompt"
220-
placeholder="e.g. show me failed runs from the last 7 days"
242+
placeholder={
243+
mode === "edit"
244+
? "e.g. add a filter for failed runs, change the limit to 50"
245+
: "e.g. show me failed runs from the last 7 days"
246+
}
221247
value={prompt}
222248
onChange={(e) => setPrompt(e.target.value)}
223249
disabled={isLoading}
@@ -231,16 +257,50 @@ export function AIQueryInput({ onQueryGenerated, autoSubmitPrompt }: AIQueryInpu
231257
}}
232258
/>
233259
<div className="flex justify-end gap-2 px-2 pb-2">
234-
<Button
235-
type="submit"
236-
variant="tertiary/small"
237-
disabled={isLoading || !prompt.trim()}
238-
LeadingIcon={isLoading ? Spinner : AISparkleIcon}
239-
className="pl-1.5"
240-
iconSpacing="gap-1.5"
241-
>
242-
{isLoading ? "Generating" : "Generate"}
243-
</Button>
260+
{isLoading ? (
261+
<Button
262+
type="button"
263+
variant="tertiary/small"
264+
disabled={true}
265+
LeadingIcon={Spinner}
266+
className="pl-1.5"
267+
iconSpacing="gap-1.5"
268+
>
269+
{mode === "edit" ? "Editing..." : "Generating..."}
270+
</Button>
271+
) : (
272+
<>
273+
<Button
274+
type="button"
275+
variant="tertiary/small"
276+
disabled={!prompt.trim()}
277+
LeadingIcon={SparklesIcon}
278+
className="pl-1.5"
279+
iconSpacing="gap-1.5"
280+
onClick={() => {
281+
setMode("new");
282+
submitQuery(prompt, "new");
283+
}}
284+
>
285+
New query
286+
</Button>
287+
<Button
288+
type="button"
289+
variant="tertiary/small"
290+
disabled={!prompt.trim() || !canEdit}
291+
LeadingIcon={PencilSquareIcon}
292+
className={cn("pl-1.5", !canEdit && "opacity-50")}
293+
iconSpacing="gap-1.5"
294+
tooltip={!canEdit ? "Write a query first to enable editing" : undefined}
295+
onClick={() => {
296+
setMode("edit");
297+
submitQuery(prompt, "edit");
298+
}}
299+
>
300+
Edit query
301+
</Button>
302+
</>
303+
)}
244304
</div>
245305
</form>
246306
</div>

‎apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.query/route.tsx‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -463,6 +463,7 @@ export default function Page() {
463463
setScope(exampleScope);
464464
}}
465465
onQueryGenerated={setQuery}
466+
currentQuery={query}
466467
/>
467468
</ResizablePanel>
468469
</>
@@ -585,10 +586,12 @@ function QueryHelpSidebar({
585586
onClose,
586587
onTryExample,
587588
onQueryGenerated,
589+
currentQuery,
588590
}: {
589591
onClose: () => void;
590592
onTryExample: (query: string, scope: QueryScope) => void;
591593
onQueryGenerated: (query: string) => void;
594+
currentQuery: string;
592595
}) {
593596
return (
594597
<div className="grid h-full max-h-full grid-rows-[auto_1fr] overflow-hidden bg-background-bright">
@@ -625,7 +628,7 @@ function QueryHelpSidebar({
625628
value="ai"
626629
className="min-h-0 flex-1 overflow-y-auto p-3 scrollbar-thin scrollbar-track-transparent scrollbar-thumb-charcoal-600"
627630
>
628-
<AITabContent onQueryGenerated={onQueryGenerated} />
631+
<AITabContent onQueryGenerated={onQueryGenerated} currentQuery={currentQuery} />
629632
</ClientTabsContent>
630633
<ClientTabsContent
631634
value="guide"
@@ -650,7 +653,13 @@ function QueryHelpSidebar({
650653
);
651654
}
652655

653-
function AITabContent({ onQueryGenerated }: { onQueryGenerated: (query: string) => void }) {
656+
function AITabContent({
657+
onQueryGenerated,
658+
currentQuery,
659+
}: {
660+
onQueryGenerated: (query: string) => void;
661+
currentQuery: string;
662+
}) {
654663
const [autoSubmitPrompt, setAutoSubmitPrompt] = useState<string | undefined>();
655664

656665
const examplePrompts = [
@@ -666,10 +675,13 @@ function AITabContent({ onQueryGenerated }: { onQueryGenerated: (query: string)
666675
<div>
667676
<Header3 className="mb-2 text-text-bright">Generate query with AI</Header3>
668677
<Paragraph variant="small" className="mb-3 text-text-dimmed">
669-
Describe the data you want to query in natural language. The AI will generate a valid TSQL
670-
query for you.
678+
Describe the data you want to query in natural language, or edit the existing query.
671679
</Paragraph>
672-
<AIQueryInput onQueryGenerated={onQueryGenerated} autoSubmitPrompt={autoSubmitPrompt} />
680+
<AIQueryInput
681+
onQueryGenerated={onQueryGenerated}
682+
autoSubmitPrompt={autoSubmitPrompt}
683+
currentQuery={currentQuery}
684+
/>
673685
</div>
674686

675687
<div className="border-t border-grid-dimmed pt-4">

‎apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.query.ai-generate.tsx‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import { querySchemas } from "~/v3/querySchemas";
1111

1212
const RequestSchema = z.object({
1313
prompt: z.string().min(1, "Prompt is required"),
14+
mode: z.enum(["new", "edit"]).default("new"),
15+
currentQuery: z.string().optional(),
1416
});
1517

1618
export async function action({ request, params }: ActionFunctionArgs) {
@@ -79,7 +81,7 @@ export async function action({ request, params }: ActionFunctionArgs) {
7981
);
8082
}
8183

82-
const { prompt } = submission.data;
84+
const { prompt, mode, currentQuery } = submission.data;
8385

8486
const service = new AIQueryService(
8587
querySchemas,
@@ -105,7 +107,7 @@ export async function action({ request, params }: ActionFunctionArgs) {
105107
};
106108

107109
try {
108-
const result = service.streamQuery(prompt);
110+
const result = service.streamQuery(prompt, { mode, currentQuery });
109111

110112
// Process the stream
111113
for await (const part of result.fullStream) {
@@ -120,13 +122,6 @@ export async function action({ request, params }: ActionFunctionArgs) {
120122
args: part.args,
121123
});
122124
break;
123-
case "tool-result":
124-
sendEvent({
125-
type: "tool_result",
126-
tool: part.toolName,
127-
result: part.result,
128-
});
129-
break;
130125
case "error":
131126
sendEvent({
132127
type: "result",
@@ -145,7 +140,10 @@ export async function action({ request, params }: ActionFunctionArgs) {
145140
success: true,
146141
query,
147142
});
148-
} else if (finalText.toLowerCase().includes("cannot") || finalText.toLowerCase().includes("unable")) {
143+
} else if (
144+
finalText.toLowerCase().includes("cannot") ||
145+
finalText.toLowerCase().includes("unable")
146+
) {
149147
sendEvent({
150148
type: "result",
151149
success: false,
@@ -200,4 +198,3 @@ function extractQueryFromText(text: string): string | null {
200198

201199
return null;
202200
}
203-

‎apps/webapp/app/utils/dataExport.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,3 +76,4 @@ export function downloadFile(content: string, filename: string, mimeType: string
7676
URL.revokeObjectURL(url);
7777
}
7878

79+

0 commit comments

Comments
 (0)