Exploring multi-agent AI systems

import tkinter as tk
from tkinter import messagebox
from tkinter import ttk
from tkinter import scrolledtext
import json
import os
import subprocess
from functools import partial

SAVE_DIR = "notes"
AGENT_INSTRUCTIONS_DIR = "agent_instructions"

# Create necessary directories
for directory in (SAVE_DIR, AGENT_INSTRUCTIONS_DIR):
    os.makedirs(directory, exist_ok=True)

class MultilineDialog:
    def __init__(self, parent, title, prompt):
        self.result = None
        
        # Create dialog window
        self.dialog = tk.Toplevel(parent)
        self.dialog.title(title)
        self.dialog.transient(parent)
        self.dialog.grab_set()
        
        # Configure dialog
        self.dialog.geometry("400x300")
        
        # Add prompt label
        tk.Label(self.dialog, text=prompt, wraplength=380).pack(pady=5, padx=10)
        
        # Add text area
        self.text_area = scrolledtext.ScrolledText(
            self.dialog,
            wrap=tk.WORD,
            width=40,
            height=10,
            font=("Arial", 10)
        )
        self.text_area.pack(expand=True, fill='both', padx=10, pady=5)
        
        # Add buttons
        button_frame = tk.Frame(self.dialog)
        button_frame.pack(fill='x', padx=10, pady=5)
        
        self.ok_button = tk.Button(button_frame, text="OK", command=self._on_ok)
        self.ok_button.pack(side='right', padx=5)
        tk.Button(button_frame, text="Cancel", command=self._on_cancel).pack(side='right')
        
        # Center dialog on parent
        self.dialog.geometry("+%d+%d" % (
            parent.winfo_rootx() + 50,
            parent.winfo_rooty() + 50))
        
        # Bind Enter key to OK button
        self.dialog.bind('<Return>', lambda e: self._on_ok())
        
        # Set focus
        self.text_area.focus_set()
        
        # Wait for dialog
        self.dialog.wait_window()
    
    def _on_ok(self):
        self.result = self.text_area.get("1.0", "end-1c").strip()
        self.dialog.destroy()
    
    def _on_cancel(self):
        self.result = None  # Explicitly set result to None on cancel
        self.dialog.destroy()

class TextEditor:
    def __init__(self, parent):
        self.frame = tk.Frame(parent)
        self.frame.pack(side="left", expand=True, fill="both")
        
        # Text widget with efficient configuration
        self.text_widget = tk.Text(
            self.frame,
            wrap="word",
            font=("Arial", 17),
            undo=True,
            autoseparators=True,
            maxundo=-1
        )
        self.text_widget.pack(expand=True, fill="both")
        
        # Optimized event bindings
        self._setup_bindings()
    
    def _setup_bindings(self):
        """Set up optimized keyboard bindings."""
        for sequence, callback in [
            ('<Command-z>', self.undo),
            ('<Control-z>', self.undo),
            ('<Command-Shift-Z>', self.redo),
            ('<Control-y>', self.redo)
        ]:
            self.text_widget.bind_all(sequence, callback)
    
    def get_text(self):
        return self.text_widget.get("1.0", "end-1c")
    
    def set_text(self, text):
        self.text_widget.delete("1.0", tk.END)
        self.text_widget.insert("1.0", text)
    
    def get_selection(self):
        try:
            return self.text_widget.get(tk.SEL_FIRST, tk.SEL_LAST)
        except tk.TclError:
            return None
    
    def undo(self, event=None):
        try:
            self.text_widget.edit_undo()
        except tk.TclError:
            pass
        return "break"
    
    def redo(self, event=None):
        try:
            self.text_widget.edit_redo()
        except tk.TclError:
            pass
        return "break"

class WritingApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Writing App")
        
        # Configure root for better performance
        self.root.update_idletasks()
        
        # Main Layout
        self.frame = tk.Frame(root)
        self.frame.pack(fill="both", expand=True)
        
        self._init_components()
        self._load_initial_data()
    
    def _init_components(self):
        """Initialize all UI components efficiently."""
        self._init_notes_pane()
        self.editor = TextEditor(self.frame)
        self._init_quick_insert_pane()
        self._init_buttons()
        self._init_bindings()
        
        # Track current note
        self.current_note = None
    
    def _init_notes_pane(self):
        self.notes_frame = tk.Frame(self.frame, width=200, bg="#f0f0f0")
        self.notes_frame.pack(side="left", fill="y")
        
        tk.Label(self.notes_frame, text="Notes", font=("Arial", 12, "bold"), 
                bg="#000000").pack(pady=5)
        
        self.notes_listbox = tk.Listbox(self.notes_frame)
        self.notes_listbox.pack(fill="both", expand=True, padx=5, pady=5)
        
        # Optimize listbox selection handling
        self.notes_listbox.bind("<<ListboxSelect>>", 
                              lambda e: self.root.after(50, self.switch_note))
        
        tk.Button(self.notes_frame, text="New Note", 
                 command=lambda: self.root.after(1, self.create_new_note)).pack(pady=5)
    
    def _init_quick_insert_pane(self):
        self.quick_insert_frame = tk.Frame(self.frame, width=200, bg="#e0e0e0")
        self.quick_insert_frame.pack(side="right", fill="y")
        
        tk.Label(self.quick_insert_frame, text="Agents",
                font=("Arial", 12, "bold"), bg="#000000").pack(pady=5)
        
        self.quick_insert_listbox = tk.Listbox(self.quick_insert_frame)
        self.quick_insert_listbox.pack(fill="both", expand=True, padx=5, pady=5)
        
        # Optimize quick insert handling
        self.quick_insert_listbox.bind("<<ListboxSelect>>",
                                     lambda e: self.root.after(50, self.launch_chatbot_with_prompt))
    
    def _init_buttons(self):
        """Initialize buttons with optimized commands."""
        self.button_frame = tk.Frame(self.root)
        self.button_frame.pack(pady=5)
        
        buttons = [
            ("Chat with Selection", self.chat_with_selection),
            ("Save (Cmd+S / Ctrl+S)", self.save_current_note)
        ]
        
        for text, command in buttons:
            tk.Button(self.button_frame, text=text,
                     command=lambda cmd=command: self.root.after(1, cmd)).pack(side=tk.LEFT, padx=5)
    
    def _init_bindings(self):
        """Set up optimized keyboard bindings."""
        for sequence in ("<Command-s>", "<Control-s>"):
            self.root.bind(sequence, lambda e: self.root.after(1, self.save_current_note))
    
    def _load_initial_data(self):
        """Load initial data efficiently."""
        self.root.after(100, self.refresh_notes_list)
        self.root.after(100, self.load_agents)
    
    def create_new_note(self):
        dialog = MultilineDialog(self.root, "New Note", "Enter note title:")
        if dialog.result:
            file_path = os.path.join(SAVE_DIR, f"{dialog.result}.txt")
            if not os.path.exists(file_path):
                with open(file_path, "w", encoding="utf-8"):
                    pass
            self.refresh_notes_list()
    
    def switch_note(self):
        selected_indices = self.notes_listbox.curselection()
        if not selected_indices:
            return
        
        selected_note = self.notes_listbox.get(selected_indices[0])
        file_path = os.path.join(SAVE_DIR, f"{selected_note}.txt")
        
        if os.path.exists(file_path):
            self.load_note(file_path, selected_note)
    
    def load_note(self, file_path, note_name):
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                note_text = f.read()
            
            self.editor.set_text(note_text)
            self.current_note = note_name
        except Exception as e:
            messagebox.showerror("Error", f"Failed to load note: {e}")
    
    def save_current_note(self, event=None):
        if not self.current_note:
            return
        
        file_path = os.path.join(SAVE_DIR, f"{self.current_note}.txt")
        try:
            with open(file_path, "w", encoding="utf-8") as f:
                f.write(self.editor.get_text())
            messagebox.showinfo("Save", f"'{self.current_note}' saved successfully!")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to save note: {e}")
    
    def refresh_notes_list(self):
        self.notes_listbox.delete(0, tk.END)
        notes = sorted(f.replace(".txt", "") for f in os.listdir(SAVE_DIR) 
                      if f.endswith(".txt"))
        for note in notes:
            self.notes_listbox.insert(tk.END, note)
    
    def load_agents(self):
        self.quick_insert_listbox.delete(0, tk.END)
        inserts = sorted(f.replace(".txt", "") for f in os.listdir(AGENT_INSTRUCTIONS_DIR) if f.endswith(".txt"))
        for insert in inserts:
            self.quick_insert_listbox.insert(tk.END, insert)

    
    def chat_with_selection(self):
        selection = self.editor.get_selection()
        if not selection:
            messagebox.showwarning("Selection Error", 
                                 "Please select some text before starting a chat.")
            return
        
        dialog = MultilineDialog(self.root, "Chat Prompt", "Enter any additional context or questions:")
        if dialog.result is not None:  # Only proceed if not canceled
            additional_note = dialog.result
            full_prompt = f"{selection}\n\n{additional_note}".strip()
            self.root.after(1, lambda: subprocess.Popen(["python", "chat_app.py", full_prompt]))
        
    def launch_chatbot_with_prompt(self):
        selected_indices = self.quick_insert_listbox.curselection()
        if not selected_indices:
            return

        selected_text_key = self.quick_insert_listbox.get(selected_indices[0])
        file_path = os.path.join(AGENT_INSTRUCTIONS_DIR, f"{selected_text_key}.txt")

        try:
            with open(file_path, "r", encoding="utf-8") as f:
                quick_insert_text = f"##### {f.read().strip()} --------->"

            current_note_text = self.editor.get_text()
            dialog = MultilineDialog(self.root, "Additional Note", "Enter additional details (optional):")

            if dialog.result is not None:  # Proceed only if not canceled
                additional_note = dialog.result
                full_prompt = f"{quick_insert_text}\n\n{current_note_text}\n\n{additional_note}".strip()

                # Launch chat_app.py with the prompt and agent ID
                self.root.after(1, lambda: subprocess.Popen([
                    "python",
                    "chat_app.py",
                    full_prompt,
                    selected_text_key  # Pass the agent ID as an additional argument
                ]))

                # Clear the selection in the quick_insert_listbox
                self.quick_insert_listbox.selection_clear(0, tk.END)

        except Exception as e:
            messagebox.showerror("Error", f"Failed to launch chat: {e}")
            
