feat: add download_examples command (#87)

* fix: examples cannot be used while using pip
Fixes #52

* docs: update README

* fix: revert deletion

* test: add unit tests for MemOS CLI functions

* Update Makefile

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: Copilot <[email protected]>

Author: Ki-Seki

Makefile

@@ -24,4 +24,4 @@ serve:
 	poetry run uvicorn memos.api.start_api:app
 
 openapi:
-	poetry run python scripts/export_openapi.py --output docs/openapi.json
+	poetry run memos export_openapi --output docs/openapi.json

README.md

@@ -174,6 +174,14 @@ curl -fsSL https://ollama.com/install.sh | sh
 
 To use functionalities based on the `transformers` library, ensure you have [PyTorch](https://pytorch.org/get-started/locally/) installed (CUDA version recommended for GPU acceleration).
 
+#### Download Examples
+
+To download example code, data and configurations, run the following command:
+
+```bash
+memos download_examples
+```
+
 ## 💬 Community & Support
 
 Join our community to ask questions, share your projects, and connect with other developers.

pyproject.toml

@@ -59,6 +59,9 @@ python-dotenv = "^1.1.1"
 langgraph = "^0.5.1"
 langmem = "^0.0.27"
 
+[tool.poetry.scripts]
+memos = "memos.cli:main"
+
 [[tool.poetry.source]]
 name = "mirrors"
 url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/"

scripts/export_openapi.py

@@ -0,0 +1,113 @@
+"""
+MemOS CLI Tool
+This script provides command-line interface for MemOS operations.
+"""
+
+import argparse
+import json
+import os
+import zipfile
+
+from io import BytesIO
+
+
+def export_openapi(output: str) -> bool:
+    """Export OpenAPI schema to JSON file."""
+    from memos.api.start_api import app
+
+    # Create directory if it doesn't exist
+    if os.path.dirname(output):
+        os.makedirs(os.path.dirname(output), exist_ok=True)
+
+    with open(output, "w") as f:
+        json.dump(app.openapi(), f, indent=2)
+        f.write("\n")
+
+    print(f"✅ OpenAPI schema exported to: {output}")
+    return True
+
+
+def download_examples(dest: str) -> bool:
+    import requests
+
+    """Download examples from the MemOS repository."""
+    zip_url = "https://github.com/MemTensor/MemOS/archive/refs/heads/main.zip"
+    print(f"📥 Downloading examples from {zip_url}...")
+
+    try:
+        response = requests.get(zip_url)
+        response.raise_for_status()
+
+        with zipfile.ZipFile(BytesIO(response.content)) as z:
+            extracted_files = []
+            for file in z.namelist():
+                if "MemOS-main/examples/" in file and not file.endswith("/"):
+                    # Remove the prefix and extract to dest
+                    relative_path = file.replace("MemOS-main/examples/", "")
+                    extract_path = os.path.join(dest, relative_path)
+
+                    # Create directory if it doesn't exist
+                    os.makedirs(os.path.dirname(extract_path), exist_ok=True)
+
+                    # Extract the file
+                    with z.open(file) as source, open(extract_path, "wb") as target:
+                        target.write(source.read())
+                    extracted_files.append(extract_path)
+
+        print(f"✅ Examples downloaded to: {dest}")
+        print(f"📁 {len(extracted_files)} files extracted")
+
+    except requests.RequestException as e:
+        print(f"❌ Error downloading examples: {e}")
+        return False
+    except Exception as e:
+        print(f"❌ Error extracting examples: {e}")
+        return False
+
+    return True
+
+
+def main():
+    """Main CLI entry point."""
+    parser = argparse.ArgumentParser(
+        prog="memos",
+        description="MemOS Command Line Interface",
+    )
+
+    # Create subparsers for different commands
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # Download examples command
+    examples_parser = subparsers.add_parser("download_examples", help="Download example files")
+    examples_parser.add_argument(
+        "--dest",
+        type=str,
+        default="./examples",
+        help="Destination directory for examples (default: ./examples)",
+    )
+
+    # Export API command
+    api_parser = subparsers.add_parser("export_openapi", help="Export OpenAPI schema to JSON file")
+    api_parser.add_argument(
+        "--output",
+        type=str,
+        default="openapi.json",
+        help="Output path for OpenAPI schema (default: openapi.json)",
+    )
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Handle commands
+    if args.command == "download_examples":
+        success = download_examples(args.dest)
+        exit(0 if success else 1)
+    elif args.command == "export_openapi":
+        success = export_openapi(args.output)
+        exit(0 if success else 1)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

src/memos/cli.py

@@ -0,0 +1,106 @@
+"""
+Tests for the MemOS CLI tool.
+"""
+
+import zipfile
+
+from io import BytesIO
+from unittest.mock import MagicMock, mock_open, patch
+
+import pytest
+import requests
+
+from memos.cli import download_examples, export_openapi, main
+
+
+class TestExportOpenAPI:
+    """Test the export_openapi function."""
+
+    @patch("memos.api.start_api.app")
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("os.makedirs")
+    def test_export_openapi_success(self, mock_makedirs, mock_file, mock_app):
+        """Test successful OpenAPI export."""
+        mock_openapi_data = {"openapi": "3.0.0", "info": {"title": "Test API"}}
+        mock_app.openapi.return_value = mock_openapi_data
+
+        result = export_openapi("/test/path/openapi.json")
+
+        assert result is True
+        mock_makedirs.assert_called_once_with("/test/path", exist_ok=True)
+        mock_file.assert_called_once_with("/test/path/openapi.json", "w")
+
+    @patch("memos.api.start_api.app")
+    @patch("builtins.open", side_effect=OSError("Permission denied"))
+    def test_export_openapi_error(self, mock_file, mock_app):
+        """Test OpenAPI export when file writing fails."""
+        mock_app.openapi.return_value = {"test": "data"}
+
+        with pytest.raises(IOError):
+            export_openapi("/invalid/path/openapi.json")
+
+
+class TestDownloadExamples:
+    """Test the download_examples function."""
+
+    def create_mock_zip_content(self):
+        """Create mock zip file content for testing."""
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
+            zip_file.writestr("MemOS-main/examples/test_example.py", "# Test example content")
+            zip_file.writestr(
+                "MemOS-main/examples/subfolder/another_example.py", "# Another example"
+            )
+        return zip_buffer.getvalue()
+
+    @patch("requests.get")
+    @patch("os.makedirs")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_download_examples_success(self, mock_file, mock_makedirs, mock_requests):
+        """Test successful examples download."""
+        mock_response = MagicMock()
+        mock_response.content = self.create_mock_zip_content()
+        mock_requests.return_value = mock_response
+
+        result = download_examples("/test/dest")
+
+        assert result is True
+        mock_requests.assert_called_once_with(
+            "https://github.com/MemTensor/MemOS/archive/refs/heads/main.zip"
+        )
+        mock_response.raise_for_status.assert_called_once()
+
+    @patch("requests.get")
+    def test_download_examples_error(self, mock_requests):
+        """Test download examples when request fails."""
+        mock_requests.side_effect = requests.RequestException("Network error")
+
+        result = download_examples("/test/dest")
+
+        assert result is False
+
+
+class TestMainCLI:
+    """Test the main CLI function."""
+
+    @patch("memos.cli.download_examples")
+    def test_main_download_examples(self, mock_download):
+        """Test main function with download_examples command."""
+        mock_download.return_value = True
+
+        with patch("sys.argv", ["memos", "download_examples", "--dest", "/test/dest"]):
+            with pytest.raises(SystemExit) as exc_info:
+                main()
+            assert exc_info.value.code == 0
+            mock_download.assert_called_once_with("/test/dest")
+
+    @patch("memos.cli.export_openapi")
+    def test_main_export_openapi(self, mock_export):
+        """Test main function with export_openapi command."""
+        mock_export.return_value = True
+
+        with patch("sys.argv", ["memos", "export_openapi", "--output", "/test/openapi.json"]):
+            with pytest.raises(SystemExit) as exc_info:
+                main()
+            assert exc_info.value.code == 0
+            mock_export.assert_called_once_with("/test/openapi.json")