Skip to content

Commit

Permalink
Added diff view for SQL query in non-confirmation mode (#74)
Browse files Browse the repository at this point in the history
  • Loading branch information
kektobiologist authored Nov 11, 2024
1 parent 89e2418 commit c9b5e4c
Show file tree
Hide file tree
Showing 5 changed files with 58 additions and 22 deletions.
24 changes: 23 additions & 1 deletion apps/src/base/appController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@ export function Action(metadata: ActionMetadata) {
Reflect.defineMetadata('actionMetadata', metadata, target, propertyKey);
};
}

// export type ActionReturn = string | void | Promise<string | void>
// export type CheckMethods<T> = { [K in keyof T]: T[K] extends Function ?
// ((...args: any) => ActionReturn) : T[K] }
export abstract class AppController<T> {
protected app: App<T>;

Expand Down Expand Up @@ -109,8 +113,26 @@ export abstract class AppController<T> {
}

async runAction(fn: string, args: any) {
let renderInfo = {
text: null,
code: null,
oldCode: null
}
// get render info if it exists
const metadata = Reflect.getMetadata('actionMetadata', this, fn);
if (metadata) {
let renderBody = metadata['renderBody']
if (typeof renderBody === 'function') {
renderInfo = await renderBody(args, await this.app.getState())
}
}
// @ts-ignore: Check if controller has function and execute!
return await this[fn](args);
let result = await this[fn](args);
return {
type: 'BLANK',
...result,
renderInfo
}
}
}

