Initial working inline edits

apps/ui/src/app/MainChat.tsx

@@ -19,7 +19,8 @@ function LeftPanel() {
   const [isExpanded, setIsExpanded] = useState(false);
 
   const runTestFunc = () => {
-    trpc.testFunc.query();
+    console.log(chatStore.get('messages'));
+    // trpc.testFunc.query();
   };
 
   if (!isExpanded) {

apps/ui/src/app/Messages/ToolMessageRender.tsx

@@ -13,17 +13,12 @@ export function ToolMessageRender<T extends ToolType>({
   if (!type) return;
   const renderTemplate = TOOL_RENDER_TEMPLATES[type];
 
-  const description = useMemo(() => {
-    const result = renderTemplate.description(message);
-    return result;
-  }, [message.contents]);
-
   return (
     <Alert>
       <renderTemplate.Icon className="w-4 h-4" />
       <AlertTitle>{renderTemplate.title(message)}</AlertTitle>
       <AlertDescription className="break-words whitespace-pre-wrap">
-        {description}
+        {renderTemplate.description(message)}
       </AlertDescription>
       <div className={`flex gap-2 ${message.loading && 'hidden'}`}>
         <button

apps/ui/src/llms/messages/LLMOutputParser.ts

@@ -1,5 +1,4 @@
-import { sleep } from '@taffy/shared-helpers';
-import { CustomMessage } from './Messages';
+import { chatStore, getToolMessages } from '../../stores/chat-store';
 import {
   TOOL_END_MATCH_REGEX,
   TOOL_START_MATCH_REGEX,
@@ -10,9 +9,6 @@ import { TOOL_RENDER_TEMPLATES } from './tools';
 const logger = console;
 
 export class LLMOutputParser {
-  private inTool = false;
-  private messages: ToolMessage[] = [new ToolMessage()];
-
   async handleTextStream(
     stream: AsyncIterable<string>,
     onMsgUpdate?: () => void
@@ -46,33 +42,30 @@ export class LLMOutputParser {
     //The tool matches are more fuzzy than the actual instructions given to the models, to try to render as far as possible.
     const toolStartMatch = line.match(TOOL_START_MATCH_REGEX);
     const toolEndMatch = line.match(TOOL_END_MATCH_REGEX);
-    if (toolStartMatch && !this.inTool) {
-      this.messages.push(new ToolMessage());
-      this.inTool = true;
+    if (toolStartMatch) {
+      chatStore.set('messages', [
+        ...chatStore.get('messages'),
+        new ToolMessage(),
+      ]);
     }
 
-    const latestMsg = this.messages.at(-1);
+    const latestMsg = getToolMessages().at(-1);
     if (!latestMsg) {
       throw new Error('Expected at least one message!');
     }
-    if (this.inTool) {
-      latestMsg.contents += `${line}\n`;
-      latestMsg.loading = true;
-    }
+
+    latestMsg.contents += `${line}\n`;
 
     if (toolEndMatch) {
-      this.inTool = false;
-      latestMsg.loading = false;
       if (!latestMsg.type) return;
       const renderTemplate = TOOL_RENDER_TEMPLATES[latestMsg.type];
       if (!renderTemplate) return;
       if (!renderTemplate.onFocus) return;
       renderTemplate.onFocus(latestMsg as any);
       //TODO: Stop parsing until user finishes focus action
     }
-  }
 
-  getMessages(): CustomMessage[] {
-    return this.messages;
+    //Trigger rerender, cos I'm not sure how else to ensure proper updating
+    chatStore.set('messages', [...chatStore.get('messages')]);
   }
 }

apps/ui/src/llms/messages/ToolMessage.ts

@@ -5,14 +5,14 @@ import { makeObservable, observable, computed } from 'mobx';
 
 export const TOOL_START_MATCH_REGEX = /{TOOL (\w+)(?: (.*))?}/;
 export const TOOL_END_MATCH_REGEX = /{END_TOOL\s?(\w*)}/;
+export const THINKING_START_MATCH_REGEX =
+  /{THINKING_START}\n(.*?)\n(?:{THINKING_END})?/g;
 
 const logger = console;
 
 export class ToolMessage<
   ToolName extends ToolType = ToolType
 > extends BaseMessage {
-  loading: boolean = false;
-
   get role(): 'user' | 'assistant' | 'system' {
     return this.type ? TOOL_TEMPLATES[this.type].role : 'assistant';
   }
@@ -25,7 +25,7 @@ export class ToolMessage<
       props: computed,
       contents: observable,
       body: computed,
-      loading: observable,
+      loading: computed,
     });
     this.contents = contents ?? '';
   }
@@ -36,6 +36,11 @@ export class ToolMessage<
     return (curType as any) === type;
   }
 
+  get loading(): boolean {
+    const ended = Boolean(this.contents.match(TOOL_END_MATCH_REGEX));
+    return !ended;
+  }
+
   //TODO: Memoize
   get type(): ToolName | undefined {
     const toolStartMatch = this.contents.match(TOOL_START_MATCH_REGEX);
@@ -59,7 +64,8 @@ export class ToolMessage<
   get body(): string {
     let bodyMatch = this.contents
       .replace(TOOL_START_MATCH_REGEX, '')
-      .replace(TOOL_END_MATCH_REGEX, '');
+      .replace(TOOL_END_MATCH_REGEX, '')
+      .replace(THINKING_START_MATCH_REGEX, '');
     if (bodyMatch.startsWith('\n')) {
       bodyMatch = bodyMatch.substring(1);
     }
@@ -79,6 +85,15 @@ export class ToolMessage<
     } as any);
   }
 
+  get thoughts(): string | undefined {
+    const thoughtMatches = this.contents.matchAll(THINKING_START_MATCH_REGEX);
+    const thoughts: string[] = [];
+    for (const match of thoughtMatches) {
+      thoughts.push(match[1].trim());
+    }
+    return thoughts.join('\n');
+  }
+
   toRawMessages(): RawMessage[] {
     return [
       {

apps/ui/src/llms/messages/tools/toolTemplates.tsx

@@ -52,17 +52,32 @@ export const TOOL_TEMPLATES = {
   //   1. Read the .gitignore file using ASSISTANT_READ_FILE
   //   2. Update the .gitignore file using ASSISTANT_WRITE_FILE`,
   // },
+  //   ASSISTANT_WRITE_FILE: {
+  //     role: 'assistant',
+  //     desc: `Ask the user for permission to create/overwrite a file. You will need to provide the FULL FILE CONTENTS, because the action suggested to the user will be a full override of the existing file. DO NOT INCLUDE LINE NUMBERS IN THE OUTPUT.`,
+  //     propDesc: {
+  //       filePath:
+  //         "The path to which the file is written. If the file path doesn't exist, directories will be recursively created until we are able to create the file.",
+  //     },
+  //     sampleProps: {
+  //       filePath: 'src/utils/helloWorld.ts',
+  //     },
+  //     sampleBody: `export function helloWorld() {
+  // ${'console'}.log("Hello World!");
+  // }`,
+  //   },
+  //I'm not sure about the thinking start version
   ASSISTANT_WRITE_FILE: {
     role: 'assistant',
     desc: `Ask the user for permission to create/overwrite a file. You will need to provide the FULL FILE CONTENTS, because the action suggested to the user will be a full override of the existing file. DO NOT INCLUDE LINE NUMBERS IN THE OUTPUT. Before tackling a challenging part of the code, you can walk yourself through the coding process in a THINKING block
 
-{THINKING_START}
-In order to write this function, I will need to...
-{THINKING_END}
+  {THINKING_START}
+  In order to write this function, I will need to...
+  {THINKING_END}
 
-The thinking blocks will not be included in the outputted code.
-Within the thinking blocks, consider the user's existing code style and practices and follow those.
-`,
+  The thinking blocks will not be included in the outputted code.
+  Within the thinking blocks, consider the user's existing code style and practices and follow those.
+  `,
     propDesc: {
       filePath:
         "The path to which the file is written. If the file path doesn't exist, directories will be recursively created until we are able to create the file.",
@@ -72,7 +87,7 @@ Within the thinking blocks, consider the user's existing code style and practice
     },
     sampleBody: `export function helloWorld() {
 {THINKING_START}
-The hello world function should log hello world. 
+The hello world function should log hello world.
 {THINKING_END}
   ${'console'}.log("Hello World!");
 }`,
@@ -181,12 +196,15 @@ export const TOOL_RENDER_TEMPLATES: {
     title: () => 'Requesting permission to write the following files',
     description: (data) => {
       if (!data.props) return;
+      const thoughtsString = data.thoughts ? `💡${data.thoughts}` : '';
+      const infoString = data.loading
+        ? `${data.body.length} characters loaded so far`
+        : 'Loading Complete!';
+      const fullStr = infoString + '\n\n' + thoughtsString;
       return (
         <>
           <div>File Path - {data.props.filePath} </div>
-          {data.loading
-            ? `${data.body.length} characters loaded so far`
-            : 'Loading complete!'}
+          {fullStr}
         </>
       );
     },

apps/ui/src/stores/chat-store.ts

@@ -48,7 +48,6 @@ const setLlm = () => {
 };
 
 export async function continuePrompt(llm: LLM | null = chatStore.get('llm')) {
-  const curMsgs = chatStore.get('messages');
   if (!llm) {
     llm = setLlm();
   }
@@ -57,9 +56,7 @@ export async function continuePrompt(llm: LLM | null = chatStore.get('llm')) {
   const parser = new LLMOutputParser();
   const stream = llm.prompt(rawMessages);
 
-  await parser.handleTextStream(stream, () => {
-    chatStore.set('messages', [...curMsgs, ...parser.getMessages()]);
-  });
+  await parser.handleTextStream(stream);
 }
 
 function getRawMessages(): RawMessage[] {

apps/ui/src/stores/run-prompt.ts

@@ -38,7 +38,7 @@ async function updateChatInline(input: string) {
   });
   const startLine = +fileContextMsg.props.startLine;
   let preContents =
-    curContents?.split('\n').slice(0, startLine).join('\n') ?? '';
+    curContents?.split('\n').slice(0, startLine - 1).join('\n') ?? '';
   preContents += '\n{THINKING_START}\n';
 
   const preAssistantPrompt = new ToolMessage(