nora-lib 0.0.7.dev2__tar.gz → 0.1.8__tar.gz

{nora_lib-0.0.7.dev2 → nora_lib-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.7.dev2
+Version: 0.1.8
 Summary: For making and coordinating agents and tools
 Home-page: https://github.com/allenai/nora_lib
 Requires-Python: >=3.9

{nora_lib-0.0.7.dev2 → nora_lib-0.1.8}/nora_lib/interactions/interactions_service.py

@@ -1,5 +1,5 @@
 import requests
-from typing import Optional
+from typing import Optional, List
 import json
 
 from nora_lib.interactions.models import (
@@ -8,6 +8,7 @@ from nora_lib.interactions.models import (
     Message,
     ReturnedMessage,
     ThreadRelationsResponse,
+    VirtualThread,
 )
 
 
@@ -21,8 +22,13 @@ class InteractionsService:
         self.timeout = timeout
         self.headers = {"Authorization": f"Bearer {token}"} if token else None
 
-    def save_message(self, message: Message) -> None:
-        """Save a message to the Interactions API"""
+    def save_message(
+        self, message: Message, virtual_thread_id: Optional[str] = None
+    ) -> None:
+        """
+        Save a message to the Interaction Store
+        :param virtual_thread_id: Optional ID of a virtual thread to associate with the message
+        """
         message_url = f"{self.base_url}/interaction/v1/message"
         response = requests.post(
             message_url,
@@ -31,9 +37,25 @@ class InteractionsService:
             timeout=self.timeout,
         )
         response.raise_for_status()
+        if virtual_thread_id:
+            # Use an event to tag the message with the virtual thread ID
+            event = Event(
+                type=VirtualThread.EVENT_TYPE,
+                actor_id=message.actor_id,
+                message_id=message.message_id,
+                data={
+                    VirtualThread.ID_FIELD: virtual_thread_id,
+                    VirtualThread.EVENT_TYPE_FIELD: VirtualThread.EVENT_TYPE,
+                },
+                timestamp=message.ts,
+            )
+            self.save_event(event)
 
-    def save_event(self, event: Event) -> None:
-        """Save an event to the Interactions API"""
+    def save_event(self, event: Event, virtual_thread_id: Optional[str] = None) -> None:
+        """
+        Save an event to the Interaction Store
+        :param virtual_thread_id: Optional ID of a virtual thread to associate with the event
+        """
         event_url = f"{self.base_url}/interaction/v1/event"
         response = requests.post(
             event_url,
@@ -42,6 +64,77 @@ class InteractionsService:
             timeout=self.timeout,
         )
         response.raise_for_status()
+        if virtual_thread_id:
+            # Use an event to tag the event with the virtual thread ID
+            # Attach it to the same message as this event, along with the event type
+            event = Event(
+                type=VirtualThread.EVENT_TYPE,
+                actor_id=event.actor_id,
+                message_id=event.message_id,
+                data={
+                    VirtualThread.ID_FIELD: virtual_thread_id,
+                    VirtualThread.EVENT_TYPE_FIELD: event.type,
+                },
+                timestamp=event.timestamp,
+            )
+            self.save_event(event)
+
+    def get_virtual_thread_content(
+        self, message_id: str, virtual_thread_id: str
+    ) -> List[ReturnedMessage]:
+        """Fetch all messages and events in a virtual thread
+        Returns all messages and events in the same thread as the given message,
+        but filtered to only include those associated with the given virtual thread.
+        :param message_id: The ID of a message in the virtual thread
+        :param virtual_thread_id: The ID of the virtual thread
+        """
+        message_search_url = f"{self.base_url}/interaction/v1/search/message"
+        # Fetch all events and filter on the client side
+        # Need an IStore schema change to do this server-side
+        request_body = {
+            "id": message_id,
+            "relations": {
+                "preceding_messages": {
+                    "max": 100,
+                    "relations": {"events": {}},
+                },
+                "events": {},
+            },
+        }
+
+        response = requests.post(
+            message_search_url,
+            json=request_body,
+            headers=self.headers,
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        result = ReturnedMessage.model_validate(response.json()["message"])
+        all_messages = result.preceding_messages + [result]
+        virtual_thread_content = []
+        for msg in all_messages:
+            event_types_in_virtual_thread = set(
+                event.data[VirtualThread.EVENT_TYPE_FIELD]
+                for event in msg.events
+                if event.type == VirtualThread.EVENT_TYPE
+                and event.data.get(VirtualThread.ID_FIELD) == virtual_thread_id
+            )
+            if not event_types_in_virtual_thread:
+                continue
+            virtual_thread_content.append(msg)
+            msg.events = [
+                event
+                for event in msg.events
+                if event.type != VirtualThread.EVENT_TYPE
+                and event.type in event_types_in_virtual_thread
+            ]
+            if VirtualThread.EVENT_TYPE not in event_types_in_virtual_thread:
+                # An event has been tagged with the virtual thread ID
+                # but the message itself is not in the virtual thread
+                # Somewhat pathological case, probably shouldn't happen
+                # Set the message text to empty string
+                msg.text = ""
+        return virtual_thread_content
 
     def save_annotation(self, annotation: AnnotationBatch) -> None:
         """Save an annotation to the Interactions API"""
@@ -79,11 +172,18 @@ class InteractionsService:
 
         return res
 
-    def fetch_all_threads_by_channel(self, channel_id: str, min_timestamp: str) -> dict:
+    def fetch_all_threads_by_channel(
+        self,
+        channel_id: str,
+        min_timestamp: str,
+        thread_event_types: Optional[list[str]] = None,
+    ) -> dict:
         """Fetch a message from the Interactions API"""
         message_url = f"{self.base_url}/interaction/v1/search/channel"
         request_body = self._channel_lookup_request(
-            channel_id=channel_id, min_timestamp=min_timestamp
+            channel_id=channel_id,
+            min_timestamp=min_timestamp,
+            thread_event_types=thread_event_types,
         )
         response = requests.post(
             message_url,
@@ -95,7 +195,7 @@ class InteractionsService:
         return response.json()
 
     def fetch_thread_messages_and_events_for_message(
-        self, message_id: str, event_types: list[str]
+        self, message_id: str, event_types: List[str]
     ) -> ThreadRelationsResponse:
         """Fetch messages sorted by timestamp and events for agent context"""
         message_url = f"{self.base_url}/interaction/v1/search/message"
@@ -166,7 +266,9 @@ class InteractionsService:
         return response.json()
 
     @staticmethod
-    def _channel_lookup_request(channel_id: str, min_timestamp: str) -> dict:
+    def _channel_lookup_request(
+        channel_id: str, min_timestamp: str, thread_event_types: Optional[list[str]]
+    ) -> dict:
         """Interaction service API request to get threads and messages for a channel"""
         return {
             "id": channel_id,
@@ -176,7 +278,8 @@ class InteractionsService:
                         "messages": {
                             "filter": {"min_timestamp": min_timestamp},
                             "apply_annotations_from_actors": ["*"],
-                        }
+                        },
+                        "events": {"filter": {"type": thread_event_types or []}},
                     }
                 }
             },

{nora_lib-0.0.7.dev2 → nora_lib-0.1.8}/nora_lib/interactions/models.py

@@ -87,6 +87,7 @@ class ReturnedMessage(BaseModel):
     message_id: Optional[str] = None
     annotated_text: Optional[str] = None
     events: List[Event] = Field(default_factory=list)
+    preceding_messages: List["ReturnedMessage"] = Field(default_factory=list)
     thread_id: Optional[str] = None
     channel_id: Optional[str] = None
     annotations: List[Annotation] = Field(default_factory=list)
@@ -144,6 +145,19 @@ class ThreadRelationsResponse(BaseModel):
     )  # includes events associated with each message
 
 
