A Writing Assistant with Memory is More Useful

pip install MemoryOS -U

import os
import uuid
from openai import OpenAI
from memos.api.client import MemOSClient

os.environ["MEMOS_API_KEY"] = "mpg-xx" # Get MemOS_API_KEY from cloud service console
os.environ["OPENAI_API_KEY"] = "sk-xx" # Replace with your own API_KEY

conversation_counter = 0

def generate_conversation_id():
    global conversation_counter
    conversation_counter += 1
    return f"conversation_{conversation_counter:03d}"

class WritingAssistant:
    """AI Writing Assistant, helps users write with memory capability"""
    
    def __init__(self):
        self.memos_client = MemOSClient(api_key=os.getenv("MEMOS_API_KEY"))
        self.openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    
    def search_memory(self, query, user_id, conversation_id):
        """Search relevant memories based on query"""
        response = self.memos_client.search_memory(query, user_id, conversation_id)
        return [memory_detail.memory_value for memory_detail in response.data.memory_detail_list]

    def build_system_prompt(self, memories):
        """Build a system prompt that includes formatted memories"""
        base_prompt = """
          You are a professional writing assistant who can remember the user’s writing style and preferences.
          You can call conversation memories to provide more personalized replies.
          Please use these memories to understand the user’s background, preferences, and past interactions.
          If memories are provided, naturally reference them where relevant, but do not explicitly mention having memory capabilities.
        """

        if memories:
            # Format memories as a numbered list
            formatted_memories = "## Memories:\n"
            for i, memory in enumerate(memories, 1):
                formatted_memories += f"{i}. {memory}\n"
            
            return f"{base_prompt}\n\n{formatted_memories}"
        else:
            return base_prompt
        

    def add_message(self, messages, user_id, conversation_id):
        """Add messages to MemOS so they can be processed into memories"""
        self.memos_client.add_message(messages, user_id, conversation_id)

    def get_message(self, user_id, conversation_id):
        """Retrieve the raw messages stored in MemOS (for debugging/inspection)"""
        response = self.memos_client.get_message(user_id, conversation_id)
        return response.data.message_detail_list

    def chat(self, query, user_id, conversation_id):
        """Main chat function to handle dialogue with memory integration"""
        # 1. Search relevant memories
        memories = self.search_memory(query, user_id, conversation_id)
        
        # Build system prompt with memories
        system_prompt = self.build_system_prompt(memories)
        
        # 2. Use OpenAI to generate an answer
        response = self.openai_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": query}
            ]
        )
        answer = response.choices[0].message.content

        # 3. Save the conversation into memories
        messages = [
            {"role": "user", "content": query},
            {"role": "assistant", "content": answer}
        ]
        self.memos_client.add_message(messages, user_id, conversation_id)
        
        return answer

ai_assistant = WritingAssistant()
user_id = "memos_writing_user_123"

def demo_questions():
    return [
      "Help me write a notification email for a team dinner",
      "Help me write a client email summarizing the new features of our upcoming finance app",
    ]

def pre_configured_conversations():
    """Return pre-configured dialogue pairs"""
    return [
        {
            "user": "I work in the marketing department at an internet company. Keep the tone light when writing emails, and start with 'Dear XX'."
        },
        {
            "user": "When writing summaries, I prefer to list three bullet points first."
        }
    ]

def execute_pre_conversations():
    """Execute pre-configured dialogues"""
    conversations = pre_configured_conversations()
    conversation_id = generate_conversation_id()
    
    print(f"\n🔄 Executing pre-configured dialogues（conversation_id={conversation_id}）...")
    print("=" * 60)
    
    for i, conv in enumerate(conversations, 1):
        print(f"\n💬 Dialogue {i}")
        print(f"👤 User: {conv['user']}")
        
        # Execute dialogue
        answer = ai_assistant.chat(conv['user'], user_id, conversation_id)
        print(f"🤖 Assistant: {answer}")
        print("-" * 40)
    
    print("\n✅ Pre-configured dialogues completed!")
    print("=" * 60)

def main():
    print("📝 Welcome to the MemOS writing assistant example!")
    print("💡 With MemOS, your writing assistant better understands your style and preferences! ✍️ \n")
    
    # Ask whether to execute pre-configured dialogues first
    while True:
        pre_chat = input("🤔 Do you want to execute the pre-configured dialogues first? This will consume 2 add and 2 search calls. Execute? (y/n): ").strip().lower()
        
        if pre_chat in ['y', 'yes', 'Y']:
            execute_pre_conversations()
            break
        elif pre_chat in ['n', 'no', 'N']:
            print("📝 Starting a brand-new writing assistant dialogue...")
            break
        else:
            print("⚠️ Please enter 'y' for yes or 'n' for no")

    print("\n⚡️ Each question you enter next will take place in a brand-new conversation (with a new conversation ID). MemOS will automatically recall your historical behavioral memories across conversations to provide you with continuous and personalized service.")    
    print("\n🎯 Here are some example questions. You can continue chatting with the writing assistant:")
    for i, question in enumerate(demo_questions(), 1):
        print(f"  {i}. {question}")

    while True:
        user_query = input("\n🤔 Please enter your writing request (or type 'exit' to quit): ").strip()
        
        if user_query.lower() in ['quit', 'exit', 'q']:
            print("👋 Thank you for using the writing assistant. Happy writing!")
            break
        
        if not user_query:
            continue
        
        print("🤖 Creating...")
        conversation_id = generate_conversation_id()
        answer = ai_assistant.chat(user_query, user_id, conversation_id)
        print(f"\n💬 conversation_id: {conversation_id}\n💡 [Assistant]: {answer}\n")
        print("-" * 60)


if __name__ == "__main__":
    main()