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

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

@@ -1,10 +1,13 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.1.dev0
+Version: 0.0.3
 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
 Provides-Extra: dev
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"

nora_lib-0.0.3/README.md

@@ -0,0 +1,37 @@
+# nora_lib
+
+For making and coordinating agents and tools.
+
+# Development
+
+When preparing a PR, make sure you first run the verify script
+to confirm the code is buildable, passes type- and formatting-checks,
+as well as unit tests.
+
+```bash
+cd <project_root>
+./verify.sh
+```
+
+The script will tell you what's wrong if anything fails.
+
+Update the `version=` field in `setup.py` when you make changes
+as part of your changeset.
+
+# Publication
+
+After your PR merges to `main` you will need to publish
+the library to public pypi for it to be useable by client applications.
+
+This requires one environment variable to be set, which can be found in
+the 1Pass NORA Vault under the secret named "NORA pypi token".
+
+```bash
+export AI2_NORA_PYPI_TOKEN=<SECRET IN NORA VAULT>
+cd <project_root>
+git checkout main
+git pull origin main
+git tag v<YOUR_NEW_VERSION>
+git push origin --tags
+./publish.sh
+```

nora_lib-0.0.3/nora_lib/context/context_service.py

@@ -0,0 +1,88 @@
+from datetime import datetime, timezone
+from typing import List, Optional
+from uuid import UUID
+
+from nora_lib.interactions.interactions_service import InteractionsService
+from nora_lib.interactions.models import (
+    ReturnedMessage,
+    ReturnedAgentContextMessage,
+    ReturnedAgentContextEvent,
+    EventType,
+    AgentMessageData,
+    Event,
+)
+from nora_lib.context.models import WrappedTaskObject
+
+
+class ContextService:
+    """
+    Save and retrieve task agent context from interaction store
+    """
+
+    def __init__(
+        self,
+        agent_actor_id: str,  # uuid representing this agent in interaction store
+        interactions_base_url: str,
+        interactions_bearer_token: Optional[str],
+        timeout: int = 30,
+    ):
+        self.interactions_service = self._get_interactions_service(
+            interactions_base_url, interactions_bearer_token, timeout
+        )
+        self.agent_actor_id = agent_actor_id
+
+    def _get_interactions_service(self, url, token, timeout) -> InteractionsService:
+        return InteractionsService(url, timeout, token)
+
+    def fetch_context(
+        self, request: WrappedTaskObject
+    ) -> List[ReturnedAgentContextMessage]:
+        message_id = request.message_id
+
+        returned_messages: List[ReturnedMessage] = (
+            self.interactions_service.fetch_messages_and_events_for_forked_thread(
+                message_id, EventType.AGENT_CONTEXT
+            )
+        )
+
+        messages_with_filtered_events: List[ReturnedAgentContextMessage] = []
+        for message in returned_messages:
+            events_saved_by_this_agent: List[ReturnedAgentContextEvent] = []
+            if message.events:
+                for event in message.events:
+                    context_event = ReturnedAgentContextEvent.model_validate(event)
+                    if context_event.actor_id == self.agent_actor_id:
+                        events_saved_by_this_agent.append(context_event)
+
+                events_saved_by_this_agent.sort(
+                    key=lambda event: datetime.fromisoformat(event.timestamp)
+                )
+
+            updated_message = ReturnedAgentContextMessage(
+                message_id=message.message_id,
+                actor_id=message.actor_id,
+                text=message.text,
+                ts=message.ts,
+                annotated_text=message.annotated_text,
+                events=events_saved_by_this_agent,
+            )
+
+            messages_with_filtered_events.append(updated_message)
+
+        return messages_with_filtered_events
+
+    def save_context(self, event_data: WrappedTaskObject):
+        agent_data = AgentMessageData(
+            message_data=event_data.model_dump(),
+            data_sender_actor_id=event_data.sender_actor_id,
+            virtual_thread_id=event_data.virtual_thread_id,
+        )
+        event = Event(
+            type=EventType.AGENT_CONTEXT,
+            actor_id=UUID(self.agent_actor_id),
+            timestamp=datetime.now(timezone.utc),
+            data=agent_data.model_dump(),
+            message_id=event_data.message_id,
+        )
+
+        self.interactions_service.save_event(event)

nora_lib-0.0.3/nora_lib/context/models.py

