nora-lib 0.0.5.dev2__tar.gz → 0.0.6__tar.gz

{nora_lib-0.0.5.dev2 → nora_lib-0.0.6}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.5.dev2
+Version: 0.0.6
 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
 Requires-Dist: requests
+Requires-Dist: boto3
 Provides-Extra: dev
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pytest; extra == "dev"

{nora_lib-0.0.5.dev2 → nora_lib-0.0.6}/nora_lib/interactions/interactions_service.py

@@ -2,6 +2,7 @@ from datetime import datetime
 import logging
 import requests
 from typing import List, Optional
+import json
 
 from nora_lib.interactions.models import (
     AnnotationBatch,
@@ -31,7 +32,7 @@ class InteractionsService:
             message_url,
             json=message.model_dump(),
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
 
@@ -42,7 +43,7 @@ class InteractionsService:
             event_url,
             json=event.model_dump(),
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
 
@@ -53,7 +54,7 @@ class InteractionsService:
             annotation_url,
             json=annotation.model_dump(),
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
 
@@ -68,7 +69,7 @@ class InteractionsService:
             message_url,
             json=request_body,
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
         res_dict = response.json()["message"]
@@ -92,7 +93,7 @@ class InteractionsService:
             message_url,
             json=request_body,
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
         return response.json()
@@ -130,7 +131,7 @@ class InteractionsService:
                 e,
             )
 
-        messages_with_events.sort(key=lambda x: datetime.fromisoformat(x.ts))
+        messages_with_events.sort(key=lambda x: x.ts)
         return messages_with_events
 
     def fetch_thread_messages_and_events_for_message(
@@ -143,7 +144,7 @@ class InteractionsService:
             message_url,
             json=request_body,
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
         json_response = response.json()
@@ -176,11 +177,46 @@ class InteractionsService:
             thread_search_url,
             json=request_body,
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
         return response.json()
 
+    def fetch_messages_and_events_for_forked_thread(
+        self, message_id: str, event_type: str
+    ) -> List[ReturnedMessage]:
+        """Build a history of messages for a given message including associated events.
+        This includes messages from pre-forked threads."""
+        returned_messages: List[ReturnedMessage] = []
+
+        messages_for_thread: ThreadRelationsResponse = (
+            self.fetch_thread_messages_and_events_for_message(message_id, [event_type])
+        )
+        if messages_for_thread.messages:
+            returned_messages.extend(messages_for_thread.messages)
+
+        # Lookup any thread_fork events (conversation across surfaces)
+        thread_fork_events = self.fetch_messages_and_events_for_thread(
+            messages_for_thread.thread_id, EventType.THREAD_FORK.value
+        )
+        for forked_thread_event in thread_fork_events.get("thread", {}).get(
+            "events", []
+        ):
+            event_data = ThreadForkEventData.model_validate(
+                forked_thread_event.get("data", {})
+            )
+            forked_thread: ThreadRelationsResponse = (
+                self.fetch_thread_messages_and_events_for_message(
+                    event_data.previous_message_id, [event_type]
+                )
+            )
+            if forked_thread.messages:
+                returned_messages.extend(forked_thread.messages)
+
+        returned_messages.sort(key=lambda x: x.ts)
+
+        return returned_messages
+
     def fetch_events_for_message(
         self,
         message_id: str,
@@ -199,7 +235,7 @@ class InteractionsService:
             message_search_url,
             json=request_body,
             headers=self.headers,
-            timeout=int(self.timeout),
+            timeout=self.timeout,
         )
         response.raise_for_status()
         return response.json()
@@ -237,3 +273,32 @@ class InteractionsService:
                 }
             },
         }
+
+    @staticmethod
+    def prod() -> "InteractionsService":
+        return InteractionsService(
+            base_url="https://s2ub.prod.s2.allenai.org/service/noraretrieval",
+            timeout=30,
+            token=InteractionsService._fetch_bearer_token(
+                "nora/prod/interaction-bearer-token"
+            ),
+        )
+
+    @staticmethod
+    def dev() -> "InteractionsService":
+        return InteractionsService(
+            base_url="https://s2ub.dev.s2.allenai.org/service/noraretrieval",
+            timeout=30,
+            token=InteractionsService._fetch_bearer_token(
+                "nora/dev/interaction-bearer-token"
+            ),
+        )
+
+    @staticmethod
+    def _fetch_bearer_token(secret_id: str) -> str:
+        import boto3
+
+        secrets_manager = boto3.client("secretsmanager", region_name="us-west-2")
+        return json.loads(
+            secrets_manager.get_secret_value(SecretId=secret_id)["SecretString"]
+        )["token"]

