大模型应用（五）模型记忆力代码实战，memory最佳实践

此人未设置昵称
1,737 阅读6分钟

前言

前面的博文中有介绍模型记忆力，理论请见：juejin.cn/post/733543…

首先讲memory是很宽泛的概念，通常理解一个大模型的应用=思考力+行动力+记忆力+设定。通常意义上所谓的memory就是记忆力。

我接触的memory代码，是langchain上看到的，它给了我很大启发，后来还在我们整个团队内部分享，间接促成了我们对于memory的专项化建设。

对记忆力的抽象

简单的将记忆力分三种

  • 短期的，明确的记忆。不管信息有用没有，短期的记忆都应该是清楚地。
  • 长期的，模糊的记忆，记住的信息，可能缺失，也可能不准。就像你记不住三年前你立的flag。
  • 长期的，明确的记忆。比如你的年龄，身高等。可能好几年都不会有人问你年龄，但你通常是能记住的。

针对上面的分析，我们对应如下trait：

pub trait Memory{
    //短期记忆<-------------
    //加载上下文
    async fn load_context(&self,user:&str, max: usize) -> anyhow::Result<Vec<ChatCompletionRequestMessage>>;
    //追加会话日志，可以在上下文中获取到
    async fn add_session_log(&self,user:&str,record: Vec<ChatCompletionRequestMessage>);

    //长期明确记忆<-------------
    //拉取用户标签
    async fn get_user_tag(&self,user:&str,tag:&str) -> anyhow::Result<String>;
    //给用户贴标签
    async fn set_user_tage(&self,user:&str,kvs:HashMap<String,String>);

    //长期模糊记忆<-------------
    //召回长期记忆
    async fn recall_summary(&self,user:&str,query: &str,n:usize)-> anyhow::Result<Vec<String>>;
    //将记忆进行总结
    async fn summary_history(&self,user:&str);
}

代码实现

详细代码传送门

短期记忆

通常我们把短期记忆放在上下文里面，存储的放在任意数据库中。

我这里实现很简单，直接放在了内存里作为一个列表存储，重启就没了。

长期-明确记忆

这种的，我们只需要按照标签存储，如姓名，年龄。

我这里用了一个本地的KV数据库，保证我们重启后，还能够记住这些信息。

async fn add_session_log(&self,mut record: Vec<ChatCompletionRequestMessage>) {
    //fixme 存才并发错误
    #[allow(invalid_reference_casting)]
    unsafe {
        let sc = &mut *(&self.short_context as *const Vec<ChatCompletionRequestMessage> as *mut Vec<ChatCompletionRequestMessage>);
        sc.append(&mut record);
    }
}

async fn get_user_tag(&self,uid:&str, tag: &str) -> anyhow::Result<String> {
    let result = UIB.get(format!("{}-{}",uid,tag).as_str());
    return result
}

长期-模糊记忆

虽然人类的长期记忆是模糊的，但是我们的数据库存了所有对话信息，不管过去多久，实际上我们都能明确的找到用户的会话记录。也就是说我们可以很清楚的记得长期记忆。

但是受限于大模型的tokens，我们不能把这些记忆都放进去，并且很多内容是无效的，这个时候就需要我们对历史进行总结，并且在这些总结中，找到相关的内容。

这里就需要两个基本功能，总结和召回：

  • 总结：总结历史会话我们还是用大模型进行总结，并且定期离线，或者在线发起总结。
  • 召回：我这里用阿里云的DashVector+embedding召回
async fn recall_summary(&self,uid:&str, query: &str, n: usize) -> anyhow::Result<Vec<String>> {
    //这里是topn召回
    self.summary.top_n(uid,query,n).await
}

async fn summary_history(&self,uid:&str){
    let prompt:ChatCompletionRequestMessage = ChatCompletionRequestSystemMessageArgs::default()
        .content("你是阅读理解专家。你的目标是，从过往的对话，提取关键信息，进行总结。总结后的内容必须言简意赅，简洁明了。最总通过`summery_tool`记录下来。")
        .build()
        .unwrap()
        .into();

    let mut context = vec![prompt];
    let mut chat_context = match self.load_context(1000).await{
        Ok(o) => o,
        Err(e) => {
            wd_log::log_field("error",e).error("summary history load context error");
            return;
        }
    };
    context.append(&mut chat_context);

    let summery_tool = ChatCompletionToolArgs::default()
        .r#type(ChatCompletionToolType::Function)
        .function(FunctionObjectArgs::default()
            .name("summery_tool")
            .description("记录总结内容")
            .parameters(json!({
                "type": "object",
                "properties": {
                    "list": {
                        "type": "array",
                        "items":{
                            "type":"string",
                            "description": "条目",
                        },
                        "description": "总结得到的列表",
                    },
                },
                "required": ["location"],
            }))
            .build().unwrap())
        .build().unwrap();


    let chat_req = CreateChatCompletionRequestArgs::default()
        .max_tokens(4096u16)
        .model("gpt-4")
        .messages(context)
        .tools([summery_tool])
        .build()
        .unwrap();

    let client = Client::new();
    let resp = client.chat().create(chat_req).await.unwrap();
    let mut summery_list = vec![];
    for i in resp.choices {
        if i.message.tool_calls.is_some(){
            for j in i.message.tool_calls.unwrap(){
                println!("总结的内容：{}",j.function.arguments);
                if let Ok(sl) = serde_json::from_str::<SummeryList>(j.function.arguments.as_str()) {
                    summery_list = sl.list;
                    break
                }
            }
        }else{
            wd_log::log_field("role",i.message.role)
                .field("message",format!("{:?}",i))
                .info("summery not is tool");
        }
    }
    if let Err(e) = self.summary.insert(uid.to_string(),summery_list).await {
        wd_log::log_field("error",e).error("insert to dash vector error")
    }
}

