nora-lib-impl 1.0.0__py3-none-any.whl

nora_lib/impl/tasks/state.py

@@ -0,0 +1,156 @@
+"""
+Handles IO for asynchronous task-related state.
+
+A StateManager instance must be initialized with
+a concrete subclass of `AsyncTaskState`, as implemented
+by dependent projects.
+"""
+
+import logging
+from uuid import UUID
+from datetime import datetime, timezone
+from typing import Optional, Any
+from pydantic import BaseModel
+
+from nora_lib.tasks.models import AsyncTaskState, R
+from nora_lib.tasks.state import (
+    IStateManager,
+    TaskStateFetchException,
+    NoSuchTaskException,
+)
+from nora_lib.impl.interactions.interactions_service import InteractionsService
+from nora_lib.impl.interactions.models import Event, ReturnedEvent
+from nora_lib.impl.pubsub import PubsubService
+from nora_lib.impl.context.agent_context import AgentContext
+
+TASK_STATE_CHANGE_TOPIC = "istore:event:task_state"
+
+
+class RemoteStateManagerFactory:
+    """
+    Stores task state in the interaction store
+    """
+
+    def __init__(
+        self,
+        agent_name: str,
+        actor_id: UUID,
+        interactions_service: InteractionsService,
+        pubsub_service: PubsubService,
+    ):
+        """
+        :param agent_name: Used to form the event type that will hold the task state in the interactions store
+        :param actor_id: Associated with the events written to the interactions store
+        :param interactions_service:
+        """
+        self.agent_name = agent_name
+        self.actor_id = actor_id
+        self.interactions_service = interactions_service
+        self.pubsub_service = pubsub_service
+
+    def for_message(self, message_id: str) -> IStateManager[R]:
+        return RemoteStateManager(
+            self.agent_name,
+            self.actor_id,
+            self.interactions_service,
+            self.pubsub_service,
+            message_id,
+        )
+
+    def for_agent_context(self, context: AgentContext) -> IStateManager[R]:
+        return RemoteStateManager(
+            self.agent_name,
+            self.actor_id,
+            self.interactions_service,
+            PubsubService(context.pubsub.base_url, context.pubsub.namespace),
+            context.message.message_id,
+        )
+
+
+class RemoteStateManager(IStateManager[R]):
+    """
+    Stores task state in the interaction store
+    """
+
+    _TASK_STATE_EVENT_TYPE = "agent:{}:task_state"
+
+    def __init__(
+        self,
+        agent_name: str,
+        actor_id: UUID,
+        interactions_service: InteractionsService,
+        pubsub_service: PubsubService,
+        message_id: str,
+    ):
+        """
+        :param agent_name: Agent that saved the task
+        :param actor_id: ID for the agent (ignored when reading)
+        :param message_id: The message that initiated the request for task status
+        """
+        self.agent_name = agent_name
+        self.actor_id = actor_id
+        self.message_id = message_id
+        self.interactions_service = interactions_service
+        self.pubsub_service = pubsub_service
+
+    def read_state(self, task_id: str) -> AsyncTaskState[R]:
+        event_type = RemoteStateManager._TASK_STATE_EVENT_TYPE.format(self.agent_name)
+        response = (
+            self.interactions_service.fetch_thread_messages_and_events_for_message(
+                self.message_id, [event_type]
+            )
+        )
+        latest_state: Optional[AsyncTaskState[R]] = None
+        latest_timestamp = None
+        for msg in response.messages or []:
+            for event in msg.events or []:
+                try:
+                    state = AsyncTaskState[Any].model_validate(event.data)
+                except Exception as e:
+                    # Event json blob has unexpected format
+                    raise TaskStateFetchException(
+                        f"Event of type {event_type} for message {self.message_id} does not deserialize to AsyncTaskState: {e}"
+                    )
+                if state.task_id != task_id:
+                    continue
+                if latest_state is None or (
+                    latest_timestamp and event.timestamp > latest_timestamp
+                ):
+                    latest_state = state
+                    latest_timestamp = event.timestamp
+
+        if not latest_state:
+            raise NoSuchTaskException(task_id)
+        return latest_state
+
+    def write_state(self, state: AsyncTaskState[R]) -> None:
+        event_type = RemoteStateManager._TASK_STATE_EVENT_TYPE.format(self.agent_name)
+        event = Event(
+            type=event_type,
+            actor_id=self.actor_id,
+            timestamp=datetime.now(tz=timezone.utc),
+            message_id=self.message_id,
+            data=state.model_dump(),
+        )
+        event_id = self.interactions_service.save_event(event)
+        returned_event = ReturnedEvent(
+            event_id=event_id,
+            type=event.type,
+            actor_id=event.actor_id,
+            message_id=event.message_id,
+            timestamp=event.timestamp,
+        )
+        payload = TaskStateChangeNotification(
+            agent=self.agent_name, event=returned_event
+        )
+        try:
+            self.pubsub_service.publish(TASK_STATE_CHANGE_TOPIC, payload.model_dump())
+        except Exception as e:
+            logging.exception(
+                f"Failed to publish event to pubsub topic {TASK_STATE_CHANGE_TOPIC} at {self.pubsub_service.base_url}"
+            )
+
+
+class TaskStateChangeNotification(BaseModel):
+    agent: str
+    event: ReturnedEvent