@@ -0,0 +1,21 @@
+from typing import Generic, Optional, TypeVar
+from pydantic import BaseModel, Field
+
+
+R = TypeVar("R", bound=BaseModel)
+
+
+class WrappedTaskObject(BaseModel, Generic[R]):
+    """Encloses request or response object with additional metadata"""
+
+    message_id: str = Field(
+        description="id of originating message; key for istore retrieval"
+    )
+    sender_actor_id: str = Field(
+        description="string representation of uuid identifying agent sending data"
+    )
+    virtual_thread_id: Optional[str] = Field(
+        description="Tool-defined local thread to associate follow up requests"
+    )
+    task_id: Optional[str] = Field(description="Reference to a long-running task")
+    data: R = Field(description="Tool-defined request or response")

nora_lib-0.0.3/nora_lib/interactions/interactions_service.py

@@ -0,0 +1,117 @@
+from datetime import datetime
+import requests
+from typing import List, Optional
+
+from nora_lib.interactions.models import (
+    Event,
+    EventType,
+    ReturnedMessage,
+    ThreadRelationsResponse,
+    ThreadForkEventData,
+    thread_message_lookup_request,
+)
+
+
+class InteractionsService:
+    """
+    Service which saves interactions to the Interactions API
+    """
+
+    def __init__(self, base_url, timeout, token):
+        self.base_url = base_url
+        self.timeout = timeout
+        self.headers = {"Authorization": f"Bearer {token}"}
+
+    def save_event(self, event: Event) -> None:
+        """Save an event to the Interactions API"""
+        event_url = f"{self.base_url}/interaction/v1/event"
+        response = requests.post(
+            event_url,
+            json=event.model_dump(),
+            headers=self.headers,
+            timeout=int(self.timeout),
+        )
+        response.raise_for_status()
+
+    def fetch_thread_messages_and_events_for_message(
+        self, message_id: str, event_type: str
+    ) -> ThreadRelationsResponse:
+        """Fetch messages and associated events from the same thread as provided messagev id"""
+        message_url = f"{self.base_url}/interaction/v1/search/message"
+        request_body = thread_message_lookup_request(message_id, event_type=event_type)
+        response = requests.post(
+            message_url,
+            json=request_body,
+            headers=self.headers,
+            timeout=int(self.timeout),
+        )
+        response.raise_for_status()
+        json_response = response.json()
+
+        return ThreadRelationsResponse.model_validate(
+            json_response.get("message", {}).get("thread", {})
+        )
+
+    def fetch_messages_and_events_for_thread(
+        self,
+        thread_id: str,
+        event_type: Optional[str] = None,
+        min_timestamp: Optional[str] = None,
+    ) -> dict:
+        """Fetch messages and events for the thread containing a given message from the Interactions API"""
+        THREAD_SEARCH_URL = f"{self.base_url}/interaction/v1/search/thread"
+        request_body = {
+            "id": thread_id,
+            "relations": {
+                "messages": (
+                    {"filter": {"min_timestamp": min_timestamp}}
+                    if min_timestamp
+                    else {}
+                ),
+                "events": {"filter": {"type": event_type}} if event_type else {},
+            },
+        }
+
+        response = requests.post(
+            THREAD_SEARCH_URL,
+            json=request_body,
+            headers=self.headers,
+            timeout=int(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: datetime.fromisoformat(x.ts))
+
+        return returned_messages

nora_lib-0.0.3/nora_lib/interactions/models.py

@@ -0,0 +1,114 @@
+"""
+Model for interactions to be sent to the interactions service.
+"""
+
+from datetime import datetime
+from enum import Enum
+from typing import Optional, List
+from uuid import UUID
+
+from pydantic import BaseModel, Field, field_serializer
+
+
+class EventType(str, Enum):
+    """Event types for the interactions service"""
+
+    AGENT_CONTEXT = "agent:message_context"
+    THREAD_FORK = "thread_fork"
+
+
+class Event(BaseModel):
+    """event object to be sent to the interactions service; requires association with a message, thread or channel id"""
+
+    type: str
+    actor_id: UUID = Field(
+        description="identifies actor writing the event to the interaction service"
+    )
+    timestamp: datetime
+    text: Optional[str] = None
+    data: Optional[dict] = Field(default_factory=dict)
+    message_id: Optional[str] = None
+    thread_id: Optional[str] = None
+    channel_id: Optional[str] = None
+
+    @field_serializer("actor_id")
+    def serialize_actor_id(self, actor_id: UUID):
+        return str(actor_id)
+
+    @field_serializer("timestamp")
+    def serialize_timestamp(self, timestamp: datetime):
+        return timestamp.isoformat()
+
+
+class AgentMessageData(BaseModel):
+    """capture requests to and responses from tools within Events"""
+
+    message_data: dict  # dict of agent/tool request/response format
+    data_sender_actor_id: Optional[str] = None  # agent sending the data
+    virtual_thread_id: Optional[str] = None  # tool-provided thread
+    tool_call_id: Optional[str] = None  # llm-provided thread
+    tool_name: Optional[str] = None  # llm identifier for tool
+
+
+class ReturnedAgentContextEvent(BaseModel):
+    """Event format returned by interaction service for agent context events"""
+
+    actor_id: str  # agent that saved this context
+    timestamp: str
+    data: AgentMessageData
+    type: str
+
+
+class ReturnedAgentContextMessage(BaseModel):
+    """Message format returned by interaction service for search by thread"""
+
+    message_id: str
+    actor_id: str
+    text: str
+    ts: str
+    annotated_text: Optional[str] = None
+    events: Optional[List[ReturnedAgentContextEvent]] = None
+
+
+class ThreadForkEventData(BaseModel):
+    """Event data for a thread fork event"""
+
+    previous_message_id: str
+
+
+class ReturnedMessage(BaseModel):
+    """Message format returned by interaction service for search by thread"""
+
+    message_id: str
+    actor_id: str
+    text: str
+    ts: str
+    annotated_text: Optional[str] = None
+    events: Optional[List[dict]] = None
+
+
+class ThreadRelationsResponse(BaseModel):
+    """Thread format returned by interaction service for thread relations in a search response"""
+
+    thread_id: str
+    events: Optional[List[Event]] = None  # events associated only with the thread
+    messages: Optional[List[ReturnedMessage]] = (
+        None  # includes events associated with each message
+    )
+
+
+def thread_message_lookup_request(message_id: str, event_type: str) -> dict:
+    """retrieve messages and events for the thread associated with a message"""
+    return {
+        "id": message_id,
+        "relations": {
+            "thread": {
+                "relations": {
+                    "messages": {
+                        "relations": {"events": {"filter": {"type": event_type}}},
+                        "apply_annotations_from_actors": ["*"],
+                    },
+                }
+            }
+        },
+    }

nora_lib-0.0.3/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.3/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.3}/nora_lib.egg-info/PKG-INFO