{nora_lib-0.0.5.dev2 → nora_lib-0.0.6}/nora_lib/interactions/models.py

@@ -43,7 +43,7 @@ class Message(BaseModel):
     channel_id: str
     surface: Surface
     ts: datetime
-    annotations: Optional[List[Annotation]] = None
+    annotations: List[Annotation] = Field(default_factory=list)
 
     @field_serializer("actor_id")
     def serialize_actor_id(self, actor_id: UUID):
@@ -58,7 +58,9 @@ class EventType(str, Enum):
     """Event types for the interactions service"""
 
     AGENT_CONTEXT = "agent:message_context"
+    S2_ANNOTATION = "s2_annotation"
     THREAD_FORK = "thread_fork"
+    UI_INTERACTION = "ui_interaction"
 
 
 class Event(BaseModel):
@@ -88,14 +90,14 @@ class ReturnedMessage(BaseModel):
     """Message format returned by interaction service"""
 
     message_id: str
-    actor_id: str
+    actor_id: UUID
     text: str
-    ts: str
+    ts: datetime
     annotated_text: Optional[str] = None
     events: List[Event] = Field(default_factory=list)
     thread_id: Optional[str] = None
     channel_id: Optional[str] = None
-    annotations: Optional[List[Annotation]] = None
+    annotations: List[Annotation] = Field(default_factory=list)
 
 
 class AgentMessageData(BaseModel):
@@ -111,8 +113,8 @@ class AgentMessageData(BaseModel):
 class ReturnedAgentContextEvent(BaseModel):
     """Event format returned by interaction service for agent context events"""
 
-    actor_id: str  # agent that saved this context
-    timestamp: str
+    actor_id: UUID  # agent that saved this context
+    timestamp: datetime
     data: AgentMessageData
     type: str
 
@@ -121,11 +123,11 @@ class ReturnedAgentContextMessage(BaseModel):
     """Message format returned by interaction service for search by thread"""
 
     message_id: str
-    actor_id: str
+    actor_id: UUID
     text: str
     ts: str
     annotated_text: Optional[str] = None
-    events: Optional[List[ReturnedAgentContextEvent]] = None
+    events: List[ReturnedAgentContextEvent] = Field(default_factory=list)
 
 
 class ThreadForkEventData(BaseModel):

nora_lib-0.0.6/nora_lib/tasks/state.py

@@ -0,0 +1,163 @@
+"""
+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 uuid import UUID
+from datetime import datetime, timezone
+from typing import Generic, Optional, Type, Any
+from abc import ABC, abstractmethod
+
+from nora_lib.tasks.models import AsyncTaskState, R, TASK_STATUSES
+from nora_lib.interactions.interactions_service import InteractionsService
+from nora_lib.interactions.models import Event
+
+
+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)
+
+
+class RemoteStateManagerFactory:
+    """
+    Stores task state in the interaction store
+    """
+
+    def __init__(
+        self, agent_name: str, actor_id: UUID, interactions_service: InteractionsService
+    ):
+        """
+        :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
+
+    def for_message(self, message_id: str) -> IStateManager[R]:
+        return RemoteStateManager(
+            self.agent_name, self.actor_id, self.interactions_service, 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,
+        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
+
+    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 = Event(
+            type=RemoteStateManager._TASK_STATE_EVENT_TYPE.format(self.agent_name),
+            actor_id=self.actor_id,
+            timestamp=datetime.now(tz=timezone.utc),
+            message_id=self.message_id,
+            data=state.model_dump(),
+        )
+        self.interactions_service.save_event(event)

{nora_lib-0.0.5.dev2 → nora_lib-0.0.6}/nora_lib.egg-info/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.5.dev2
+Version: 0.0.6
 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
 Requires-Dist: requests
+Requires-Dist: boto3
 Provides-Extra: dev
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pytest; extra == "dev"

{nora_lib-0.0.5.dev2 → nora_lib-0.0.6}/nora_lib.egg-info/requires.txt

@@ -1,5 +1,6 @@
 pydantic<3,>=2
 requests
+boto3
 
 [dev]
 mypy

{nora_lib-0.0.5.dev2 → nora_lib-0.0.6}/setup.py

@@ -1,13 +1,13 @@
 import setuptools
 
-runtime_requirements = ["pydantic>=2,<3", "requests"]
+runtime_requirements = ["pydantic>=2,<3", "requests", "boto3"]
 
 # For running tests, linting, etc
 dev_requirements = ["mypy", "pytest", "black", "types-requests"]
 
 setuptools.setup(
     name="nora_lib",
-    version="0.0.5.dev2",
+    version="0.0.6",
     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.5.dev2/nora_lib/tasks/state.py

@@ -1,53 +0,0 @@
-"""
-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)