Upload 2 files

aworld/output/storage/artifact_repository.py

@@ -0,0 +1,286 @@
+import json
+import os
+import time
+import uuid
+from enum import Enum
+from pathlib import Path
+from typing import Dict, Any, Optional, List
+
+from pydantic import BaseModel
+
+from aworld.output import Artifact
+
+
+class ArtifactRepository:
+    def __init__(self):
+        """
+        Initialize the artifact repository
+        """
+        pass
+
+    def load_index(self) -> Dict[str, Any]:
+        """Load or create index file"""
+        pass
+
+    def save_index(self, index_data: Dict[str, Any]) -> None:
+        """Save index to file"""
+
+    def store_artifact(self,
+                       artifact: Artifact
+                       ) -> str:
+        """
+        Store artifact and return its version identifier
+
+        Args:
+            artifact_id: Unique identifier of the artifact
+            data: Data to be stored
+            metadata: Optional metadata
+
+        Returns:
+            Version identifier
+        """
+        pass
+
+    def delete_artifact(self, artifact_id: str) -> None:
+        """
+        Delete artifact from repository
+        """
+        pass
+
+    def retrieve_latest_artifact(self, artifact_id: str) -> Optional[Dict[str, Any]]:
+        pass
+
+    def get_artifact_versions(self, artifact_id: str) -> List[Dict[str, Any]]:
+        """
+        Get information about all versions of an artifact
+
+        Args:
+            artifact_id: Artifact identifier
+
+        Returns:
+            List of version information
+        """
+        pass
+    def super_path(self) -> str:
+        pass
+
+    def artifact_path(self, artifact_id):
+        return self.super_path() + f"/artifact/{artifact_id}/index.json"
+
+    def generate_tree_data(self, workspace_name: str) -> dict:
+        """
+        Abstract method: Generate a directory tree structure for the workspace.
+        Args:
+            workspace_name: Name of the workspace (for root node)
+        Returns:
+            Directory tree as dict
+        """
+        raise NotImplementedError()
+
+
+class CommonEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, Enum):
+            return obj.name
+        if isinstance(obj, BaseModel):
+            return obj.model_dump_json()
+        return json.JSONEncoder.default(self, obj)
+
+class EnumDecoder(json.JSONDecoder):
+    def decode(self, s, **kwargs):
+        parsed_json = super().decode(s, **kwargs)
+        for key, value in parsed_json.items():
+            if isinstance(value, dict) and value.get("__enum__"):
+                enum_type = globals()[value["__enum_type__"]]
+                enum_value = enum_type[value["__enum_value__"]]
+                parsed_json[key] = enum_value
+        return parsed_json
+
+
+class LocalArtifactRepository(ArtifactRepository):
+    """Artifact storage layer: manages versioned artifacts through content-addressable storage"""
+
+    def __init__(self, storage_path: str):
+        """
+        Initialize the artifact repository
+        
+        Args:
+            storage_path: Directory path for storing data
+        """
+        super().__init__()
+        self.storage_path = Path(storage_path)
+        self.storage_path.mkdir(parents=True, exist_ok=True)
+        self.index_path = self.storage_path / "index.json"
+        self.index = self.load_index()
+
+    def load_index(self) -> Dict[str, Any]:
+        """Load or create index file"""
+        if self.index_path.exists():
+            try:
+                with open(self.index_path, 'r') as f:
+                    return json.load(f)
+            except json.JSONDecodeError:
+                return {"artifacts": [], "versions": []}
+        else:
+            index = {"artifacts": [], "versions": []}
+            self._save_index(index)
+            return index
+
+    def save_index(self, workspace_data) -> None:
+        self._save_index(workspace_data)
+
+    def _save_index(self, index: Dict[str, Any]) -> None:
+        """Save index to file"""
+        with open(self.index_path, 'w') as f:
+            json.dump(index, f, indent=2, ensure_ascii=False, cls=CommonEncoder)
+
+    def store_artifact(self,
+                       artifact: Artifact
+                       ) -> str:
+        """
+        Store artifact and return its version identifier
+        
+        Args:
+            artifact: Artifact to be stored
+            
+        Returns:
+            Version identifier
+        """
+        # Create version record
+        version = {
+            "hash": artifact.artifact_id,
+            "timestamp": time.time(),
+            "metadata": artifact.metadata or {}
+        }
+
+        # Update index
+
+        data = artifact.to_dict()
+        # Store content
+        content_path = Path(self.artifact_path(artifact.artifact_id))
+        content_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(content_path, 'w') as f:
+            json.dump(data, f, indent=2, ensure_ascii=False, cls=CommonEncoder)
+
+        if artifact.attachments:
+            for attachment in artifact.attachments:
+                attachment_path = content_path.parent / attachment.filename
+                attachment_path.parent.mkdir(parents=True, exist_ok=True)
+                with open(attachment_path, 'w') as f:
+                    f.write(attachment.content)
+
+
+
+        return "success"
+
+    def retrieve_latest_artifact(self, artifact_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve artifact based on version ID
+
+        Args:
+            version_id: Version identifier
+
+        Returns:
+            Stored data, or None if it doesn't exist
+        """
+        artifact_path = self.artifact_path(artifact_id)
+
+        if not Path(artifact_path).exists():
+            return None
+
+        with open(artifact_path, 'r') as f:
+            return json.load(f)
+
+    def get_artifact_versions(self, artifact_id: str) -> List[Dict[str, Any]]:
+        """
+        Get information about all versions of an artifact
+        
+        Args:
+            artifact_id: Artifact identifier
+            
+        Returns:
+            List of version information
+        """
+        if artifact_id not in self.index["artifacts"]:
+            return []
+
+        versions = []
+        for version_id in self.index["artifacts"][artifact_id]:
+            version_info = self.index["versions"][version_id].copy()
+            version_info["id"] = version_id
+            versions.append(version_info)
+
+        return versions
+    
+    def delete_artifact(self, artifact_id: str) -> bool:
+        """
+        Delete the specified artifact and its attachments from storage
+        Args:
+            artifact_id: Artifact identifier
+        Returns:
+            Whether deletion was successful
+        """
+        content_path = Path(self.artifact_path(artifact_id))
+        if not Path(content_path).exists():
+            return False
+            
+        # Delete the artifact file
+        os.remove(content_path)
+        
+        # Delete the artifact directory and all its contents
+        artifact_dir = content_path.parent
+        if artifact_dir.exists():
+            for item in artifact_dir.glob('**/*'):
+                if item.is_file():
+                    os.remove(item)
+            for item in reversed(list(artifact_dir.glob('**/*'))):
+                if item.is_dir():
+                    os.rmdir(item)
+            os.rmdir(artifact_dir)
+        
+        return True
+
+    def super_path(self):
+        return str(self.storage_path)
+
+    def generate_tree_data(self, workspace_name: str) -> dict:
+        """
+        Generate a directory tree structure based on the actual local workspace folder structure.
+        Args:
+            workspace_name: Name of the workspace (for root node)
+        Returns:
+            Directory tree as dict
+        """
+        import os
+        def build_tree(path: str, parent_id: str, depth: int = 1) -> dict:
+            import uuid
+            node = {
+                "name": os.path.basename(path) or workspace_name,
+                "id": str(uuid.uuid4()),
+                "type": "dir" if os.path.isdir(path) else "file",
+                "parentId": parent_id,
+                "depth": depth,
+                "expanded": False,
+                "children": []
+            }
+            if os.path.isdir(path):
+                for entry in sorted(os.listdir(path)):
+                    full_path = os.path.join(path, entry)
+                    node["children"].append(build_tree(full_path, node["id"], depth + 1))
+            return node
+        root_path = str(self.storage_path)
+        import os
+        if not os.path.exists(root_path):
+            return {
+                "name": workspace_name,
+                "id": "-1",
+                "children": []
+            }
+        tree = build_tree(root_path, "-1", 1)
+        tree["name"] = workspace_name
+        tree["id"] = "-1"
+        tree["parentId"] = None
+        tree["depth"] = 0
+        return tree
+
+

aworld/output/storage/oss_artifact_repository.py

@@ -0,0 +1,283 @@
+import hashlib
+import json
+import time
+import uuid
+from typing import Dict, Any, Optional, List, Literal
+from .artifact_repository import ArtifactRepository, CommonEncoder
+from aworld.output.artifact import Artifact, ArtifactAttachment
+
+
+class OSSArtifactRepository(ArtifactRepository):
+    """
+    Artifact storage implementation based on Alibaba Cloud OSS, similar to LocalArtifactRepository but using OSS as backend.
+    """
+    def __init__(self,
+                 access_key_id: str,
+                 access_key_secret: str,
+                 endpoint: str,
+                 bucket_name: str,
+                 storage_path: str = "aworld/workspaces/"):
+        """
+        Initialize OSS artifact repository
+        Args:
+            access_key_id: OSS access key ID
+            access_key_secret: OSS access key secret
+            endpoint: OSS service endpoint
+            bucket_name: OSS bucket name
+            storage_path: Storage prefix, defaults to "aworld/workspaces/"
+        """
+        import oss2
+
+        super().__init__()
+        self.auth = oss2.Auth(access_key_id, access_key_secret)
+        self.bucket = oss2.Bucket(self.auth, endpoint, bucket_name)
+        self.prefix = storage_path.rstrip('/') + '/'
+        self.index_key = f"{self.prefix}index.json"
+        self.index = self.load_index()
+
+    def load_index(self) -> Dict[str, Any]:
+        """
+        Load or create index file from OSS
+        Returns:
+            Index dictionary
+        """
+        import oss2
+
+        try:
+            result = self.bucket.get_object(self.index_key)
+            content = result.read().decode('utf-8')
+            return json.loads(content)
+        except oss2.exceptions.NoSuchKey:
+            index = {"artifacts": [], "versions": []}
+            self._save_index(index)
+            return index
+        except Exception as e:
+            print(f"Failed to load index file: {e}")
+            return {"artifacts": [], "versions": []}
+
+    def save_index(self, index: Dict[str, Any]) -> None:
+        """
+        Save the current index to OSS
+        """
+        self._save_index(index)
+
+    def _save_index(self, index: Dict[str, Any]) -> None:
+        """
+        Save index file to OSS
+        Args:
+            index: Index dictionary
+        """
+        try:
+            content = json.dumps(index, indent=2, ensure_ascii=False, cls=CommonEncoder)
+            self.bucket.put_object(self.index_key, content.encode('utf-8'))
+        except Exception as e:
+            print(f"Failed to save index file: {e}")
+            raise
+
+    def artifact_path(self, artifact_id: str) -> str:
+        """
+        Get the OSS path for a given artifact
+        Args:
+            artifact_id: Artifact identifier
+        Returns:
+            OSS path string
+        """
+        return f"{self.prefix}artifact/{artifact_id}/index.json"
+
+    def attachment_path(self, artifact_id: str, filename: str) -> str:
+        """
+        Get the OSS path for an artifact attachment
+        Args:
+            artifact_id: Artifact identifier
+            filename: Attachment filename
+        Returns:
+            OSS path string
+        """
+        return f"{self.prefix}artifact/{artifact_id}/attachments/{filename}"
+
+    def store_artifact(self, artifact: Artifact) -> str:
+        """
+        Store artifact and its attachments to OSS
+        Args:
+            artifact: Artifact instance to be stored
+        Returns:
+            Version identifier (always 'success' for now)
+        """
+        import oss2
+
+        try:
+            # Prepare version record
+            version = {
+                "hash": artifact.artifact_id,
+                "timestamp": time.time(),
+                "metadata": artifact.metadata or {}
+            }
+            # Store artifact content
+            data = artifact.to_dict()
+            content_key = self.artifact_path(artifact.artifact_id)
+            content = json.dumps(data, indent=2, ensure_ascii=False, cls=CommonEncoder)
+            self.bucket.put_object(content_key, content.encode('utf-8'))
+            # Store attachments if any
+            if artifact.attachments:
+                for attachment in artifact.attachments:
+                    if isinstance(attachment, ArtifactAttachment):
+                        attachment_key = self.attachment_path(artifact.artifact_id, attachment.filename)
+                        self.bucket.put_object(attachment_key, attachment.content.encode('utf-8'))
+            # Update index
+            artifact_exists = False
+            for item in self.index["artifacts"]:
+                if item['artifact_id'] == artifact.artifact_id:
+                    item['version'] = version
+                    artifact_exists = True
+                    break
+            if not artifact_exists:
+                self.index["artifacts"].append({
+                    'artifact_id': artifact.artifact_id,
+                    'type': 'artifact',
+                    'version': version
+                })
+            self._save_index(self.index)
+            return "success"
+        except Exception as e:
+            print(f"Storage failed: {e}")
+            raise
+
+    def retrieve_latest_artifact(self, artifact_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve the latest version of artifact based on artifact ID
+        Args:
+            artifact_id: Artifact identifier
+        Returns:
+            Stored data as dict, or None if it doesn't exist
+        """
+        import oss2
+
+        try:
+            content_key = self.artifact_path(artifact_id)
+            try:
+                result = self.bucket.get_object(content_key)
+                content = result.read().decode('utf-8')
+                return json.loads(content)
+            except oss2.exceptions.NoSuchKey:
+                print(f"Content file doesn't exist: {content_key}")
+                return None
+        except Exception as e:
+            print(f"Failed to retrieve latest artifact: {e}")
+            return None
+
+    def get_artifact_versions(self, artifact_id: str) -> List[Dict[str, Any]]:
+        """
+        Get information about all versions of an artifact (currently only latest version is tracked)
+        Args:
+            artifact_id: Artifact identifier
+        Returns:
+            List of version information
+        """
+        try:
+            for artifact in self.index["artifacts"]:
+                if artifact['artifact_id'] == artifact_id:
+                    version_info = artifact["version"].copy()
+                    version_info["artifact_id"] = artifact_id
+                    return [version_info]
+            return []
+        except Exception as e:
+            print(f"Failed to get version information: {e}")
+            return []
+
+    def delete_artifact(self, artifact_id: str) -> bool:
+        """
+        Delete the specified artifact and its attachments from OSS
+        Args:
+            artifact_id: Artifact identifier
+        Returns:
+            Whether deletion was successful
+        """
+        import oss2
+
+        try:
+            # Delete artifact content
+            content_key = self.artifact_path(artifact_id)
+            try:
+                self.bucket.delete_object(content_key)
+            except oss2.exceptions.NoSuchKey:
+                pass
+            # Delete attachments (list objects under attachments/)
+            attachment_prefix = f"{self.prefix}artifact/{artifact_id}/attachments/"
+            for obj in oss2.ObjectIterator(self.bucket, prefix=attachment_prefix):
+                self.bucket.delete_object(obj.key)
+            # Remove from index
+            for i, artifact in enumerate(self.index["artifacts"]):
+                if artifact['artifact_id'] == artifact_id:
+                    del self.index["artifacts"][i]
+                    self._save_index(self.index)
+                    return True
+            return False
+        except Exception as e:
+            print(f"Failed to delete artifact: {e}")
+            return False
+
+    def list_artifacts(self) -> List[Dict[str, Any]]:
+        """
+        List all artifacts in the repository
+        Returns:
+            List of artifact information
+        """
+        try:
+            return [
+                {
+                    "artifact_id": artifact["artifact_id"],
+                    "type": artifact["type"],
+                    "timestamp": artifact["version"]["timestamp"],
+                    "metadata": artifact["version"]["metadata"]
+                }
+                for artifact in self.index["artifacts"]
+            ]
+        except Exception as e:
+            print(f"Failed to list artifacts: {e}")
+            return []
+    
+
+    def generate_tree_data(self, workspace_name: str) -> dict:
+        """
+        Generate a directory tree structure based on the OSS workspace folder structure.
+        Args:
+            workspace_name: Name of the workspace (for root node)
+        Returns:
+            Directory tree as dict
+        """
+        import oss2
+
+        all_keys = [obj.key for obj in oss2.ObjectIterator(self.bucket, prefix=self.prefix)]
+        # remove root prefix
+        rel_keys = [key[len(self.prefix):] for key in all_keys if key != self.index_key]
+        # build tree
+        root = {
+            "name": workspace_name,
+            "id": "-1",
+            "type": "dir",
+            "parentId": None,
+            "depth": 0,
+            "expanded": False,
+            "children": []
+        }
+        node_map = {"": root}  # 路径到节点的映射
+        for key in rel_keys:
+            parts = [p for p in key.split('/') if p]
+            cur_path = ""
+            parent_path = ""
+            for depth, part in enumerate(parts):
+                parent_path = cur_path
+                cur_path = f"{cur_path}/{part}" if cur_path else part
+                if cur_path not in node_map:
+                    node = {
+                        "name": part,
+                        "id": str(uuid.uuid4()),
+                        "type": "dir" if depth < len(parts) - 1 else "file",
+                        "parentId": node_map[parent_path]["id"],
+                        "depth": depth + 1,
+                        "expanded": False,
+                        "children": []
+                    }
+                    node_map[parent_path]["children"].append(node)
+                    node_map[cur_path] = node
+        return root