if __name__ == "__main__":
    root = tk.Tk()
    app = WritingApp(root)
    root.mainloop()

import tkinter as tk
from tkinter import scrolledtext, messagebox
from litellm import completion
import os
import sys
from functools import partial
import queue
import threading
import weave; weave.init("writing-agent")

# Set API Keys
os.environ["OPENAI_API_KEY"] = ""
os.environ["GEMINI_API_KEY"] = ""
os.environ["ANTHROPIC_API_KEY"] = ""

def load_last_model():
    try:
        with open("last_model.txt", "r") as f:
            return f.read().strip()
    except:
        return "gemini/gemini-1.5-flash"

@weave.op
def run_inference(model_id, messages, temperature=0.7, agent_id=None):
    """
    Dedicated function to handle model inference
    
    Args:
        model_id (str): The ID of the model to use
        messages (list): List of conversation messages
        temperature (float): Temperature parameter for generation
        agent_id (str): The ID of the agent to use (optional)
        
    Returns:
        str: The model's response text
    
    Raises:
        Exception: If there's an error during inference
    """
    try:
        response = completion(
            model=model_id,
            messages=messages,
            temperature=temperature
        )
        return response["choices"][0]["message"]["content"]
    except Exception as e:
        raise Exception(f"Inference error: {str(e)}")
