有记忆的写作助手更好用

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" # 从云服务控制台获取MemOS_API_KEY
os.environ["OPENAI_API_KEY"] = "sk-xx" # 替换为你自己的API_KEY

conversation_counter = 0

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

class WritingAssistant:
    """AI写作助手，帮助用户写作，具备记忆能力"""
    
    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):
      """查询相关记忆"""
      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):
        """构建包含格式化记忆的系统提示"""
        base_prompt = """
          你是一位专业的的写作助手, 能够记住用户的写作风格和偏好。
          你可调用对话记忆，助力提供更具个性化的回复。
          请借助这些记忆，理解用户的场景背景、偏好倾向及过往互动情况。
          如果提供了记忆，请在相关时自然地引用它们，但不要明确提及拥有记忆功能
        """

        if memories:
            # 将记忆格式化为编号列表
            formatted_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):
      """添加消息"""
      self.memos_client.add_message(messages, user_id, conversation_id)

    def get_message(self, user_id, conversation_id):
      """获取消息"""
      response = self.memos_client.get_message(user_id, conversation_id)

      return response.data.message_detail_list

    def chat(self, query, user_id, conversation_id):
        """处理包含记忆集成的对话的主要聊天函数"""
        # 1. 搜索相关记忆
        memories = self.search_memory(query, user_id, conversation_id)
        
        # 构建包含记忆的系统提示
        system_prompt = self.build_system_prompt(memories)
        
        # 2. 使用OpenAI生成回答
        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. 将对话保存到记忆中
        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 [
      "帮我写一封团队聚餐的通知邮件",
      "帮我写一封客户邮件，内容是总结我们即将上线的理财App新功能",
    ]

def pre_configured_conversations():
    """返回预配置的对话对"""
    return [
        {
            "user": "我在互联网公司市场部工作，写邮件时语气要轻松点，开头加'亲爱的XX'"
        },
        {
            "user": "写总结时我习惯先列三条要点"
        }
    ]

def execute_pre_conversations():
    """执行预配置的对话"""
    conversations = pre_configured_conversations()
    conversation_id = generate_conversation_id()
    
    print(f"\n🔄 正在执行预配置对话（conversation_id={conversation_id})...")
    print("=" * 60)
    
    for i, conv in enumerate(conversations, 1):
        print(f"\n💬 对话 {i}")
        print(f"👤 用户: {conv['user']}")
        
        # 执行对话
        answer = ai_assistant.chat(conv['user'], user_id, conversation_id)
        print(f"🤖 助手: {answer}")
        print("-" * 40)
    
    print("\n✅ 预配置对话执行完毕！")
    print("=" * 60)

def main():
    print("📝 欢迎查看MemOS在写作助手中的使用示例！")
    print("💡 在MemOS的加持下，让您的写作助手更懂您的风格和偏好！ ✍️ \n")
    
    # 询问用户是否要先执行预配置对话
    while True:
        pre_chat = input("🤔 您想先执行预配置对话吗？预计消耗2次add和2次search的调用额度，是否执行？(y/n): ").strip().lower()
        
        if pre_chat in ['y', 'yes', '是', 'Y']:
            execute_pre_conversations()
            break
        elif pre_chat in ['n', 'no', '否', 'N']:
            print("📝 开始全新的写作助手对话...")
            break
        else:
            print("⚠️  请输入 'y' 表示是或 'n' 表示否")

    print("\n⚡️ 接下来您输入的每一个问题，都将会在一个全新的会话中展开（新的 conversation id）。MemOS将会跨会话自动召回您的历史行为记忆，为您提供连续、个性化的服务。")    
    print("\n🎯 以下是一些示例问题，您可以继续跟写作助手对话:")
    for i, question in enumerate(demo_questions(), 1):
      print(f"  {i}. {question}")

    while True:
        user_query = input("\n🤔 请输入您的写作需求 (或输入 'exit' 退出): ").strip()
        
        if user_query.lower() in ['quit', 'exit', 'q', '退出']:
            print("👋 感谢使用写作助手，祝您写作愉快！")
            break
        
        if not user_query:
            continue
        
        print("🤖 正在创作...")
        conversation_id = generate_conversation_id()
        answer = ai_assistant.chat(user_query, user_id, conversation_id)
        print(f"\n💬 conversation_id: {conversation_id}\n💡 [助手]: {answer}\n")
        print("-" * 60)


if __name__ == "__main__":
    main()