nora_lib/progress/models.py

@@ -0,0 +1,47 @@
+from enum import Enum
+import uuid
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from nora_lib.serializers import (
+    UuidWithSerializer,
+    DatetimeWithSerializer,
+)
+
+
+class RunState(str, Enum):
+    """State of a step"""
+
+    CREATED = "created"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+
+
+class StepProgress(BaseModel):
+    """Data class for step progress. This goes into the `data` field in `Event`."""
+
+    # A short message, e.g. "searching for $query". Recommend < 100 chars.
+    short_desc: str
+    # Detailed message.
+    long_desc: Optional[str] = None
+    # Updates on the same unit of work have the same step_id.
+    step_id: UuidWithSerializer = Field(default_factory=uuid.uuid4)
+    # Inner steps can be constituent to some outer step, effectively a tree.
+    parent_step_id: Optional[UuidWithSerializer] = None
+    # Populated if this step is due to an async task.
+    task_id: Optional[str] = None
+
+    # Enum of possible states.
+    run_state: RunState = RunState.CREATED
+    # DB timestamp when this step was defined/created.
+    created_at: Optional[DatetimeWithSerializer] = None
+    # When this step started running.
+    started_at: Optional[DatetimeWithSerializer] = None
+    # Estimated finish time, if available.
+    finish_est: Optional[DatetimeWithSerializer] = None
+    # When this step stopped running, whether that was due to success or failure.
+    finished_at: Optional[DatetimeWithSerializer] = None
+    # Error message in case of terminal step failure.
+    error_message: Optional[str] = None

nora_lib/progress/reporter.py

@@ -0,0 +1,122 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+from datetime import datetime, timezone
+
+from nora_lib.progress.models import StepProgress, RunState
+
+
+class StepProgressWriter(ABC):
+    """Writes step progress to a store or service"""
+
+    def write(self, step_progress: StepProgress):
+        """Default no-op implementation"""
+        pass
+
+
+class StepProgressReporter:
+    """
+    Tracks the lifecycle of a task and records incremental progress to some external store
+
+    Usage:
+    # Create/define a step
+    find_papers_progress = StepProgressReporter(...)
+    find_papers_progress.create()
+
+    # Do something
+    ...
+
+    # Start the step
+    find_papers_progress.start()
+
+    # Describe a part of the work of that outer step as a child step.
+    ...
+    count_citation_progress = find_papers_progress.create_child_step(short_desc="Count citations")
+    count_citation_progress.create()
+    ...
+    count_citation_progress.start()
+    ...
+    count_citation_progress.finish(is_success=True)
+    ...
+
+    # Finish the outer step
+    find_papers_progress.finish(is_success=False, error_message="Something went wrong")
+
+    # Alternatively, you can use this as a context. Step state transitions are managed for you,
+    # so you should NOT call any of the create/start/finish methods.
+
+    with StepProgressReporter(...) as spr:
+        # Do something
+        ...
+
+    # This step will be automatically created, started, and finished when the context exits.
+    # If an exception is raised, the step will be marked as failed
+    # and the exception message will be recorded in the error_message field
+    """
+
+    def __init__(self, step_progress: StepProgress, writer: StepProgressWriter):
+        self.step_progress = step_progress
+        self.writer = writer
+
+    def __enter__(self):
+        if self.step_progress.created_at is None:
+            self.create()
+        self.start()
+        return self
+
+    def __exit__(self, error_type, value, traceback):
+        is_success = error_type is None
+        self.finish(is_success=is_success, error_message=str(value))
+        return True
+
+    def create(self):
+        """Create a step, but don't start it yet. This is useful for defining plans."""
+        if self.step_progress.run_state in [
+            RunState.RUNNING,
+            RunState.SUCCEEDED,
+            RunState.FAILED,
+        ]:
+            return
+
+        self.step_progress.run_state = RunState.CREATED
+        self.step_progress.created_at = datetime.now(timezone.utc)
+        self.writer.write(self.step_progress)
+
+    def start(self):
+        """Start a step"""
+        if self.step_progress.run_state in [
+            RunState.RUNNING,
+            RunState.SUCCEEDED,
+            RunState.FAILED,
+        ]:
+            return
+
+        self.step_progress.started_at = datetime.now(timezone.utc)
+        self.step_progress.run_state = RunState.RUNNING
+        self.writer.write(self.step_progress)
+
+    def finish(self, is_success: bool, error_message: Optional[str] = None):
+        """Finish a step whether it was successful or not"""
+        if self.step_progress.run_state != RunState.RUNNING:
+            return
+        else:
+            self.step_progress.finished_at = datetime.now(timezone.utc)
+            self.step_progress.run_state = (
+                RunState.SUCCEEDED if is_success else RunState.FAILED
+            )
+            self.step_progress.error_message = error_message if error_message else None
+            self.writer.write(self.step_progress)
+
+    def create_child_step(
+        self, short_desc: str, **step_progress_kwargs
+    ) -> "StepProgressReporter":
+        """Create a child step"""
+        child_step_progress = StepProgress(
+            short_desc=short_desc,
+            parent_step_id=self.step_progress.step_id,
+            task_id=self.step_progress.task_id,
+            **step_progress_kwargs,
+        )
+        child_reporter = StepProgressReporter(
+            step_progress=child_step_progress, writer=self.writer
+        )
+        return child_reporter

