nora-lib 0.0.1.dev0__tar.gz → 0.0.2__tar.gz

{nora_lib-0.0.1.dev0 → nora_lib-0.0.2}/PKG-INFO

@@ -1,9 +1,10 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.1.dev0
+Version: 0.0.2
 Summary: For making and coordinating agents and tools
 Home-page: https://github.com/allenai/nora_lib
 Requires-Python: >=3.9
+Requires-Dist: pydantic<3,>=2
 Provides-Extra: dev
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pytest; extra == "dev"

nora_lib-0.0.2/nora_lib/tasks/models.py

@@ -0,0 +1,25 @@
+from typing import Any, Dict, Generic, Optional, TypeVar
+from pydantic import BaseModel, Field
+
+
+R = TypeVar("R", bound=BaseModel)
+
+
+TASK_STATUSES = dict(STARTED="STARTED", FAILED="FAILED", COMPLETED="COMPLETED")
+
+
+class AsyncTaskState(BaseModel, Generic[R]):
+    """Models the current state of an asynchronous request."""
+
+    task_id: str = Field(
+        "Identifies the long-running task so that its status and eventual result"
+        "can be checked in follow-up calls."
+    )
+    estimated_time: str = Field(
+        description="How long we expect this task to take from start to finish."
+    )
+    task_status: str = Field(description="Current human-readable status of the task.")
+    task_result: Optional[R] = Field(description="Final result of the task.")
+    extra_state: Dict[str, Any] = Field(
+        description="Any extra task-specific state can go in here as free-form JSON-serializable dictionary."
+    )

nora_lib-0.0.2/nora_lib/tasks/state.py

@@ -0,0 +1,53 @@
+"""
+Handles IO for asynchronous task-related state.
+Currently just reads/writes from local disk,
+the best and most robust mechanism.
+
+A StateManager instance must be initialized with
+a concrete subclass of `AsyncTaskState`, as implemented
+by dependent projects.
+"""
+
+import json
+import os
+from typing import Generic, Optional, Type
+
+from nora_lib.tasks.models import AsyncTaskState, R, TASK_STATUSES
+
+
+class NoSuchTaskException(Exception):
+    def __init__(self, task_id: str):
+        self._task_id = task_id
+
+    def __str__(self):
+        return f"No record found for task {self._task_id}"
+
+
+class StateManager(Generic[R]):
+    def __init__(self, task_state_class: Type[AsyncTaskState[R]], state_dir) -> None:
+        self._task_state_class = task_state_class
+        self._state_dir = state_dir
+
+    def read_state(self, task_id: str) -> AsyncTaskState[R]:
+        task_state_path = os.path.join(self._state_dir, f"{task_id}.json")
+        if not os.path.isfile(task_state_path):
+            raise NoSuchTaskException(task_id)
+
+        with open(task_state_path, "r") as f:
+            return self._task_state_class(**json.loads(f.read()))
+
+    def write_state(self, state: AsyncTaskState[R]) -> None:
+        task_state_path = os.path.join(self._state_dir, f"{state.task_id}.json")
+        with open(task_state_path, "w") as f:
+            json.dump(state.model_dump(), f)
+
+    def update_status(self, task_id: str, new_status: str) -> None:
+        state = self.read_state(task_id)
+        state.task_status = new_status
+        self.write_state(state)
+
+    def save_result(self, task_id: str, task_result: R) -> None:
+        state = self.read_state(task_id)
+        state.task_status = TASK_STATUSES["COMPLETED"]
+        state.task_result = task_result
+        self.write_state(state)

{nora_lib-0.0.1.dev0 → nora_lib-0.0.2}/nora_lib.egg-info/PKG-INFO

@@ -1,9 +1,10 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.1.dev0
+Version: 0.0.2
 Summary: For making and coordinating agents and tools
 Home-page: https://github.com/allenai/nora_lib
 Requires-Python: >=3.9
+Requires-Dist: pydantic<3,>=2
 Provides-Extra: dev
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pytest; extra == "dev"

{nora_lib-0.0.1.dev0 → nora_lib-0.0.2}/nora_lib.egg-info/SOURCES.txt

