diff --git a/Makefile b/Makefile index c250316ca..fcffaf7d1 100644 --- a/Makefile +++ b/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 diff --git a/README.md b/README.md index 00af88745..273d1603b 100644 --- a/README.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index ad82c63f0..2b5bed7b0 100644 --- a/pyproject.toml +++ b/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/" diff --git a/scripts/export_openapi.py b/scripts/export_openapi.py deleted file mode 100644 index 818b16466..000000000 --- a/scripts/export_openapi.py +++ /dev/null @@ -1,16 +0,0 @@ -import argparse -import json - -from memos.api.start_api import app - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Export OpenAPI schema to JSON file.") - parser.add_argument( - "--output", type=str, default="docs/openapi.json", help="Output path for OpenAPI schema." - ) - args = parser.parse_args() - with open(args.output, "w") as f: - json.dump(app.openapi(), f, indent=2) - f.write("\n") - print("Export completed successfully") diff --git a/src/memos/cli.py b/src/memos/cli.py new file mode 100644 index 000000000..fb3475ff3 --- /dev/null +++ b/src/memos/cli.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() diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 000000000..9750af121 --- /dev/null +++ b/tests/test_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")