工具实现

给用户贴标签

为了获取用户的长期准确记忆，我们不可能给用户一个问卷，让用户把所有信息都填一遍，所以，信息的收集应该来自对话。调起一个tool进行收集。所以我们需要实现一个function call。

pub struct UserTagsNode{
    inner:Arc<ShortLongMemoryMap>,
}
impl UserTagsNode {
    //注入到agent中，也就是对自身功能的描述
    pub fn as_openai_tool(&self) -> ChatCompletionTool {
        ChatCompletionToolArgs::default()
            .r#type(ChatCompletionToolType::Function)
            .function(FunctionObjectArgs::default()
                .name(self.id())
                .description("记录用户姓名，年龄")
                .parameters(Value::from_str(r#"{"type":"object","properties":{"tag":{"type":"string","description":"用户标签","enum":["name","age"]},"name":{"type":"string","description":"姓名"},"age":{"type":"integer","description":"年龄"}},"required":["tag"]}"#).unwrap())
                .build().unwrap())
            .build()
            .unwrap()
            .into()
    }
}

#[async_trait::async_trait]
impl Node for UserTagsNode {
    fn id(&self) -> String {
        "user_tag".into()
    }
    
    //按照前面博文的实现，我们需要给tool一个具体的节点实现
    async fn go(&self, ctx: Arc<Context>, mut args: TaskInput) -> anyhow::Result<TaskOutput> {
        let uid = user_id_from_ctx(ctx.as_ref());
        if let Some(s) = args.get_value::<ChatCompletionMessageToolCall>() {
            let input = s.function.arguments;
            let req = serde_json::from_str::<UserTagsNodeReq>(input.as_str())?;
            match req.tag.as_str() {
                "name"=> self.inner.set_user_tage(uid.as_str(),HashMap::from([("name".into(),req.name)])).await,
                "age"=> self.inner.set_user_tage(uid.as_str(),HashMap::from([("age".into(),req.age.to_string())])).await,
                _=> return anyhow::anyhow!("tag[{}] unknown",req.tag).err()
            };
            let msg: ChatCompletionRequestMessage = ChatCompletionRequestToolMessageArgs::default()
                .tool_call_id(s.id)
                .content("success")
                .build()
                .unwrap()
                .into();
            return go_next_or_over(ctx,msg)
        }else{
            anyhow::anyhow!("args not find").err()
        }
    }
}

主动总结

一般来说总结应该是离线的，但我们也应该允许用户主动总结，同样来一个tool支持。

pub struct SummeryNode{
    inner:Arc<ShortLongMemoryMap>,
}
impl SummeryNode {
    //对tool的描述
    pub fn as_openai_tool(&self) -> ChatCompletionTool {
        ChatCompletionToolArgs::default()
            .r#type(ChatCompletionToolType::Function)
            .function(FunctionObjectArgs::default()
                .name(self.id())
                .description("用户主动总结")
                .parameters(Value::from_str(r#"{"type":"object","properties":{"is":{"type":"boolean","description":"是否总结"}}}"#).unwrap())
                .build().unwrap())
            .build()
            .unwrap()
            .into()
    }
}

#[async_trait::async_trait]
impl Node for SummeryNode {
    fn id(&self) -> String {
        "summery".into()
    }
    //具体的执行节点
    async fn go(&self, ctx: Arc<Context>, mut args: TaskInput) -> anyhow::Result<TaskOutput> {
        let uid = user_id_from_ctx(ctx.as_ref());
        let input = args.get_value::<ChatCompletionMessageToolCall>().unwrap();
        self.inner.summary_history(uid.as_str()).await;
        let msg: ChatCompletionRequestMessage = ChatCompletionRequestToolMessageArgs::default()
            .tool_call_id(input.id)
            .content("success")
            .build()
            .unwrap()
            .into();
        go_next_or_over(ctx,msg)
    }
}

测试

将我们的memory，替换到上一篇博文中的single agent中，并且起一个窗口进行交互。具体代码我就不贴了。

详细请见传送门

效果

为了避免短期上下文对效果的影响，我们通过重启的方式丢弃掉短期上下文记忆。

  1. 首先起一个新用户，直接对话，如下图：
    • 这里可以看到prompt中关于用户和总结的信息都是空的
    • 我们在对话中告知了姓名和年龄，并且模型正确的记录了用户标签。

image.png

  1. 退出程序，重新启动，如下图：
    • 上面我们说的用户信息已经被记住，并正确组装到prompt中
    • 我们主动发起总结，记录总结信息

image.png

  1. 退出程序，重新启动，如下图：
    • 可以看到prompt中，总结被召回
    • 和模型对话，无论是年龄，还是过生日都能够正确回答。

image.png

尾语

上面代码中的总结召回还过于粗糙，实际应该在每轮对话中使用。我这里先实现memory的部分，完整版在prompt实现中贴出。

记忆力是大模型应用的最重要的事情之一，还有很多需要思考的东西，如果你有想法欢迎留言或评论。

avatar
文章
阅读
粉丝