@@ -7,4 +7,10 @@ nora_lib.egg-info/PKG-INFO
 nora_lib.egg-info/SOURCES.txt
 nora_lib.egg-info/dependency_links.txt
 nora_lib.egg-info/requires.txt
-nora_lib.egg-info/top_level.txt
+nora_lib.egg-info/top_level.txt
+nora_lib/tasks/__init__.py
+nora_lib/tasks/models.py
+nora_lib/tasks/state.py
+tests/test_placeholder.py
+tests/tasks/__init__.py
+tests/tasks/test_state.py

{nora_lib-0.0.1.dev0 → nora_lib-0.0.2}/nora_lib.egg-info/requires.txt

@@ -1,3 +1,4 @@
+pydantic<3,>=2
 
 [dev]
 mypy

{nora_lib-0.0.1.dev0 → nora_lib-0.0.2}/nora_lib.egg-info/top_level.txt

@@ -1 +1,2 @@
 nora_lib
+tests

{nora_lib-0.0.1.dev0 → nora_lib-0.0.2}/setup.py

@@ -1,13 +1,13 @@
 import setuptools
 
-runtime_requirements = []
+runtime_requirements = ["pydantic>=2,<3"]
 
 # For running tests, linting, etc
 dev_requirements = ["mypy", "pytest", "black"]
 
 setuptools.setup(
     name="nora_lib",
-    version="0.0.1.dev0",
+    version="0.0.2",
     description="For making and coordinating agents and tools",
     url="https://github.com/allenai/nora_lib",
     packages=setuptools.find_packages(exclude=(["tests"])),

nora_lib-0.0.2/tests/tasks/test_state.py

@@ -0,0 +1,77 @@
+import tempfile
+import unittest
+
+from pydantic import BaseModel, Field
+
+from nora_lib.tasks.models import AsyncTaskState, TASK_STATUSES
+from nora_lib.tasks.state import NoSuchTaskException, StateManager
+
+
+class MyTaskResult(BaseModel):
+    a: str
+    b: int
+
+
+class MyAsyncTaskState(AsyncTaskState[MyTaskResult]):
+    pass
+
+
+class TestState(unittest.TestCase):
+    def test__can_read_and_write_state(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(MyAsyncTaskState, tmpdir)
+            state = MyAsyncTaskState(
+                task_id="asdf",
+                estimated_time="40 days and 40 nights",
+                task_status="STARTED",
+                task_result=None,
+                extra_state={"foo": "bar"},
+            )
+            manager.write_state(state)
+            fetched_state = manager.read_state("asdf")
+
+            self.assertEqual(state, fetched_state)
+
+    def test__raises_an_error_if_referencing_nonexistent_task(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(MyAsyncTaskState, tmpdir)
+            with self.assertRaises(NoSuchTaskException):
+                manager.read_state("asdf")
+
+    def test__allows_specific_update_of_status_field(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(MyAsyncTaskState, tmpdir)
+            state = MyAsyncTaskState(
+                task_id="asdf",
+                estimated_time="40 days and 40 nights",
+                task_status="STARTED",
+                task_result=None,
+                extra_state={"foo": "bar"},
+            )
+            manager.write_state(state)
+            manager.update_status(state.task_id, "fail")
+            fetched_state = manager.read_state("asdf")
+            state.task_status = "fail"
+
+            self.assertEqual(state, fetched_state)
+
+    def test__allows_specific_update_of_result_field(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(MyAsyncTaskState, tmpdir)
+            state = MyAsyncTaskState(
+                task_id="asdf",
+                estimated_time="40 days and 40 nights",
+                task_status="STARTED",
+                task_result=None,
+                extra_state={"foo": "bar"},
+            )
+            manager.write_state(state)
+
+            result = MyTaskResult(a="asdf", b=123)
+            manager.save_result(state.task_id, result)
+            fetched_state = manager.read_state("asdf")
+
+            state.task_status = TASK_STATUSES["COMPLETED"]
+            state.task_result = result
+
+            self.assertEqual(state, fetched_state)

nora_lib-0.0.2/tests/test_placeholder.py

@@ -0,0 +1,6 @@
+import unittest
+
+
+class TestPlaceholder(unittest.TestCase):
+    def test_placeholder(self):
+        self.assertEqual(1, 1)