ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworld/memory/models.py

@@ -0,0 +1,499 @@
+import uuid
+from abc import abstractmethod
+from datetime import datetime
+
+from pydantic import BaseModel, ConfigDict, Field
+from typing import Any, Dict, List, Optional, Literal
+
+from aworld.models.model_response import ToolCall
+
+class MemoryItem(BaseModel):
+    id: str = Field(description="id")
+    content: Any = Field(description="content")
+    created_at: Optional[str] = Field(None, description="created at")
+    updated_at: Optional[str] = Field(None, description="updated at")
+    metadata: dict = Field(
+        description="metadata, use to store additional information, such as user_id, agent_id, run_id, task_id, etc.")
+    tags: list[str] = Field(description="tags")
+    histories: list["MemoryItem"] = Field(default_factory=list)
+    deleted: bool = Field(default=False)
+    memory_type: Literal["init", "message", "summary", "agent_experience", "user_profile", "fact", "conversation_summary"] = Field(default="message")
+    version: int = Field(description="version")
+
+    def __init__(self, **data):
+        # Set default values for optional fields
+        if "id" not in data:
+            data["id"] = str(uuid.uuid4())
+        if "created_at" not in data:
+            data["created_at"] = datetime.now().isoformat()
+        if "updated_at" not in data:
+            data["updated_at"] = data["created_at"]
+        if "metadata" not in data:
+            data["metadata"] = {}
+        if "tags" not in data:
+            data["tags"] = []
+        if "version" not in data:
+            data["version"] = 1
+
+        super().__init__(**data)
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "MemoryItem":
+        """Create a MemoryItem instance from a dictionary.
+
+        Args:
+            data (dict): A dictionary containing the memory item data.
+
+        Returns:
+            MemoryItem: An instance of MemoryItem.
+        """
+        return cls(**data)
+
+    @property
+    def user_id(self) -> str:
+        return self.metadata.get('user_id')
+
+    @property
+    def session_id(self) -> str:
+        return self.metadata.get('session_id')
+
+    @property
+    def task_id(self) -> str:
+        return self.metadata.get('task_id')
+
+    @property
+    def agent_id(self) -> str:
+        return self.metadata.get('agent_id')
+
+    @property
+    def agent_name(self) -> str:
+        return self.metadata.get('agent_name')
+
+    @property
+    def application_id(self) -> str:
+        return self.metadata.get('application_id', 'default')
+
+    @property
+    def embedding_text(self) -> Optional[str]:
+        return self.content
+
+    def mark_has_summary(self):
+        self.metadata['summary'] = True
+
+    @property
+    def has_summary(self) -> bool:
+        return self.metadata.get('summary', False)
+    
+    @property
+    def content_length(self) -> int:
+        return len(self.content)
+
+    @property
+    def status(self) -> str:
+        return self.metadata.get('status', 'ACCEPTED')
+
+    @status.setter
+    def status(self, value: Literal["DRAFT", "ACCEPTED", "DISCARD"]) -> None:
+        self.metadata['status'] = value
+
+    @abstractmethod
+    def to_openai_message(self) -> dict:
+        pass
+
+
+class MessageMetadata(BaseModel):
+    """
+    Metadata for memory messages, including user, session, task, and agent information.
+    Args:
+        user_id (str): The ID of the user.
+        session_id (str): The ID of the session.
+        task_id (str): The ID of the task.
+        agent_id (str): The ID of the agent.
+    """
+    agent_id: str = Field(description="The ID of the agent")
+    agent_name: Optional[str] = Field(description="The name of the agent")
+    session_id: Optional[str] = Field(default=None,description="The ID of the session")
+    task_id: Optional[str] = Field(default=None,description="The ID of the task")
+    user_id: Optional[str] = Field(default=None, description="The ID of the user")
+    summary_content: Optional[str] = Field(default=None, description="The summary of the memory item")
+
+    model_config = ConfigDict(extra="allow")
+
+    @property
+    def to_dict(self) -> Dict[str, Any]:
+        return self.model_dump()
+
+class AgentExperienceItem(BaseModel):
+    skill: str = Field(description="The skill demonstrated in the experience")
+    actions: List[str] = Field(description="The actions taken by the agent")
+
+
+class AgentExperience(MemoryItem):
+    """
+    Represents an agent's experience, including skills and actions.
+    All custom attributes are stored in content and metadata.
+    Args:
+        agent_id (str): The ID of the agent.
+        skill (str): The skill demonstrated in the experience.
+        actions (List[str]): The actions taken by the agent.
+        metadata (Optional[Dict[str, Any]]): Additional metadata.
+    """
+    def __init__(self, agent_id: str, skill: str, actions: List[str], metadata: Optional[Dict[str, Any]] = None) -> None:
+        meta = metadata.copy() if metadata else {}
+        meta['agent_id'] = agent_id
+        agent_experience = AgentExperienceItem(skill=skill, actions=actions)
+        super().__init__(content=agent_experience, metadata=meta, memory_type="agent_experience")
+
+    @property
+    def agent_id(self) -> str:
+        return self.metadata['agent_id']
+
+    @property
+    def skill(self) -> str:
+        return self.content.skill
+
+    @property
+    def actions(self) -> List[str]:
+        return self.content.actions
+
+    @property
+    def embedding_text(self):
+        return f"skill:{self.skill}, actions:{self.actions}"
+
+    def to_openai_message(self) -> dict:
+        return {
+            "role": "system",
+            "content": self.content
+        }
+
+
+class UserProfileItem(BaseModel):
+    key: str = Field(description="The key of the profile")
+    value: Any = Field(description="The value of the profile")
+
+class UserProfile(MemoryItem):
+    """
+    Represents a user profile key-value pair.
+    All custom attributes are stored in content and metadata.
+    Args:
+        user_id (str): The ID of the user.
+        key (str): The profile key.
+        value (Any): The profile value.
+        metadata (Optional[Dict[str, Any]]): Additional metadata.
+    """
+    def __init__(self, user_id: str, key: str, value: Any, metadata: Optional[Dict[str, Any]] = None, **kwargs) -> None:
+        meta = metadata.copy() if metadata else {}
+        meta['user_id'] = user_id
+        user_profile = UserProfileItem(key=key, value=value)
+        super().__init__(content=user_profile, metadata=meta, memory_type="user_profile", **kwargs)
+
+    @property
+    def user_id(self) -> str:
+        return self.metadata['user_id']
+
+    @property
+    def key(self) -> str:
+        return self.content.key
+
+    @property
+    def value(self) -> Any:
+        return self.content.value
+
+    @property
+    def item(self) -> UserProfileItem:
+        return self.content
+
+    @property
+    def embedding_text(self):
+        return f"key:{self.key} value:{self.value}"
+    
+    def to_openai_message(self) -> dict:
+        return {
+            "role": "system",
+            "content": self.content
+        }
+
+class Fact(MemoryItem):
+    """
+    Represents Fact from conversation.
+    Args:
+        user_id (str): The ID of the user.
+        content (str): fact.
+        metadata (Optional[Dict[str, Any]]): Additional metadata.
+    """
+    def __init__(self, user_id: str = None, agent_id: str = None, content: str = None, metadata: Optional[Dict[str, Any]] = None, **kwargs) -> None:
+        meta = metadata.copy() if metadata else {}
+        if user_id:
+            meta['user_id'] = user_id
+        elif metadata.get('user_id'):
+            meta['user_id'] = metadata.get('user_id')
+
+        if 'memory_type' in kwargs:
+            kwargs.pop("memory_type")
+        super().__init__(content=content, metadata=meta, memory_type="fact", **kwargs)
+
+    @property
+    def key(self) -> str:
+        return self.content.key
+
+    @property
+    def value(self) -> Any:
+        return self.content.value
+
+    @property
+    def embedding_text(self):
+        return self.content
+
+    def to_openai_message(self) -> dict:
+        return {
+            "role": "user",
+            "content": self.content
+        }
+
+class MemorySummary(MemoryItem):
+    """
+    Represents a memory summary.
+    All custom attributes are stored in content and metadata.
+    Args:
+        item_ids (str): The IDS of the agent.
+        summary (str): The summary text.
+        metadata (Optional[Dict[str, Any]]): Additional metadata.
+    """
+    def __init__(self, item_ids: list[str], summary: str, metadata: MessageMetadata, **kwargs) -> None:
+        meta = metadata.to_dict
+        meta['item_ids'] = item_ids
+        meta['role'] = "user"
+        super().__init__(content=summary, metadata=meta, memory_type="summary", **kwargs)
+
+    @property
+    def summary_item_ids(self):
+        return self.metadata['item_ids']
+
+    def to_openai_message(self) -> dict:
+        return {
+            "role": "user",
+            "content": self.content
+        }
+
+
+class ConversationSummary(MemoryItem):
+    """
+    Represents a conversation summary.
+    All custom attributes are stored in content and metadata.
+    Args:
+        user_id (str): The ID of the user.
+        session_id (str): The ID of the session.
+        summary (str): The summary text of the conversation.
+        metadata (MessageMetadata): Metadata object containing additional information.
+    """
+
+    def __init__(self, user_id: str, session_id: str, summary: str, metadata: MessageMetadata, **kwargs) -> None:
+        meta = metadata.to_dict
+        meta['user_id'] = user_id
+        meta['session_id'] = session_id
+        super().__init__(content=summary, metadata=meta, memory_type="conversation_summary", **kwargs)
+
+    def to_openai_message(self) -> dict:
+        return {
+            "role": "assistant",
+            "content": self.content
+        }
+
+
+class MemoryMessage(MemoryItem):
+    """
+    Represents a memory message with role, user, session, task, and agent information.
+    Args:
+        role (str): The role of the message sender.
+        metadata (MessageMetadata): Metadata object containing user, session, task, and agent IDs.
+        content (Optional[Any]): Content of the message.
+    """
+    def __init__(self, role: str, metadata: MessageMetadata, content: Optional[Any] = None, memory_type="message", **kwargs) -> None:
+        meta = metadata.to_dict
+        meta['role'] = role
+        super().__init__(content=content, metadata=meta, memory_type=memory_type, **kwargs)
+
+    @property
+    def role(self) -> str:
+        return self.metadata['role']
+
+    @property
+    def user_id(self) -> str:
+        return self.metadata['user_id']
+
+    @property
+    def session_id(self) -> str:
+        return self.metadata['session_id']
+
+    @property
+    def task_id(self) -> str:
+        return self.metadata['task_id']
+
+    def set_task_id(self, task_id):
+        self.metadata['task_id'] = task_id
+
+    @property
+    def agent_id(self) -> str:
+        return self.metadata['agent_id']
+    
+    @abstractmethod
+    def to_openai_message(self) -> dict:
+        pass
+
+class MemorySystemMessage(MemoryMessage):
+    """
+    Represents a system message with role and content.
+    Args:
+        metadata (MessageMetadata): Metadata object containing user, session, task, and agent IDs.
+        content (str): The content of the message.
+    """
+    def __init__(self, content: str, metadata: MessageMetadata, **kwargs) -> None:
+        super().__init__(role="system", metadata=metadata, content=content, memory_type="init", **kwargs)
+
+    def to_openai_message(self) -> dict:
+        return {
+            "role": self.role,
+            "content": self.content
+        }
+
+    @property
+    def embedding_text(self) -> Optional[str]:
+        return None
+
+
+class MemoryHumanMessage(MemoryMessage):
+    """
+    Represents a human message with role and content.
+    Args:
+        metadata (MessageMetadata): Metadata object containing user, session, task, and agent IDs.
+        content (str): The content of the message.
+    """
+    def __init__(self, metadata: MessageMetadata, content: Any, memory_type = "init", **kwargs) -> None:
+        super().__init__(role="user", metadata=metadata, content=content, memory_type=memory_type, **kwargs)
+    
+    def to_openai_message(self) -> dict:
+        return {
+            "role": self.role,
+            "content": self.content
+        }
+
+class MemoryAIMessage(MemoryMessage):
+    """
+    Represents an AI message with role and content.
+    Args:
+        metadata (MessageMetadata): Metadata object containing user, session, task, and agent IDs.
+        content (str): The content of the message.
+    """
+    def __init__(self, content: str, tool_calls: Optional[List[ToolCall]] = [], metadata: MessageMetadata = None, **kwargs) -> None:
+        meta = metadata.to_dict
+        if tool_calls:
+            meta['tool_calls'] = [tool_call.to_dict() for tool_call in tool_calls]
+        super().__init__(role="assistant", metadata=MessageMetadata(**meta), content=content, **kwargs)
+
+    @property
+    def tool_calls(self) -> List[ToolCall]:
+        if "tool_calls" not in self.metadata or not self.metadata['tool_calls']:
+            return None
+        tc = [ToolCall(**tool_call) for tool_call in self.metadata['tool_calls']]
+        return tc if len(tc) > 0 else None
+      
+    def to_openai_message(self) -> dict:
+        return {
+            "role": self.role,
+            "content": self.content,
+            "tool_calls": [tool_call.to_dict() for tool_call in self.tool_calls or []] or None
+        }
+
+class MemoryToolMessage(MemoryMessage):
+    """
+    Represents a tool message with role, content, tool_call_id, and status.
+    Args:
+        metadata (MessageMetadata): Metadata object containing user, session, task, and agent IDs.
+        tool_call_id (str): The ID of the tool call.
+        status (Literal["success", "error"]): The status of the tool call.
+        content (str): The content of the message.
+    """
+    def __init__(self, tool_call_id: str, content: Any, status: Literal["success", "error"] = "success", metadata: MessageMetadata = None, **kwargs) -> None:
+        metadata.tool_call_id = tool_call_id
+        metadata.status = status
+        super().__init__(role="tool", metadata=metadata, content=content, **kwargs)
+
+    @property
+    def tool_call_id(self) -> str:
+        return self.metadata['tool_call_id']
+
+    @property
+    def status(self) -> str:
+        return self.metadata['status']
+
+    @property
+    def embedding_text(self) -> Optional[str]:
+        return None
+
+    def to_openai_message(self) -> dict:
+        return {
+            "role": self.role,
+            "content": self.content,
+            "tool_call_id": self.tool_call_id,
+        }
+
+
+class LongTermExtractParams(BaseModel):
+    session_id: str = Field(description="The ID of the session")
+    task_id: Optional[str] = Field(description="The ID of the task")
+    memories: List[MemoryItem] = Field(default_factory=list, description="The list of memories to process")
+
+    application_id: Optional[str] = Field(default=None, description="The ID of the application")
+    extract_type: Literal["user_profile", "agent_experience"] = Field(description="The type of long-term extract")
+
+    def to_openai_messages(self) -> List[dict]:
+        return [memory.to_openai_message() for memory in self.memories]
+
+class UserProfileExtractParams(LongTermExtractParams):
+    user_id: Optional[str] = Field(description="The ID of the user")
+
+    def __init__(self, user_id: str, session_id: str, task_id: str, memories: List[MemoryItem] = None, application_id: str = None, **kwargs) -> None:
+        kwargs = {
+            "user_id": user_id,
+            "session_id": session_id,
+            "task_id": task_id,
+            "memories": memories or [],
+            "application_id": application_id,
+            "extract_type": "user_profile",
+            **kwargs
+        }
+        super().__init__(**kwargs)
+
+    model_config = ConfigDict(extra="allow")
+
+class AgentExperienceExtractParams(LongTermExtractParams):
+    agent_id: str = Field(default=None, description="The ID of the agent")
+
+    def __init__(self, agent_id: str, session_id: str, task_id: str, memories: List[MemoryItem] = None,
+                 application_id: str = None,**kwargs) -> None:
+        super().__init__(session_id=session_id,
+                         task_id=task_id,
+                         memories=memories,
+                         application_id=application_id,
+                         extract_type="agent_experience", **kwargs)
+        self.agent_id = agent_id
+
+    model_config = ConfigDict(extra="allow")
+
+class LongTermMemoryTriggerParams(BaseModel):
+    """
+    Metadata for memory messages, including user, session, task, and agent information.
+    Args:
+        user_id (str): The ID of the user.
+        session_id (str): The ID of the session.
+        task_id (str): The ID of the task.
+        agent_id (str): The ID of the agent.
+    """
+    agent_id: str = Field(default=None, description="The ID of the agent")
+    session_id: str = Field(default=None, description="The ID of the session")
+    task_id: str = Field(default=None, description="The ID of the task")
+    user_id: Optional[str] = Field(default=None, description="The ID of the user")
+    application_id: Optional[str] = Field(default="default", description="The ID of the application, namespace for memory")
+    force: Optional[bool] = Field(default=False, description="Whether to force trigger long-term memory")
+
+    model_config = ConfigDict(extra="allow")