Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 19 additions & 7 deletions moatless/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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":
Expand Down
9 changes: 7 additions & 2 deletions moatless/settings.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
7 changes: 6 additions & 1 deletion moatless/storage/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
117 changes: 83 additions & 34 deletions moatless/storage/memory_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 + "/")
]
23 changes: 17 additions & 6 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand All @@ -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


Expand Down
47 changes: 47 additions & 0 deletions tests/storage/test_memory_storage.py
Original file line number Diff line number Diff line change
@@ -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())