@@ -1,10 +1,13 @@
 Metadata-Version: 2.1
 Name: nora_lib
-Version: 0.0.1.dev0
+Version: 0.0.3
 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
 Provides-Extra: dev
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"

nora_lib-0.0.3/nora_lib.egg-info/SOURCES.txt

@@ -0,0 +1,22 @@
+README.md
+pyproject.toml
+setup.py
+nora_lib/__init__.py
+nora_lib/py.typed
+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/context/__init__.py
+nora_lib/context/context_service.py
+nora_lib/context/models.py
+nora_lib/interactions/__init__.py
+nora_lib/interactions/interactions_service.py
+nora_lib/interactions/models.py
+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.3/nora_lib.egg-info/requires.txt

@@ -0,0 +1,8 @@
+pydantic<3,>=2
+requests
+
+[dev]
+mypy
+pytest
+black
+types-requests

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

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

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

@@ -1,13 +1,13 @@
 import setuptools
 
-runtime_requirements = []
+runtime_requirements = ["pydantic>=2,<3", "requests"]
 
 # For running tests, linting, etc
-dev_requirements = ["mypy", "pytest", "black"]
+dev_requirements = ["mypy", "pytest", "black", "types-requests"]
 
 setuptools.setup(
     name="nora_lib",
-    version="0.0.1.dev0",
+    version="0.0.3",
     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.3/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.3/tests/test_placeholder.py

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

nora_lib-0.0.1.dev0/README.md

@@ -1,3 +0,0 @@
-# nora_lib
-
-For making and coordinating agents and tools.

nora_lib-0.0.1.dev0/nora_lib.egg-info/SOURCES.txt

@@ -1,10 +0,0 @@
-README.md
-pyproject.toml
-setup.py
-nora_lib/__init__.py
-nora_lib/py.typed
-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-0.0.1.dev0/nora_lib.egg-info/requires.txt

@@ -1,5 +0,0 @@
-
-[dev]
-mypy
-pytest
-black