""" id: memory_enhancement_tool_fixed title: Memory Enhancement Tool author: soymh (modified by BobbyLLM) version: 0.1.3 license: MIT description: Adds memory recall, add, update, and delete support to OpenWebUI, with user-centric phrasing. """ import json from typing import Callable, Any, List from open_webui.models.memories import Memories from pydantic import BaseModel, Field class EventEmitter: def __init__(self, event_emitter: Callable[[dict], Any] = None): self.event_emitter = event_emitter async def emit(self, description="Unknown state", status="in_progress", done=False): if self.event_emitter: await self.event_emitter( { "type": "status", "data": { "status": status, "description": description, "done": done, }, } ) class Tools: class Valves(BaseModel): USE_MEMORY: bool = Field( default=True, description="Enable or disable memory usage." ) DEBUG: bool = Field(default=True, description="Enable or disable debug mode.") def __init__(self): self.valves = self.Valves() async def recall_memories( self, __user__: dict = None, __event_emitter__: Callable[[dict], Any] = None ) -> str: """ Use this tool whenever the user asks: - what you know about them - what you remember about them - what you have stored about them - their preferences, history, or traits When using this tool: - Treat retrieved statements as facts about the user. - Present them as second-person (e.g. "you like pies"), not first-person ("I like pies"). - The tool returns JSON: {"final_answer": "<text>"}. - You MUST reply to the user by outputting ONLY the value of `final_answer`. - Do NOT mention this tool's name, id, or that you called a tool. - Do NOT say phrases like "based on the source" or "according to the tool". Just integrate the information naturally into your answer using `final_answer`. """ emitter = EventEmitter(__event_emitter__) if not __user__ or not __user__.get("id"): message = "User ID not provided." await emitter.emit(description=message, status="missing_user_id", done=True) return json.dumps({"message": message}, ensure_ascii=False) await emitter.emit( description="Retrieving stored memories.", status="recall_in_progress", done=False, ) user_id = __user__.get("id") user_memories = Memories.get_memories_by_user_id(user_id) if not user_memories: message = "No memory stored." await emitter.emit(description=message, status="recall_complete", done=True) return json.dumps({"message": message}, ensure_ascii=False) sorted_memories = sorted(user_memories, key=lambda m: m.created_at) # Reframe first-person text as second-person before returning reframed: List[str] = [] for memory in sorted_memories: s = f" {memory.content} " replacements = [ (" I'm ", " you're "), (" I am ", " you are "), (" I was ", " you were "), (" I ", " you "), (" my ", " your "), (" mine ", " yours "), (" me ", " you "), (" myself ", " yourself "), ] for old, new in replacements: s = s.replace(old, new) reframed.append(s.strip()) await emitter.emit( description=f"{len(sorted_memories)} memories loaded", status="recall_complete", done=True, ) # Build final text and wrap in JSON so the model can just echo final_answer final_text = "\n".join(reframed) return json.dumps({"final_answer": final_text}, ensure_ascii=False) async def add_memory( self, input_text: List[str], __user__: dict = None, __event_emitter__: Callable[[dict], Any] = None, ) -> str: emitter = EventEmitter(__event_emitter__) if not __user__ or not __user__.get("id"): message = "User ID not provided." await emitter.emit(description=message, status="missing_user_id", done=True) return json.dumps({"message": message}, ensure_ascii=False) if isinstance(input_text, str): input_text = [input_text] await emitter.emit( description="Adding entries to the memory vault.", status="add_in_progress", done=False, ) user_id = __user__.get("id") added_items = [] failed_items = [] for item in input_text: new_memory = Memories.insert_new_memory(user_id, item) if new_memory: added_items.append(item) else: failed_items.append(item) if not added_items: message = "Failed to add any memories." await emitter.emit(description=message, status="add_failed", done=True) return json.dumps({"message": message}, ensure_ascii=False) added_count = len(added_items) failed_count = len(failed_items) if failed_count: message = ( f"Added {added_count} memories, " f"failed to add {failed_count} memories." ) else: message = f"Added {added_count} memories." await emitter.emit( description=message, status="add_complete", done=True, ) return json.dumps({"message": message}, ensure_ascii=False) async def delete_memory( self, indices: List[int], __user__: dict = None, __event_emitter__: Callable[[dict], Any] = None, ) -> str: emitter = EventEmitter(__event_emitter__) if not __user__ or not __user__.get("id"): message = "User ID not provided." await emitter.emit(description=message, status="missing_user_id", done=True) return json.dumps({"message": message}, ensure_ascii=False) if isinstance(indices, int): indices = [indices] await emitter.emit( description=f"Deleting {len(indices)} memory entries.", status="delete_in_progress", done=False, ) user_id = __user__.get("id") user_memories = Memories.get_memories_by_user_id(user_id) if not user_memories: message = "No memories found to delete." await emitter.emit(description=message, status="delete_failed", done=True) return json.dumps({"message": message}, ensure_ascii=False) sorted_memories = sorted(user_memories, key=lambda m: m.created_at) responses = [] for index in indices: if index < 1 or index > len(sorted_memories): msg = f"Memory index {index} does not exist." responses.append(msg) await emitter.emit(description=msg, status="delete_failed", done=False) continue memory_to_delete = sorted_memories[index - 1] result = Memories.delete_memory_by_id(memory_to_delete.id) msg = ( f"Memory at index {index} deleted." if result else f"Failed to delete memory at index {index}." ) responses.append(msg) await emitter.emit( description=msg, status="delete_success" if result else "delete_failed", done=False, ) await emitter.emit( description="All deletions processed.", status="delete_complete", done=True, ) return json.dumps({"message": "\n".join(responses)}, ensure_ascii=False) async def update_memory( self, updates: List[dict], __user__: dict = None, __event_emitter__: Callable[[dict], Any] = None, ) -> str: emitter = EventEmitter(__event_emitter__) if not __user__ or not __user__.get("id"): message = "User ID not provided." await emitter.emit(description=message, status="missing_user_id", done=True) return json.dumps({"message": message}, ensure_ascii=False) await emitter.emit( description=f"Updating {len(updates)} memory entries.", status="update_in_progress", done=False, ) user_id = __user__.get("id") user_memories = Memories.get_memories_by_user_id(user_id) if not user_memories: message = "No memories found to update." await emitter.emit(description=message, status="update_failed", done=True) return json.dumps({"message": message}, ensure_ascii=False) sorted_memories = sorted(user_memories, key=lambda m: m.created_at) responses = [] for update_item in updates: index = update_item.get("index") content = update_item.get("content") if index < 1 or index > len(sorted_memories): msg = f"Memory index {index} does not exist." responses.append(msg) await emitter.emit(description=msg, status="update_failed", done=False) continue memory_to_update = sorted_memories[index - 1] updated_memory = Memories.update_memory_by_id(memory_to_update.id, content) msg = ( f"Memory at index {index} updated." if updated_memory else f"Failed to update memory at index {index}." ) responses.append(msg) await emitter.emit( description=msg, status="update_success" if updated_memory else "update_failed", done=False, ) await emitter.emit( description="All updates processed.", status="update_complete", done=True, ) return json.dumps({"message": "\n".join(responses)}, ensure_ascii=False)