Expand Down
5 changes: 3 additions & 2 deletions apps/src/metabase/appController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,9 @@ export class MetabaseController extends AppController<MetabaseAppState> {
labelRunning: "Updating SQL query",
labelDone: "Updated query",
description: "Updates the SQL query in the Metabase SQL editor and executes it.",
renderBody: ({ sql }: { sql: string }) => {
return {text: null, code: sql}
renderBody: ({ sql }: { sql: string }, appState: MetabaseAppStateSQLEditor) => {
const sqlQuery = appState?.sqlQuery
return {text: null, code: sql, oldCode: sqlQuery}
}
})
async updateSQLQuery({ sql, executeImmediately = true }: { sql: string, executeImmediately?: boolean }) {
Expand Down
20 changes: 7 additions & 13 deletions web/src/components/common/ActionStack.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import { getApp } from "../../helpers/app";
import 'reflect-metadata';
import { parseArguments } from '../../planner/plannerActions';
import { CodeBlock } from './CodeBlock';
import { ActionRenderInfo } from '../../state/chat/reducer';

function removeThinkingTags(input: string): string {
return input ? input.replace(/<thinking>[\s\S]*?<\/thinking>/g, '') : input;
Expand All @@ -25,7 +26,10 @@ function extractMessageContent(input: string): string {
return match ? match[1] : "";
}

export type ActionStatusView = Pick<Action, 'finished' | 'function' | 'status'>
export type ActionStatusView = Pick<Action, 'finished' | 'function' | 'status'> & {
renderInfo: ActionRenderInfo
}

export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusView>, index:number, content: string, latency: number}> = ({
actions,
status,
Expand All @@ -46,16 +50,6 @@ export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusV
return action;
}

const renderActionArgs = (action: string, args: string) => {
if (controller) {
const metadata = Reflect.getMetadata('actionMetadata', controller, action);
if (metadata) {
return metadata['renderBody'](parseArguments(args, action));
}
}
return {text: null, code: null};
}

let title: string = "";
if (status == 'FINISHED') {
let titles = actions.map(action => getActionLabels(action.function.name, 'labelDone'))
Expand Down Expand Up @@ -116,7 +110,7 @@ export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusV

</HStack>
{isExpanded && actions.map((action, index) => {
const { text, code } = renderActionArgs(action.function.name, action.function.arguments)
const { text, code, oldCode } = action.renderInfo || {text: null, code: null, oldCode: undefined}
return (
<VStack className={'action'} padding={'2px'} key={index} alignItems={"start"}>
<HStack>
Expand All @@ -132,7 +126,7 @@ export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusV
</HStack>

{ code && <Box width={"100%"} p={2} bg={"#1e1e1e"} borderRadius={5}>
<CodeBlock code={code || ""} tool={currentTool} oldCode={undefined}/>
<CodeBlock code={code || ""} tool={currentTool} oldCode={oldCode || undefined}/>
</Box>
}

Expand Down
15 changes: 10 additions & 5 deletions web/src/components/common/Chat.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,10 @@ import { ActionStack, ActionStatusView, OngoingActionStack } from './ActionStack
import { ChatContent } from './ChatContent';
import { getApp } from '../../helpers/app';

function addStatusInfoToActionPlanMessages(messages: Array<ChatMessage>) {
// adds tool information like execution status and rendering info
// this stuff is in the 'tool' messages, but we're ony rendering 'assistant' messages
// so this copy needs to be done while rendering.
function addToolInfoToActionPlanMessages(messages: Array<ChatMessage>) {
const toolMessages = messages.filter(message => message.role == 'tool') as Array<ActionChatMessage>
const toolMessageMap = new Map(toolMessages.map((message: ActionChatMessage) => [message.action.id, message]))
return messages.map(message => {
Expand All @@ -20,7 +23,8 @@ function addStatusInfoToActionPlanMessages(messages: Array<ChatMessage>) {
if (toolMessage) {
return {
...toolCall,
status: toolMessage.action.status
status: toolMessage.action.status,
renderInfo: toolMessage.content.renderInfo
}
} else {
return toolCall
Expand All @@ -39,7 +43,7 @@ function addStatusInfoToActionPlanMessages(messages: Array<ChatMessage>) {
})
}

const Chat: React.FC<ReturnType<typeof addStatusInfoToActionPlanMessages>[number]> = ({
const Chat: React.FC<ReturnType<typeof addToolInfoToActionPlanMessages>[number]> = ({
index,
role,
content,
Expand All @@ -63,7 +67,8 @@ const Chat: React.FC<ReturnType<typeof addStatusInfoToActionPlanMessages>[number
actions.push({
finished: true,
function: toolCall.function,
status: toolCall.status
status: toolCall.status,
renderInfo: toolCall.renderInfo
})
})
const latency = ('latency' in debug)? Math.round(debug.latency as number /100)/10 : 0
Expand Down Expand Up @@ -175,7 +180,7 @@ export const ChatSection = () => {
// just create a map of all role='tool' messages by their id, and for each
// tool call in each assistant message, add the status from the corresponding
// tool message
const messagesWithStatus = addStatusInfoToActionPlanMessages(messages)
const messagesWithStatus = addToolInfoToActionPlanMessages(messages)
const Chats = isEmpty(messagesWithStatus) ?
<HelperMessage /> :
messagesWithStatus.map((message, key) => (<Chat key={key} {...message} />))
Expand Down
16 changes: 15 additions & 1 deletion web/src/state/chat/reducer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,16 @@ export interface ActionPlanChatMessage extends BaseChatMessage {
content: ActionPlanMessageContent
}

export type ActionChatMessageContent = Subset<ChatMessageContent, DefaultMessageContent | BlankMessageContent>

export type ActionRenderInfo = {
text: string | null,
code: string | null,
oldCode?: string | null
}

export type ActionChatMessageContent = Subset<ChatMessageContent, DefaultMessageContent | BlankMessageContent> & {
renderInfo: ActionRenderInfo
}
export interface ActionChatMessage extends BaseChatMessage {
role: 'tool'
action: Action
Expand Down Expand Up @@ -220,6 +229,11 @@ export const chatSlice = createSlice({
index: actionMessageID,
content: {
type: 'BLANK',
renderInfo: {
text: null,
code: null,
oldCode: null
}
},
createdAt: timestamp,
updatedAt: timestamp,
Expand Down

0 comments on commit c9b5e4c

Please sign in to comment.