From ec5b8ccfe1e3ea31e7572ad9c9e35c09216e1004 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=8E=98=E6=9D=83=20=E9=A9=AC?= <mashenquan@fuzhi.ai>
Date: Tue, 30 Apr 2024 12:06:56 +0800
Subject: [PATCH] feat: + `AIMessage`.`with_agent`, `agent`

---
 metagpt/const.py  |  3 +++
 metagpt/schema.py | 11 +++++++++++
 2 files changed, 14 insertions(+)

diff --git a/metagpt/const.py b/metagpt/const.py
index c01f92adc..67636fad3 100644
--- a/metagpt/const.py
+++ b/metagpt/const.py
@@ -139,3 +139,6 @@ LLM_API_TIMEOUT = 300
 
 # Assistant alias
 ASSISTANT_ALIAS = "response"
+
+# Metadata defines
+AGENT = "agent"
diff --git a/metagpt/schema.py b/metagpt/schema.py
index 08b4e3c92..b3600ec98 100644
--- a/metagpt/schema.py
+++ b/metagpt/schema.py
@@ -38,6 +38,7 @@ from pydantic import (
 )
 
 from metagpt.const import (
+    AGENT,
     MESSAGE_ROUTE_CAUSE_BY,
     MESSAGE_ROUTE_FROM,
     MESSAGE_ROUTE_TO,
@@ -342,6 +343,9 @@ class Message(BaseModel):
         m["resources"] = [Resource(**i) for i in m.get("resources", [])]
         return m
 
+    def add_metadata(self, key: str, value: str):
+        self.metadata[key] = value
+
 
 class UserMessage(Message):
     """便于支持OpenAI的消息
@@ -372,6 +376,13 @@ class AIMessage(Message):
         kwargs.pop("role", None)
         super().__init__(content=content, role="assistant", **kwargs)
 
+    def with_agent(self, name: str):
+        self.add_metadata(key=AGENT, value=name)
+
+    @property
+    def agent(self) -> str:
+        return self.metadata.get(AGENT, "")
+
 
 class Task(BaseModel):
     task_id: str = ""