class ChatApp:
    def __init__(self, root, initial_prompt=None, agent_id=None):
        self.root = root
        self.root.title("AI Chat App")
        self.agent_id = agent_id
        
        # Message queue for thread-safe UI updates
        self.message_queue = queue.Queue()
        
        self._init_ui()
        self._setup_event_handling()
        
        # Chat history
        self.conversation = []
        
        # Processing flag
        self.is_processing = False
        
        # If an initial prompt is provided, schedule it
        if initial_prompt:
            self.root.after(100, self.send_initial_prompt, initial_prompt)

    
    def _init_ui(self):
        """Initialize UI components efficiently."""
        # Main frame
        self.main_frame = tk.Frame(self.root)
        self.main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # Model selection frame
        self._init_model_frame()
        
        # Chat display
        self._init_chat_display()
        
        # Input field
        self._init_input_field()
        
        # Buttons
        self._init_buttons()
    
    def _init_model_frame(self):
        """Initialize model selection frame."""
        model_frame = tk.Frame(self.main_frame)
        model_frame.pack(fill=tk.X, pady=(0, 10))
        
        tk.Label(model_frame, text="Model ID:").pack(side=tk.LEFT)
        
        self.model_entry = tk.Entry(model_frame)
        self.model_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(5, 0))
        self.model_entry.insert(0, load_last_model())
    
    def _init_chat_display(self):
        """Initialize chat display with optimized configuration."""
        self.chat_display = scrolledtext.ScrolledText(
            self.main_frame,
            wrap=tk.WORD,
            width=80,
            height=20,
            font=("Arial", 25),
            state=tk.DISABLED,
            setgrid=True  # Improved text rendering
        )
        self.chat_display.pack(fill=tk.BOTH, expand=True, pady=(0, 10))
        
        # Configure tags for formatting
        self.chat_display.tag_configure("user", foreground="yellow")
        self.chat_display.tag_configure("assistant", foreground="lightgreen")
        self.chat_display.tag_configure("system", foreground="red")
    
    def _init_input_field(self):
        """Initialize input field with optimized configuration."""
        self.input_field = scrolledtext.ScrolledText(
            self.main_frame,
            wrap=tk.WORD,
            width=80,
            height=5,
            font=("Arial", 10),
            undo=True,
            maxundo=50
        )
        self.input_field.pack(fill=tk.BOTH, expand=True, pady=(0, 10))
        
        # Bind Ctrl+Enter with delay to prevent double-firing
        self.input_field.bind("<Control-Return>", 
                            lambda e: self.root.after(10, self._handle_send))
    
    def _init_buttons(self):
        """Initialize buttons with optimized event handling."""
        button_frame = tk.Frame(self.main_frame)
        button_frame.pack(fill=tk.X)
        
        # Send button
        self.send_button = tk.Button(
            button_frame,
            text="Send (Ctrl+Enter)",
            command=lambda: self.root.after(10, self._handle_send)
        )
        self.send_button.pack(side=tk.LEFT, padx=(0, 5))
        
        # Clear button
        tk.Button(
            button_frame,
            text="Clear Chat",
            command=lambda: self.root.after(10, self.clear_chat)
        ).pack(side=tk.LEFT)
    
    def _setup_event_handling(self):
        """Set up event handling and message processing."""
        self.root.after(100, self._process_message_queue)
    
    def _process_message_queue(self):
        """Process messages in the queue."""
        try:
            while True:
                message = self.message_queue.get_nowait()
                self._update_chat_display(message)
        except queue.Empty:
            pass
        finally:
            # Schedule next check
            self.root.after(100, self._process_message_queue)
    
    def _update_chat_display(self, message):
        """Update chat display with message."""
        text, tag = message
        
        self.chat_display.config(state=tk.NORMAL)
        self.chat_display.insert(tk.END, text, tag)
        self.chat_display.see(tk.END)
        self.chat_display.config(state=tk.DISABLED)
    
    def _handle_send(self):
        """Handle send message action."""
        if self.is_processing:
            return
        
        user_message = self.input_field.get("1.0", tk.END).strip()
        if not user_message:
            return
        
        # Clear input field immediately
        self.input_field.delete("1.0", tk.END)
        
        # Disable input during processing
        self.is_processing = True
        self.send_button.config(state=tk.DISABLED)
        self.input_field.config(state=tk.DISABLED)
        
        # Start processing in separate thread
        threading.Thread(target=self._process_message, 
                       args=(user_message,), 
                       daemon=True).start()
    

    def _process_message(self, user_message):
        """Process message in background thread."""
        try:
            # Save selected model
            model_id = self.model_entry.get().strip()
            with open("last_model.txt", "w") as f:
                f.write(model_id)
            
            # Add user message to conversation and display
            self.conversation.append({"role": "user", "content": user_message})
            self.message_queue.put((f"You:\n{user_message}\n", "user"))
            self.message_queue.put(("#" * 50 + "\n", "system"))
            
            # Get AI response using dedicated inference function, now passing agent_id
            ai_reply = run_inference(model_id, self.conversation, agent_id=self.agent_id)
            
            # Add AI response to conversation and display
            self.conversation.append({"role": "assistant", "content": ai_reply})
            self.message_queue.put((f"AI:\n{ai_reply}\n", "assistant"))
            self.message_queue.put(("#" * 50 + "\n", "system"))
            
        except Exception as e:
            self.message_queue.put((f"Error: {str(e)}\n", "system"))
        
        finally:
            # Re-enable input
            self.root.after(0, self._enable_input)

    def _enable_input(self):
        """Re-enable input controls."""
        self.is_processing = False
        self.send_button.config(state=tk.NORMAL)
        self.input_field.config(state=tk.NORMAL)
        self.input_field.focus_set()
    
    def send_initial_prompt(self, prompt):
        """Send initial prompt when app starts."""
        if prompt:
            self._process_message(prompt)
    
    def clear_chat(self):
        """Clear chat history and display."""
        self.conversation = []
        self.chat_display.config(state=tk.NORMAL)
        self.chat_display.delete("1.0", tk.END)
        self.chat_display.config(state=tk.DISABLED)

if __name__ == "__main__":
    # Get initial prompt and agent_id from command line if provided
    initial_prompt = sys.argv[1] if len(sys.argv) > 1 else None
    agent_id = sys.argv[2] if len(sys.argv) > 2 else None
    
    root = tk.Tk()
    
    # Configure root window
    root.title("AI Chat App")
    root.geometry("800x600")
    
    # Make the window resizable
    root.rowconfigure(0, weight=1)
    root.columnconfigure(0, weight=1)
    
    app = ChatApp(root, initial_prompt, agent_id)
    
    # Start the application
    root.mainloop()