diff --git a/moatless/events.py b/moatless/events.py index acf1cfc5..6b885181 100644 --- a/moatless/events.py +++ b/moatless/events.py @@ -2,7 +2,12 @@ from datetime import datetime, timezone from typing import Optional -from pydantic import BaseModel, Field, field_validator +try: + from pydantic import BaseModel, Field, field_validator + PYDANTIC_V2 = True +except ImportError: # pragma: no cover - support pydantic v1 + from pydantic import BaseModel, Field, validator as field_validator + PYDANTIC_V2 = False logger = logging.getLogger(__name__) @@ -19,12 +24,19 @@ class BaseEvent(BaseModel): model_config = {"ser_json_timedelta": "iso8601", "json_encoders": {datetime: lambda dt: dt.isoformat()}} - @field_validator("timestamp", mode="before") - @classmethod - def parse_datetime(cls, value): - if isinstance(value, str): - return datetime.fromisoformat(value) - return value + if PYDANTIC_V2: + @field_validator("timestamp", mode="before") + @classmethod + def parse_datetime(cls, value): + if isinstance(value, str): + return datetime.fromisoformat(value) + return value + else: # pragma: no cover - support pydantic v1 + @field_validator("timestamp", pre=True) + def parse_datetime(cls, value): + if isinstance(value, str): + return datetime.fromisoformat(value) + return value @classmethod def from_dict(cls, data: dict) -> "BaseEvent": diff --git a/moatless/settings.py b/moatless/settings.py index 80e6492e..12ae39c5 100644 --- a/moatless/settings.py +++ b/moatless/settings.py @@ -1,9 +1,14 @@ import logging import os +from typing import Type -from dotenv import load_dotenv +try: + from dotenv import load_dotenv +except ModuleNotFoundError: # pragma: no cover - optional dependency + def load_dotenv(*args, **kwargs): + """Fallback no-op if python-dotenv is missing.""" + return False -from litellm import Type from moatless.eventbus.base import BaseEventBus from moatless.eventbus.local_bus import LocalEventBus from moatless.runner.asyncio_runner import AsyncioRunner diff --git a/moatless/storage/__init__.py b/moatless/storage/__init__.py index c83883f0..340e9547 100644 --- a/moatless/storage/__init__.py +++ b/moatless/storage/__init__.py @@ -3,7 +3,12 @@ """ from moatless.storage.base import BaseStorage -from moatless.storage.file_storage import FileStorage + +try: + from moatless.storage.file_storage import FileStorage +except Exception: # pragma: no cover - optional dependency + FileStorage = None # type: ignore + from moatless.storage.memory_storage import MemoryStorage __all__ = ["BaseStorage", "FileStorage", "MemoryStorage"] diff --git a/moatless/storage/memory_storage.py b/moatless/storage/memory_storage.py index 18c96fdd..49af6b75 100644 --- a/moatless/storage/memory_storage.py +++ b/moatless/storage/memory_storage.py @@ -5,64 +5,113 @@ in memory, which is useful for testing or temporary storage. """ +import json import logging +from typing import Union -from moatless.storage.base import BaseStorage +from moatless.storage.base import BaseStorage, DateTimeEncoder logger = logging.getLogger(__name__) class MemoryStorage(BaseStorage): - """ - Storage implementation that uses in-memory dictionaries. + """Simple in-memory implementation of :class:`BaseStorage`.""" - This class provides a storage implementation that keeps all data - in memory, which is useful for testing or temporary storage. - """ - - def __init__(self): + def __init__(self) -> None: """Initialize an empty in-memory storage.""" - self._data: dict[str, dict] = {} + self._data: dict[str, object] = {} + + async def read_raw(self, path: str) -> str: + """Return the raw value stored under *path*.""" + normalized_key = self.normalize_path(path) + if normalized_key not in self._data: + raise KeyError(f"Key '{path}' does not exist") - async def read(self, path: str) -> dict: - """Read binary data from memory.""" + value = self._data[normalized_key] + if isinstance(value, list): + # Represent lists as JSONL + return "\n".join( + json.dumps(v, cls=DateTimeEncoder) if isinstance(v, dict) else str(v) + for v in value + ) + if isinstance(value, dict): + return json.dumps(value, cls=DateTimeEncoder) + return str(value) + + async def read_lines(self, path: str) -> list[dict]: + """Return a list of objects stored under *path*.""" normalized_key = self.normalize_path(path) if normalized_key not in self._data: raise KeyError(f"Key '{path}' does not exist") - return self._data[normalized_key] - async def write(self, key: str, data: dict) -> None: - """Write binary data to memory.""" - normalized_key = self.normalize_path(key) + value = self._data[normalized_key] + results: list[dict] = [] + if isinstance(value, list): + for item in value: + if isinstance(item, str): + item = item.strip() + if item: + results.append(json.loads(item)) + else: + results.append(item) + elif isinstance(value, str): + for line in value.splitlines(): + line = line.strip() + if line: + results.append(json.loads(line)) + elif isinstance(value, dict): + results.append(value) + return results + + async def write_raw(self, path: str, data: str) -> None: + """Write raw string *data* to *path*.""" + normalized_key = self.normalize_path(path) self._data[normalized_key] = data - async def delete(self, key: str) -> None: - """Delete data from memory.""" - normalized_key = self.normalize_path(key) + async def append(self, path: str, data: Union[dict, str]) -> None: + """Append *data* to the entry at *path*.""" + normalized_key = self.normalize_path(path) + existing = self._data.get(normalized_key) + + if existing is None: + self._data[normalized_key] = [] + existing = self._data[normalized_key] + + if not isinstance(existing, list): + if isinstance(existing, str): + lines = existing.splitlines() + self._data[normalized_key] = lines + else: + self._data[normalized_key] = [existing] + existing = self._data[normalized_key] + + assert isinstance(existing, list) + if isinstance(data, dict): + existing.append(data) + else: + existing.append(data.rstrip("\n")) + + async def delete(self, path: str) -> None: + """Delete the value at *path*.""" + normalized_key = self.normalize_path(path) if normalized_key not in self._data: - raise KeyError(f"Key '{key}' does not exist") + raise KeyError(f"Key '{path}' does not exist") del self._data[normalized_key] - async def exists(self, key: str) -> bool: - """Check if a key exists in memory.""" - normalized_key = self.normalize_path(key) + async def exists(self, path: str) -> bool: + """Return ``True`` if *path* exists.""" + normalized_key = self.normalize_path(path) return normalized_key in self._data async def list_paths(self, prefix: str = "") -> list[str]: - """ - List all keys with the given prefix. - - Args: - prefix: The key prefix to search for - - Returns: - A list of keys that match the prefix - """ + """List all keys starting with *prefix*.""" normalized_prefix = self.normalize_path(prefix) - # When prefix is empty, return all keys if not normalized_prefix: return list(self._data.keys()) - # Filter keys that start with the prefix - return [key for key in self._data.keys() if key == normalized_prefix or key.startswith(normalized_prefix + "/")] + return [ + key + for key in self._data.keys() + if key == normalized_prefix or key.startswith(normalized_prefix + "/") + ] diff --git a/tests/conftest.py b/tests/conftest.py index 9467dbb7..57538758 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,12 +4,18 @@ import tempfile import pytest -from dotenv import load_dotenv -# Import here to avoid circular imports -from moatless import settings -from moatless.storage.base import BaseStorage -from moatless.storage.file_storage import FileStorage +try: + from dotenv import load_dotenv +except ModuleNotFoundError: # pragma: no cover - optional dependency + def load_dotenv(*args, **kwargs): + """Fallback no-op if python-dotenv is not installed.""" + return False + +try: + from moatless.storage.file_storage import FileStorage +except Exception: # pragma: no cover - optional dependency + FileStorage = None # type: ignore load_dotenv() logger = logging.getLogger(__name__) @@ -35,8 +41,13 @@ def temp_dir(): @pytest.fixture def file_storage(temp_dir): """Fixture to create a storage instance with a test directory.""" + from moatless import settings + if FileStorage is not None: + settings._storage = FileStorage(base_dir=temp_dir) + else: + from moatless.storage.memory_storage import MemoryStorage - settings._storage = FileStorage(base_dir=temp_dir) + settings._storage = MemoryStorage() return settings._storage diff --git a/tests/storage/test_memory_storage.py b/tests/storage/test_memory_storage.py new file mode 100644 index 00000000..78cfb4f0 --- /dev/null +++ b/tests/storage/test_memory_storage.py @@ -0,0 +1,47 @@ +import asyncio + +from moatless.storage.memory_storage import MemoryStorage + + +def test_basic_operations(): + async def _run(): + storage = MemoryStorage() + path = "foo/bar.json" + data = {"hello": "world"} + + await storage.write(path, data) + assert await storage.exists(path) + + read = await storage.read(path) + assert read == data + + asyncio.run(_run()) + + +def test_append_and_read_lines(): + async def _run(): + storage = MemoryStorage() + path = "events.jsonl" + + await storage.append(path, {"id": 1}) + await storage.append(path, {"id": 2}) + + lines = await storage.read_lines(path) + assert lines == [{"id": 1}, {"id": 2}] + + asyncio.run(_run()) + + +def test_list_paths_and_delete(): + async def _run(): + storage = MemoryStorage() + await storage.write("a/one.json", {"v": 1}) + await storage.write("a/two.json", {"v": 2}) + + paths = await storage.list_paths("a") + assert sorted(paths) == ["a/one.json", "a/two.json"] + + await storage.delete("a/one.json") + assert not await storage.exists("a/one.json") + + asyncio.run(_run())