+class VirtualThread:
+    """Virtuals threads are an event type used to sub-divide a thread into sb-conversations"""
+
+    # The type of event that represetns a virtual thread
+    EVENT_TYPE = "virtual_thread"
+
+    # Data field in the event that contains the ID of the virtual thread id
+    ID_FIELD = "virtual_thread_id"
+
+    # Data field in the event that contains the type of other events in the virtual thread
+    EVENT_TYPE_FIELD = "event_type"
+
+
 def thread_message_lookup_request(message_id: str, event_type: str) -> dict:
     """retrieve messages and events for the thread associated with a message"""
     return {

{nora_lib-0.0.7.dev2 → nora_lib-0.1.8}/nora_lib.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.7.dev2
+Version: 0.1.8
 Summary: For making and coordinating agents and tools
 Home-page: https://github.com/allenai/nora_lib
 Requires-Python: >=3.9

{nora_lib-0.0.7.dev2 → nora_lib-0.1.8}/nora_lib.egg-info/SOURCES.txt

@@ -18,5 +18,6 @@ nora_lib/tasks/__init__.py
 nora_lib/tasks/models.py
 nora_lib/tasks/state.py
 tests/test_placeholder.py
+tests/test_virtual_threads.py
 tests/tasks/__init__.py
 tests/tasks/test_state.py

{nora_lib-0.0.7.dev2 → nora_lib-0.1.8}/setup.py

@@ -7,7 +7,7 @@ dev_requirements = ["mypy", "pytest", "black", "types-requests"]
 
 setuptools.setup(
     name="nora_lib",
-    version="0.0.7.dev2",
+    version="0.1.8",
     description="For making and coordinating agents and tools",
     url="https://github.com/allenai/nora_lib",
     packages=setuptools.find_packages(exclude=(["tests"])),

nora_lib-0.1.8/tests/test_virtual_threads.py

@@ -0,0 +1,58 @@
+import unittest
+
+from nora_lib.interactions.interactions_service import InteractionsService
+from nora_lib.interactions.models import *
+from uuid import uuid4
+
+ACTOR = uuid4()
+THREAD = str(uuid4())
+CHANNEL = str(uuid4())
+
+
+def _msg(text):
+    return Message(
+        message_id=str(uuid4()),
+        actor_id=ACTOR,
+        text=text,
+        channel_id=CHANNEL,
+        thread_id=THREAD,
+        surface=Surface.WEB,
+        ts=datetime.now(),
+    )
+
+
+def _event(msg: Message, type: str, data: dict):
+    return Event(
+        type=type,
+        actor_id=ACTOR,
+        timestamp=datetime.now(),
+        text="",
+        data=data,
+        message_id=msg.message_id,
+    )
+
+
+@unittest.skip("Requires a local instance of the interactions service")
+class TestVirtualThreads(unittest.TestCase):
+    def setUp(self):
+        self.svc = InteractionsService("http://localhost:8090", 30, None)
+
+    def test_placeholder(self):
+        virtual_thread_1 = "virtual_thread_1"
+        virtual_thread_2 = "virtual_thread_2"
+        msg1 = _msg("Hi 1")
+        msg2 = _msg("Hi 2")
+        self.svc.save_message(msg1)
+        self.svc.save_message(msg2, virtual_thread_1)
+        event1 = _event(msg2, "event1", {})
+        event2 = _event(msg2, "event2", {})
+        event3 = _event(msg2, "event3", {})
+        self.svc.save_event(event1)
+        self.svc.save_event(event2, virtual_thread_1)
+        self.svc.save_event(event3, virtual_thread_2)
+
+        content = self.svc.get_virtual_thread_content(msg2.message_id, virtual_thread_1)
+        # Should only contain the one message tagged with virtual_thread_1
+        self.assertEqual([m.message_id for m in content], [msg2.message_id])
+        # Should only contain the events tagged with virtual_thread_1
+        self.assertEqual([e.type for e in content[0].events], [event2.type])