nebius/SWE-bench-extra huggingface dataset:
Repository: UFRN-URNAI/urnai-tools

Issue:
Base class for the state representation
We need a class that lets us know the agent's current state. It needs to have the `update` method, which will update the agent's state according to the script developed.

Suggestions:

- Class name: `StateBase`.
- Class path: `urnai.states`.

Comments:

Patch:
diff --git a/urnai/states/__init__.py b/urnai/states/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/urnai/states/state_base.py b/urnai/states/state_base.py
new file mode 100644
index 0000000..8e0a903
--- /dev/null
+++ b/urnai/states/state_base.py
@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from typing import Any, List
+
+
+class StateBase(ABC):
+    """
+    Every Agent needs to own an instance of this base class 
+    in order to define its State.
+    So every time we want to create a new agent,
+    we should either use an existing State implementation or create a new one.
+    """
+
+    @abstractmethod
+    def update(self, obs) -> List[Any]:
+        """
+        This method receives as a parameter an Observation and returns a State, which 
+        is usually a list of features extracted from the Observation. The Agent uses
+        this State during training to receive a new action from its model and also to
+        make it learn, that's why this method should always return a list.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def state(self):
+        """Returns the State currently saved."""
+        pass
+
+    @property
+    @abstractmethod
+    def dimension(self):
+        """Returns the dimensions of the States returned by the update method."""
+        pass
+
+    @abstractmethod
+    def reset(self):
+        """Resets the State currently saved."""
+        pass
\ No newline at end of file


Test patch:
diff --git a/tests/units/states/__init__.py b/tests/units/states/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/units/states/test_state_base.py b/tests/units/states/test_state_base.py
new file mode 100644
index 0000000..c42b803
--- /dev/null
+++ b/tests/units/states/test_state_base.py
@@ -0,0 +1,25 @@
+import unittest
+from abc import ABCMeta
+
+from urnai.states.state_base import StateBase
+
+
+class TestStateBase(unittest.TestCase):
+
+    def test_abstract_methods(self):
+        StateBase.__abstractmethods__ = set()
+
+        class FakeState(StateBase):
+            def __init__(self):
+                super().__init__()
+
+        fake_state = FakeState()
+        update_return = fake_state.update("observation")
+        state = fake_state.state
+        dimension = fake_state.dimension
+        reset_return = fake_state.reset()
+        assert isinstance(StateBase, ABCMeta)
+        assert update_return is None
+        assert state is None
+        assert dimension is None
+        assert reset_return is None