From c58566c82e2705334bf8e04a0aca7213b929be05 Mon Sep 17 00:00:00 2001
From: Simon Knott <info@simonknott.de>
Date: Thu, 13 Feb 2025 12:45:36 +0100
Subject: [PATCH] more

---
 .../trace-viewer/src/ui/aiConversation.tsx    | 26 +++++++++++--------
 packages/trace-viewer/src/ui/llm.tsx          | 21 ++++++++-------
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/packages/trace-viewer/src/ui/aiConversation.tsx b/packages/trace-viewer/src/ui/aiConversation.tsx
index f9bc550fe0..adbc64f70f 100644
--- a/packages/trace-viewer/src/ui/aiConversation.tsx
+++ b/packages/trace-viewer/src/ui/aiConversation.tsx
@@ -7,11 +7,11 @@ import type { Conversation, LLMMessage } from './llm';
 export function AIConversation({ history, conversation }: { history: LLMMessage[], conversation: Conversation }) {
   const [input, setInput] = useState('');
 
-  const onSubmit = useCallback(async (event: React.FormEvent<HTMLFormElement>) => {
-    event.preventDefault();
-    setInput('');
-    const content = new FormData(event.target as any).get('content') as string;
-    await conversation.send(content);
+  const onSubmit = useCallback(() => {
+    setInput(content => {
+      conversation.send(content);
+      return '';
+    });
   }, [conversation]);
 
   return (
@@ -34,29 +34,33 @@ export function AIConversation({ history, conversation }: { history: LLMMessage[
         ))}
       </div>
 
-      <form onSubmit={onSubmit} className="input-form">
-        <input
-          type="text"
+      <div className="input-form">
+        <textarea
           name='content'
           value={input}
           onChange={(e) => setInput(e.target.value)}
+          onKeyDown={e => {
+            if (e.key === 'Enter' && !e.shiftKey) {
+              e.preventDefault();
+              onSubmit();
+            }
+          }}
           placeholder="Ask a question..."
           className="message-input"
         />
         {conversation.isSending() ? (
           <button type="button" className="send-button" onClick={(evt) => {
             evt.preventDefault()
-            console.log("aborting")
             conversation.abortSending();
           }}>
             Cancel
           </button>
         ) : (
-          <button type="submit" className="send-button" disabled={!input.trim()}>
+          <button className="send-button" disabled={!input.trim()} onClick={onSubmit}>
             Send
           </button>  
         )}
-      </form>
+      </div>
     </div>
   );
 }
diff --git a/packages/trace-viewer/src/ui/llm.tsx b/packages/trace-viewer/src/ui/llm.tsx
index 776b4b7dfb..29d607e2d2 100644
--- a/packages/trace-viewer/src/ui/llm.tsx
+++ b/packages/trace-viewer/src/ui/llm.tsx
@@ -169,7 +169,7 @@ class LLMChat {
   getConversation(id: string, systemPrompt: string) {
     if (!this.conversations.has(id)) {
       const conversation = new Conversation(this, systemPrompt);
-      this.conversations.set(id, conversation);
+      this.conversations.set(id, conversation); // TODO: cleanup
     }
     return this.conversations.get(id)!;
   }
@@ -178,36 +178,37 @@ class LLMChat {
 export class Conversation {
   history: LLMMessage[];
   onChange = new EventEmitter<void>();
-  private _abortController: AbortController | undefined;
+  private _abortControllers = new Set<AbortController>();
 
   constructor(private chat: LLMChat, systemPrompt: string) {
     this.history = [{ role: 'developer', content: systemPrompt }];
   }
 
   async send(content: string, displayContent?: string) {
-    if (this.isSending())
-      throw new Error('Already sending');
-
     const response: LLMMessage = { role: 'assistant', content: '' };
     this.history.push({ role: 'user', content, displayContent }, response);
+    const abortController = new AbortController();
+    this._abortControllers.add(abortController);
     this.onChange.fire();
-    this._abortController = new AbortController();
     try {
-      for await (const chunk of this.chat.api.chatCompletion(this.history, this._abortController.signal)) {
+      for await (const chunk of this.chat.api.chatCompletion(this.history, abortController.signal)) {
         response.content += chunk;
         this.onChange.fire();
       }
     } finally {
-      this._abortController = undefined;
+      this._abortControllers.delete(abortController);
+      this.onChange.fire();
     }
   }
 
   isSending(): boolean {
-    return this._abortController !== undefined;
+    return this._abortControllers.size > 0;
   }
 
   abortSending() {
-    this._abortController!.abort();
+    for (const controller of this._abortControllers)
+      controller.abort();
+    this._abortControllers.clear();
     this.onChange.fire();
   }