import gradio as gr import os import json from dotenv import load_dotenv from omoa import OllamaAgent, OllamaMixtureOfAgents, DEFAULT_PROMPTS, create_default_agents from MemoryAssistant.memory import AgentCoreMemory, AgentEventMemory from MemoryAssistant.prompts import wrap_user_message_in_xml_tags_json_mode from llama_cpp_agent.chat_history.messages import Roles # Load environment variables load_dotenv() # Ollama-specific environment variables os.environ['OLLAMA_NUM_PARALLEL'] = os.getenv('OLLAMA_NUM_PARALLEL', '4') os.environ['OLLAMA_MAX_LOADED_MODELS'] = os.getenv('OLLAMA_MAX_LOADED_MODELS', '4') MODEL_AGGREGATE = os.getenv("MODEL_AGGREGATE") MODEL_REFERENCE_1 = os.getenv("MODEL_REFERENCE_1") MODEL_REFERENCE_2 = os.getenv("MODEL_REFERENCE_2") MODEL_REFERENCE_3 = os.getenv("MODEL_REFERENCE_3") # Modify these lines to include all available models ALL_MODELS = [MODEL_AGGREGATE, MODEL_REFERENCE_1, MODEL_REFERENCE_2, MODEL_REFERENCE_3] ALL_MODELS = [model for model in ALL_MODELS if model] # Remove any None values # Global variables to store the MoA configuration moa_config = { "aggregate_agent": None, "reference_agents": [], "mixture": None } # Initialize memory components agent_core_memory = AgentCoreMemory(["persona", "user", "scratchpad"], core_memory_file="MemoryAssistant/core_memory.json") agent_event_memory = AgentEventMemory() def create_mixture(): moa_config["mixture"] = OllamaMixtureOfAgents( moa_config["reference_agents"], moa_config["aggregate_agent"] ) # Set the memory components after initialization moa_config["mixture"].agent_core_memory = agent_core_memory moa_config["mixture"].agent_event_memory = agent_event_memory def initialize_moa(): global moa_config default_agents = create_default_agents() moa_config["aggregate_agent"] = default_agents["SynthesisAgent"] moa_config["reference_agents"] = [ default_agents["AnalyticalAgent"], default_agents["HistoricalContextAgent"], default_agents["ScienceTruthAgent"] ] moa_config["mixture"] = OllamaMixtureOfAgents( moa_config["reference_agents"], moa_config["aggregate_agent"], temperature=0.6, max_tokens=2048, rounds=1 ) moa_config["mixture"].web_search_enabled = True moa_config["mixture"].agent_core_memory = agent_core_memory moa_config["mixture"].agent_event_memory = agent_event_memory print("Mixture of Agents initialized successfully!") # Call initialize_moa() at the start of the application initialize_moa() def create_agent(model, name, system_prompt, **params): supported_params = ['model', 'name', 'system_prompt'] # Add any other supported parameters here filtered_params = {k: v for k, v in params.items() if k in supported_params} return OllamaAgent(model, name, system_prompt, **filtered_params) def clear_core_memory(): if isinstance(moa_config["mixture"], OllamaMixtureOfAgents): return moa_config["mixture"].clear_core_memory() else: return "Error: MoA not initialized properly." def clear_archival_memory(): if isinstance(moa_config["mixture"], OllamaMixtureOfAgents): return moa_config["mixture"].clear_archival_memory() else: return "Error: MoA not initialized properly." def edit_archival_memory(old_content, new_content): if isinstance(moa_config["mixture"], OllamaMixtureOfAgents): return moa_config["mixture"].edit_archival_memory(old_content, new_content) else: return "Error: MoA not initialized properly." async def process_message(message, history): # Add user message to event memory agent_event_memory.add_event(Roles.user, wrap_user_message_in_xml_tags_json_mode(message)) response, web_search_performed = await moa_config["mixture"].get_response(message) # Ensure the response is a list of tuples if isinstance(response, str): formatted_response = [(None, response)] elif isinstance(response, list): formatted_response = [(None, str(item)) for item in response] else: formatted_response = [(None, str(response))] info = f"Generated response using {len(moa_config['reference_agents'])} reference agents and 1 aggregate agent." if web_search_performed: info += " Web search was performed during response generation." return formatted_response, info async def chat(message, history): response, processing_info = await process_message(message, history) # Ensure the response is a list of lists formatted_response = [[message, item[1]] if isinstance(item, tuple) else [message, str(item)] for item in response] # Append the new messages to the history updated_history = history + formatted_response # Ensure the final output is a list of lists final_output = [[msg, resp] for msg, resp in updated_history] return final_output, processing_info def update_memory(self, message, role): # Update event memory self.agent_event_memory.add_event(role, message) # Update RAG self.rag.add_document(message) def get_model_params(model_name): # Define custom parameters for each model params = { "llama2": ["temperature", "top_p", "top_k", "repeat_penalty", "num_ctx"], "mistral": ["temperature", "top_p", "top_k", "repeat_penalty", "num_ctx"], "codellama": ["temperature", "top_p", "top_k", "repeat_penalty", "num_ctx"], } return params.get(model_name, ["temperature", "top_p", "top_k", "repeat_penalty", "num_ctx"]) # Default parameters if model not found def update_model_params(model_name): params = get_model_params(model_name) components = [gr.Markdown(f"### {model_name} Parameters")] for param in params: if param == "temperature": components.append(gr.Slider(minimum=0, maximum=2, value=0.7, step=0.1, label="Temperature")) elif param == "top_p": components.append(gr.Slider(minimum=0, maximum=1, value=0.9, step=0.05, label="Top P")) elif param == "top_k": components.append(gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top K")) elif param == "repeat_penalty": components.append(gr.Slider(minimum=0.1, maximum=2, value=1.1, step=0.05, label="Repeat Penalty")) elif param == "num_ctx": components.append(gr.Slider(minimum=128, maximum=4096, value=2048, step=128, label="Context Length")) return components def update_agent_config(old_agent_name, model, new_name, prompt, **params): new_agent = create_agent(model, new_name, prompt, **params) if old_agent_name == "SynthesisAgent": moa_config["aggregate_agent"] = new_agent else: moa_config["reference_agents"] = [agent for agent in moa_config["reference_agents"] if agent.name != old_agent_name] moa_config["reference_agents"].append(new_agent) create_mixture() return f"Updated agent configuration: {old_agent_name} -> {new_name}" def edit_core_memory(section, key, value): agent_core_memory.update_core_memory(section, {key: value}) return f"Core memory updated: {section}.{key} = {value}" def search_archival_memory(query): results = moa_config["mixture"].search_archival_memory(query) return f"Archival memory search results for '{query}':\n{results}" def add_to_archival_memory(content): if isinstance(moa_config["mixture"], OllamaMixtureOfAgents): moa_config["mixture"].add_to_archival_memory(content) return f"Added to archival memory: {content}" return f"Failed to add to archival memory: {content}. MoA not initialized properly." def toggle_web_search(enabled): if isinstance(moa_config["mixture"], OllamaMixtureOfAgents): return moa_config["mixture"].toggle_web_search(enabled) return "Error: MoA not initialized properly." def create_gradio_interface(): global moa_config theme = gr.themes.Base( primary_hue="green", secondary_hue="orange", # Changed from "brown" to "orange" neutral_hue="gray", font=("Helvetica", "sans-serif"), ).set( body_background_fill="linear-gradient(to right, #1a2f0f, #3d2b1f)", body_background_fill_dark="linear-gradient(to right, #0f1a09, #261a13)", button_primary_background_fill="#3d2b1f", button_primary_background_fill_hover="#4e3827", block_title_text_color="#d3c6aa", block_label_text_color="#b8a888", input_background_fill="#f0e6d2", input_background_fill_dark="#2a1f14", input_border_color="#7d6d58", input_border_color_dark="#5c4c3d", checkbox_background_color="#3d2b1f", checkbox_background_color_selected="#5e4534", slider_color="#7d6d58", slider_color_dark="#5c4c3d", ) css = """ .gradio-container { background-image: url('file/assets/mycelium_bg.png'); background-size: cover; background-attachment: fixed; } .gr-box { border-radius: 15px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); background-color: rgba(255, 255, 255, 0.1); backdrop-filter: blur(5px); } .gr-button { border-radius: 25px; } .gr-input { border-radius: 10px; } .gr-form { border-radius: 15px; background-color: rgba(255, 255, 255, 0.05); } """ with gr.Blocks(theme=theme, css=css) as demo: gr.Markdown( """ # Mycomind Daemon: Advanced Mixture-of-Memory-RAG-Agents (MoMRA) Cognitive Assistant Harness the power of interconnected AI models inspired by mycelial networks. """ ) with gr.Tab("Configure MoA"): agent_tabs = ["Agent1", "Agent2", "Agent3", "Synthesis Agent"] all_agents = moa_config["reference_agents"] + [moa_config["aggregate_agent"]] for i, agent in enumerate(all_agents): with gr.Tab(agent_tabs[i]): with gr.Row(): with gr.Column(scale=1): model = gr.Dropdown( choices=ALL_MODELS, value=agent.model, label="Model" ) name = gr.Textbox( value=agent.name, label="Agent Name", interactive=True ) with gr.Column(scale=2): prompt = gr.Textbox( value=agent.system_prompt, label="System Prompt", lines=10, interactive=True ) with gr.Group() as params_group: gr.Markdown(f"### {agent.model} Parameters") temperature = gr.Slider(minimum=0, maximum=2, value=0.7, step=0.1, label="Temperature") top_p = gr.Slider(minimum=0, maximum=1, value=0.9, step=0.05, label="Top P") top_k = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top K") repeat_penalty = gr.Slider(minimum=0.1, maximum=2, value=1.1, step=0.05, label="Repeat Penalty") num_ctx = gr.Slider(minimum=128, maximum=4096, value=2048, step=128, label="Context Length") model.change( update_model_params, inputs=[model], outputs=[params_group] ) update_btn = gr.Button(f"Update {agent_tabs[i]}") update_status = gr.Textbox(label="Update Status", interactive=False) def update_agent_wrapper(agent_index): params = { "temperature": temperature.value, "top_p": top_p.value, "top_k": top_k.value, "repeat_penalty": repeat_penalty.value, "num_ctx": num_ctx.value } return update_agent_config(all_agents[agent_index].name, model.value, name.value, prompt.value, **params) update_btn.click( lambda: update_agent_wrapper(i), outputs=[update_status] ) with gr.Tab("Chat"): chatbot = gr.Chatbot(label="Chat History", height=400) with gr.Row(): msg = gr.Textbox(label="Your Message", placeholder="Type your message here...", lines=2, scale=4) send_btn = gr.Button("Send", variant="primary", scale=1) clear_btn = gr.Button("Clear Chat") processing_log = gr.Textbox(label="Processing Log", interactive=False) with gr.Tab("Memory Management"): with gr.Row(): with gr.Column(): archival_query = gr.Textbox(label="Archival Memory Search Query") search_archival_btn = gr.Button("Search Archival Memory") archival_results = gr.Textbox(label="Archival Memory Results", interactive=False) with gr.Column(): gr.Markdown("### Archival Memory Management") clear_archival_btn = gr.Button("Clear Archival Memory") clear_archival_status = gr.Textbox(label="Clear Archival Memory Status", interactive=False) gr.Markdown("### Edit Archival Memory") old_content = gr.Textbox(label="Old Content") new_content = gr.Textbox(label="New Content") edit_archival_btn = gr.Button("Edit Archival Memory") edit_archival_status = gr.Textbox(label="Edit Archival Memory Status", interactive=False) with gr.Column(): archival_content = gr.Textbox(label="Content to Add to Archival Memory") add_archival_btn = gr.Button("Add to Archival Memory") archival_status = gr.Textbox(label="Archival Memory Status", interactive=False) # with gr.Row(): # gr.Markdown("### Core Memory Viewer") # core_memory_viewer = gr.JSON(label="Current Core Memory", value=moa_config["mixture"].load_core_memory()) # refresh_core_memory_btn = gr.Button("Refresh Core Memory View") # with gr.Row(): # gr.Markdown("### Core Memory Editor") # core_memory_editor = gr.Textbox(label="Edit Core Memory", value=json.dumps(moa_config["mixture"].load_core_memory(), indent=2), lines=10, max_lines=20) # update_core_memory_btn = gr.Button("Update Core Memory") # core_memory_status = gr.Textbox(label="Core Memory Update Status", interactive=False) with gr.Tab("RAG Management"): with gr.Row(): with gr.Column(): upload_file = gr.File(label="Upload Document") upload_btn = gr.Button("Process Document") upload_status = gr.Textbox(label="Upload Status", interactive=False) with gr.Column(): gr.Markdown("### RAG Configuration") chunk_size = gr.Slider(minimum=128, maximum=1024, value=512, step=64, label="Chunk Size") chunk_overlap = gr.Slider(minimum=0, maximum=256, value=0, step=32, label="Chunk Overlap") k_value = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Number of Retrieved Documents (k)") with gr.Row(): gr.Markdown("### RAG Status") rag_status = gr.JSON(label="Current RAG Status") refresh_rag_status_btn = gr.Button("Refresh RAG Status") def update_rag_config(chunk_size, chunk_overlap, k_value): rag = moa_config["mixture"].rag # Update attributes if they exist if hasattr(rag, 'chunk_size'): rag.chunk_size = chunk_size if hasattr(rag, 'chunk_overlap'): rag.chunk_overlap = chunk_overlap if hasattr(rag, 'k'): rag.k = k_value # If there's a specific method to update configuration, use it if hasattr(rag, 'update_config'): rag.update_config(chunk_size=chunk_size, chunk_overlap=chunk_overlap, k=k_value) # If there's a method to reinitialize the index with new settings, call it if hasattr(rag, 'reinitialize_index'): rag.reinitialize_index() return "RAG configuration updated successfully" def get_rag_status(): rag = moa_config["mixture"].rag status = { "Index Size": rag.get_index_size() if hasattr(rag, 'get_index_size') else "Not available", "Current Configuration": rag.get_config() if hasattr(rag, 'get_config') else "Not available" } # Try to get document count if the method exists if hasattr(rag, 'get_document_count'): status["Document Count"] = rag.get_document_count() elif hasattr(rag, 'index') and hasattr(rag.index, '__len__'): status["Document Count"] = len(rag.index) else: status["Document Count"] = "Not available" return status update_rag_config_btn = gr.Button("Update RAG Configuration") update_rag_config_status = gr.Textbox(label="Update Status", interactive=False) update_rag_config_btn.click( update_rag_config, inputs=[chunk_size, chunk_overlap, k_value], outputs=[update_rag_config_status] ) refresh_rag_status_btn.click( get_rag_status, outputs=[rag_status] ) with gr.Tab("Settings"): with gr.Row(): with gr.Column(): gr.Markdown("### Web Search") web_search_toggle = gr.Checkbox(label="Enable Web Search", value=True) web_search_status = gr.Textbox(label="Web Search Status", interactive=False) with gr.Column(): gr.Markdown("### Processing Parameters") rounds_slider = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Processing Rounds") temperature_slider = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature") max_tokens_slider = gr.Slider(minimum=100, maximum=4096, value=1000, step=100, label="Max Tokens") with gr.Row(): gr.Markdown("### Additional Settings") stream_output_toggle = gr.Checkbox(label="Stream Output", value=True) debug_mode_toggle = gr.Checkbox(label="Debug Mode", value=False) #def refresh_core_memory(): # return moa_config["mixture"].load_core_memory() #def update_core_memory(new_core_memory_str): # try: # new_core_memory = json.loads(new_core_memory_str) # moa_config["mixture"].core_memory = new_core_memory # moa_config["mixture"].agent_core_memory.update_core_memory(new_core_memory) # moa_config["mixture"].agent_core_memory.save_core_memory(moa_config["mixture"].core_memory_file) # return json.dumps(new_core_memory, indent=2), "Core memory updated successfully" # except json.JSONDecodeError: # return json.dumps(moa_config["mixture"].load_core_memory(), indent=2), "Error: Invalid JSON format" # except Exception as e: # return json.dumps(moa_config["mixture"].load_core_memory(), indent=2), f"Error updating core memory: {str(e)}" def update_settings(rounds, temperature, max_tokens, stream_output, debug_mode): moa_config["mixture"].rounds = rounds moa_config["mixture"].temperature = temperature moa_config["mixture"].max_tokens = max_tokens moa_config["mixture"].stream_output = stream_output moa_config["mixture"].debug_mode = debug_mode return "Settings updated successfully" # update_core_memory_btn.click( # update_core_memory, # inputs=[core_memory_editor], # outputs=[core_memory_status] # ) # refresh_core_memory_btn.click( # refresh_core_memory, # outputs=[core_memory_viewer] # ) # update_core_memory_btn.click( # update_core_memory, # inputs=[core_memory_editor], # outputs=[core_memory_viewer, core_memory_status] # ) settings_update_btn = gr.Button("Update Settings") settings_update_status = gr.Textbox(label="Settings Update Status", interactive=False) settings_update_btn.click( update_settings, inputs=[rounds_slider, temperature_slider, max_tokens_slider, stream_output_toggle, debug_mode_toggle], outputs=[settings_update_status] ) web_search_toggle.change( toggle_web_search, inputs=[web_search_toggle], outputs=[web_search_status] ) msg.submit(chat, inputs=[msg, chatbot], outputs=[chatbot, processing_log]) send_btn.click(chat, inputs=[msg, chatbot], outputs=[chatbot, processing_log]) clear_btn.click(lambda: ([], ""), outputs=[chatbot, processing_log]) search_archival_btn.click( search_archival_memory, inputs=[archival_query], outputs=[archival_results] ) add_archival_btn.click( add_to_archival_memory, inputs=[archival_content], outputs=[archival_status] ) upload_btn.click( lambda file: moa_config["mixture"].upload_document(file.name) if file else "No file selected", inputs=[upload_file], outputs=[upload_status] ) clear_archival_btn.click( clear_archival_memory, outputs=[clear_archival_status] ) edit_archival_btn.click( edit_archival_memory, inputs=[old_content, new_content], outputs=[edit_archival_status] ) return demo if __name__ == "__main__": initialize_moa() demo = create_gradio_interface() demo.queue() demo.launch(share=True)