feat: + AIMessage.with_agent, agent

This commit is contained in:
莘权 马 2024-04-30 12:06:56 +08:00
parent d09fcd4703
commit ec5b8ccfe1
2 changed files with 14 additions and 0 deletions

View file

@ -139,3 +139,6 @@ LLM_API_TIMEOUT = 300
# Assistant alias
ASSISTANT_ALIAS = "response"
# Metadata defines
AGENT = "agent"

View file

@ -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 = ""