nora_lib/serializers.py

@@ -0,0 +1,10 @@
+from datetime import datetime
+from uuid import UUID
+
+from pydantic import PlainSerializer
+from typing_extensions import Annotated
+
+UuidWithSerializer = Annotated[UUID, PlainSerializer(lambda x: str(x), return_type=str)]
+DatetimeWithSerializer = Annotated[
+    datetime, PlainSerializer(lambda x: x.isoformat(), return_type=str)
+]

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/tasks/state.py

@@ -0,0 +1,71 @@
+"""
+Handles IO for asynchronous task-related state.
+
+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, Type
+from abc import ABC, abstractmethod
+
+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 IStateManager(ABC, Generic[R]):
+    @abstractmethod
+    def read_state(self, task_id: str) -> AsyncTaskState[R]:
+        pass
+
+    @abstractmethod
+    def write_state(self, state: AsyncTaskState[R]) -> None:
+        pass
+
+    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)
+
+
+class StateManager(IStateManager[R]):
+    """
+    Stores task state on local disk
+    """
+
+    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)
+
+
+class TaskStateFetchException(Exception):
+    def __init__(self, message: str):
+        super().__init__(message)

nora_lib_impl-1.0.0.dist-info/METADATA

@@ -0,0 +1,20 @@
+Metadata-Version: 2.2
+Name: nora_lib-impl
+Version: 1.0.0
+Summary: For making and coordinating agents and tools
+Home-page: https://github.com/allenai/nora_lib/impl
+Requires-Python: >=3.9
+Requires-Dist: pydantic<3,>=2
+Requires-Dist: requests
+Requires-Dist: boto3
+Requires-Dist: aws_requests_auth
+Provides-Extra: dev
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"
+Dynamic: home-page
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

nora_lib_impl-1.0.0.dist-info/RECORD

@@ -0,0 +1,25 @@
+nora_lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/serializers.py,sha256=8l8pZWyVfVvtm2HjkvBwIrNRclFVmtuQZg16eD1LlNE,330
+nora_lib/impl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/impl/pubsub.py,sha256=0tmFfdOf4J-YZSxdpgXct0GpO9VatzqAfr7-5DB7s2Y,4779
+nora_lib/impl/context/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/impl/context/agent_context.py,sha256=I6e24JGpDZTJaxrawx5FnkOVnpr1c2D5YTqs8zweRyU,864
+nora_lib/impl/context/context_service.py,sha256=VfWBXDqGyR2mkv8sEVs4BFBtkmmjEM5oQSGUJFH2hi4,1189
+nora_lib/impl/context/models.py,sha256=QQ4LnAemHL0N4YK9vm8DQIqFyRzaQFuX4Ngn_XKzUF0,356
+nora_lib/impl/interactions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/impl/interactions/interactions_service.py,sha256=0EJTTcRQeEbmC0CW-2v_J61cG-UQM34nFJc1s0ITqOM,19660
+nora_lib/impl/interactions/models.py,sha256=rgZUft0yz2KpYVCkQizo9evN2cnsES3NPo8pzQEx2w4,12532
+nora_lib/impl/interactions/step_progress.py,sha256=ZlxD488P7gHZdOOQhsFvk_bNJMRIyW912IWY9dTHyD8,2374
+nora_lib/impl/tasks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/impl/tasks/state.py,sha256=WJnmzqAzhqLzPbx1O6R7j1bbsqhjT5hwPb0CaUfNS_E,5374
+nora_lib/progress/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/progress/models.py,sha256=YNwI_gFM7vzQBQaH6GIIQ28MXQDdL6m1_syFWvvMSWM,1581
+nora_lib/progress/reporter.py,sha256=JK4taghGTUt0zWpEWAznGaxnNtQf4yEwnjT_wrtYQBI,4045
+nora_lib/tasks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nora_lib/tasks/models.py,sha256=Guz4PC8DCrZ_JPOK6IvxHDNS1itTTwPmSMQHDsfbwdA,925
+nora_lib/tasks/state.py,sha256=blgv4xxz3GCxd7KN9avbtAJaSvJB3lH-Cbcxywnr-gE,2142
+nora_lib_impl-1.0.0.dist-info/METADATA,sha256=WPMniVqUGucqHt98vkRxGzQqr3hwdmXhTobMRtE9ZPI,579
+nora_lib_impl-1.0.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nora_lib_impl-1.0.0.dist-info/top_level.txt,sha256=BXasX5xI5NPFNqJsU63ojSi4Y3CPeH2ueaBR-WxxoaA,9
+nora_lib_impl-1.0.0.dist-info/RECORD,,

nora_lib_impl-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

nora_lib_impl-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+nora_lib