Agent mode: frontend

mito-ai/mito_ai/providers.py

@@ -243,9 +243,15 @@ async def request_completions(
                 # Validate that the model is supported. If not fall back to gpt-4o-mini
                 if model not in self.models:
                     model = "gpt-4o-mini"
-
-                completion_function_params = get_open_ai_completion_function_params(model, request.messages, False, response_format)
-                completion = self._openAI_sync_client.chat.completions.create(**completion_function_params)
+
+                if response_format:
+                    completion_function_params = get_open_ai_completion_function_params(
+                        model, request.messages, False, response_format
+                    )
+                    completion = self._openAI_sync_client.beta.chat.completions.parse(**completion_function_params)
+                else:
+                    completion_function_params = get_open_ai_completion_function_params(model, request.messages, False)
+                    completion = self._openAI_sync_client.chat.completions.create(**completion_function_params)
 
                 if prompt_type == "agent:planning":
                     pass # TODO: Add logging for agents 

mito-ai/mito_ai/utils/open_ai_utils.py

@@ -133,4 +133,9 @@ def get_open_ai_completion_function_params(
     if model == "gpt-4o-mini":
         completion_function_params["temperature"] = 0.0
 
+    # Remove "stream" key if response_format is set.
+    # beta.chat.completions.parse will error if we try to pass it as a param. 
+    if response_format:
+        completion_function_params.pop("stream", None)
+
     return completion_function_params

mito-ai/src/Extensions/AiChat/ChatHistoryManager.tsx

@@ -13,7 +13,7 @@ export type PromptType = 'chat' | 'smartDebug' | 'codeExplain' | 'agent:planning
 // we add a message to the chat ui that tells them to set an API key.
 export interface IDisplayOptimizedChatHistory {
     message: OpenAI.Chat.ChatCompletionMessageParam
-    type: 'openai message' | 'connection error',
+    type: 'openai message' | 'openai message:agent:planning' | 'connection error',
     mitoAIConnectionErrorType?: string | null,
     codeCellID: string | undefined
 }
@@ -108,7 +108,7 @@ export class ChatHistoryManager {
         }
     }
 
-    updateMessageAtIndex(index: number, newContent: string): IOutgoingMessage {
+    updateMessageAtIndex(index: number, newContent: string, isAgentMessage: boolean = false): IOutgoingMessage {
         const activeCellID = getActiveCellID(this.notebookTracker)
         const activeCellCode = getCellCodeByID(this.notebookTracker, activeCellID)
 
@@ -121,11 +121,14 @@ export class ChatHistoryManager {
 
         this.displayOptimizedChatHistory[index] = { 
             message: getDisplayedOptimizedUserMessage(newContent, activeCellCode),
-            type: 'openai message',
+            type: isAgentMessage ? 'openai message:agent:planning' : 'openai message',
             codeCellID: activeCellID
         }
 
-        this.displayOptimizedChatHistory = this.displayOptimizedChatHistory.slice(0, index + 1);
+        // Only slice the history if it's not an agent message
+        if (!isAgentMessage) {
+            this.displayOptimizedChatHistory = this.displayOptimizedChatHistory.slice(0, index + 1);
+        }
 
         return {
             promptType: 'chat',
@@ -216,12 +219,21 @@ export class ChatHistoryManager {
             content: messageContent
         }
 
+        let type: IDisplayOptimizedChatHistory['type'];
+        if (mitoAIConnectionError) {
+            type = 'connection error';
+        } else if (promptType === 'agent:planning') {
+            type = 'openai message:agent:planning';
+        } else {
+            type = 'openai message';
+        }
+
         const activeCellID = getActiveCellID(this.notebookTracker)
 
         this.displayOptimizedChatHistory.push(
             {
                 message: aiMessage, 
-                type: mitoAIConnectionError ? 'connection error' : 'openai message',
+                type: type,
                 mitoAIConnectionErrorType: mitoAIConnectionErrorType,
                 codeCellID: activeCellID
             }

mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx

@@ -18,9 +18,11 @@ import PlayButtonIcon from '../../../icons/PlayButtonIcon';
 import CopyIcon from '../../../icons/CopyIcon';
 import copyToClipboard from '../../../utils/copyToClipboard';
 import TextButton from '../../../components/TextButton';
+import { IDisplayOptimizedChatHistory } from '../ChatHistoryManager';
 
 interface IChatMessageProps {
     message: OpenAI.Chat.ChatCompletionMessageParam
+    messageType: IDisplayOptimizedChatHistory['type']
     codeCellID: string | undefined
     messageIndex: number
     mitoAIConnectionError: boolean
@@ -33,13 +35,14 @@ interface IChatMessageProps {
     previewAICode: () => void
     acceptAICode: () => void
     rejectAICode: () => void
-    onUpdateMessage: (messageIndex: number, newContent: string) => void
+    onUpdateMessage: (messageIndex: number, newContent: string, messageType: IDisplayOptimizedChatHistory['type']) => void
     variableManager?: IVariableManager
     codeReviewStatus: CodeReviewStatus
 }
 
 const ChatMessage: React.FC<IChatMessageProps> = ({
     message,
+    messageType,
     messageIndex,
     mitoAIConnectionError,
     mitoAIConnectionErrorType,
@@ -60,14 +63,16 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
         return null;
     }
 
+    const editable = messageType === 'openai message:agent:planning' || message.role === 'user'
+
     const messageContentParts = splitStringWithCodeBlocks(message);
 
     const handleEditClick = () => {
         setIsEditing(true);
     };
 
     const handleSave = (content: string) => {
-        onUpdateMessage(messageIndex, content);
+        onUpdateMessage(messageIndex, content, messageType);
         setIsEditing(false);
     };
 
@@ -94,7 +99,8 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
         <div className={classNames(
             "message",
             { "message-user": message.role === 'user' },
-            { 'message-assistant': message.role === 'assistant' },
+            { 'message-assistant-chat': message.role === 'assistant' && messageType !== 'openai message:agent:planning' },
+            { 'message-assistant-agent': messageType === 'openai message:agent:planning' },
         )}>
             {messageContentParts.map((messagePart, index) => {
                 if (messagePart.startsWith(PYTHON_CODE_BLOCK_START_WITHOUT_NEW_LINE)) {
@@ -178,7 +184,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                                     />
                                 )}
                             </p>
-                            {message.role === 'user' && (
+                            {editable && (
                                 <div style={{ display: 'flex', justifyContent: 'flex-end', marginTop: '4px' }}>
                                     <button
                                         className="message-edit-button"