mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-16 04:33:12 +00:00
131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
import datetime
|
|
import json
|
|
from typing import List, Any
|
|
from uuid import uuid4
|
|
|
|
from pydantic import BaseModel, Field, model_validator
|
|
|
|
from experiencemaker.enumeration.role import Role
|
|
from experiencemaker.schema.reward import Reward
|
|
|
|
|
|
class ToolCall(BaseModel):
|
|
index: int = Field(default=...)
|
|
id: str = Field(default="")
|
|
name: str = Field(default="")
|
|
arguments: str = Field(default="")
|
|
type: str = Field(default="function")
|
|
result: Any = Field(default=None, exclude=True)
|
|
|
|
@model_validator(mode="before") # noqa
|
|
@classmethod
|
|
def init_tool_call(cls, data: dict):
|
|
tool_type = data.get("type", "")
|
|
tool_type_dict = data.get(tool_type, {})
|
|
|
|
for key in ["name", "arguments"]:
|
|
if key not in data:
|
|
data[key] = tool_type_dict.get(key, "")
|
|
return data
|
|
|
|
@property
|
|
def argument_dict(self):
|
|
return json.loads(self.arguments)
|
|
|
|
@property
|
|
def simple_dict(self):
|
|
return {
|
|
"id": self.id,
|
|
self.type: {
|
|
"arguments": self.arguments,
|
|
"name": self.name
|
|
},
|
|
"type": self.type,
|
|
"index": self.index,
|
|
}
|
|
|
|
@property
|
|
def tool_text_result(self):
|
|
return f"tool.{self.index}\nname={self.name}\narguments={self.arguments}\nresult={str(self.result)}".strip()
|
|
|
|
|
|
class Message(BaseModel):
|
|
role: Role = Field(default=Role.USER)
|
|
content: str | bytes = Field(default="")
|
|
reasoning_content: str = Field(default="")
|
|
tool_calls: List[ToolCall] = Field(default_factory=list)
|
|
timestamp: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
|
|
add_reasoning_content_when_content_is_empty: bool = Field(default=False)
|
|
metadata: dict = Field(default_factory=dict)
|
|
|
|
@property
|
|
def simple_dict(self) -> dict:
|
|
result = {"role": self.role.value, "content": self.content}
|
|
|
|
if self.add_reasoning_content_when_content_is_empty and not self.content:
|
|
result["content"] += self.reasoning_content
|
|
|
|
if self.tool_calls:
|
|
result["tool_calls"] = [x.simple_dict for x in self.tool_calls]
|
|
return result
|
|
|
|
|
|
class ActionMessage(Message):
|
|
role: Role = Field(default=Role.ASSISTANT)
|
|
|
|
|
|
class StateMessage(Message):
|
|
role: Role = Field(default=Role.TOOL)
|
|
tool_call_id: str = Field(default="")
|
|
|
|
def tool_result_to_content(self):
|
|
self.content += "\n\n".join([x.tool_text_result for x in self.tool_calls])
|
|
|
|
@property
|
|
def simple_dict(self) -> dict:
|
|
result = super().simple_dict
|
|
if self.tool_call_id:
|
|
result["tool_call_id"] = self.tool_call_id
|
|
return result
|
|
|
|
@property
|
|
def simple_list(self) -> list:
|
|
return [{
|
|
"role": self.role.value,
|
|
"content": str(x.result),
|
|
"tool_call_id": x.id,
|
|
} for x in self.tool_calls]
|
|
|
|
|
|
class ContextMessage(Message):
|
|
role: Role = Field(default=Role.CONTEXT_ASSISTANT)
|
|
|
|
|
|
class SummaryMessage(Message):
|
|
role: Role = Field(default=Role.SUMMARY_ASSISTANT)
|
|
|
|
|
|
class Sample(BaseModel):
|
|
steps: List[Message] = Field(default_factory=list)
|
|
metadata: dict = Field(default_factory=dict)
|
|
|
|
|
|
class Trajectory(BaseModel):
|
|
id: str = Field(default_factory=lambda: uuid4().hex)
|
|
steps: List[Message] = Field(default_factory=list)
|
|
is_terminated: bool = Field(default=False)
|
|
reward: Reward = Field(default_factory=Reward)
|
|
|
|
query: str = Field(default="")
|
|
answer: str = Field(default="")
|
|
metadata: dict = Field(default_factory=dict)
|
|
|
|
def add_step(self, step: Message):
|
|
self.steps.append(step)
|
|
|
|
def reset(self):
|
|
for name, field in self.model_fields.items():
|
|
if field.default is not None:
|
|
setattr(self, name, field.default)
|
|
elif field.default_factory is not None:
|
|
setattr(self, name, field.default_factory())
|