Repository: run-llama/llama_index
Pull request: Add get/set API to the Context and make it coroutine-safe
# Description

- Introduce `get` and `set` APIs and make them coroutine-safe
- Hide the actual storage data structures, so that the context can be implemented differently without breaking changes
- Introduce the concept of "private storaging" for data that shouldn't be shared across different steps

Fixes #15276

## Type of Change

Please delete options that are not relevant.

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] This change requires a documentation update

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [ ] Added new unit/integration tests
- [ ] Added new notebook (that tests end-to-end)
- [ ] I stared at the code and made sure it makes sense

## Suggested Checklist:

- [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added Google Colab support for the newly added notebooks.
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Feature Request]: Better API for `Context`
### Feature Description

Implement a get/set API for the `Context` instance instead of accessing the underlying `data` dictionary directly, e.g. `context.get("foo")` instead of `context.data["foo"]`

### Reason

Using a vanilla dict for context data is limiting and not very intuitive. This change would also allow for more flexible data persistence strategies, as the context data wouldn't be confined to in-memory storage in the current process.

### Value of Feature

_No response_

Hints from issue discussion:
We do already have a kvstore abstraction. I wonder if that could be looped in a in at some point

Commit message: docs

Patch to review:
diff --git a/llama-index-core/llama_index/core/workflow/context.py b/llama-index-core/llama_index/core/workflow/context.py
index 865572d568..e097214aec 100644
--- a/llama-index-core/llama_index/core/workflow/context.py
+++ b/llama-index-core/llama_index/core/workflow/context.py
@@ -1,21 +1,94 @@
 from collections import defaultdict
+import asyncio
 from typing import Dict, Any, Optional, List, Type
 
 from .events import Event
 
 
 class Context:
+    """A global object representing a context for a given workflow run.
+
+    The Context object can be used to store data that needs to be available across iterations during a workflow
+    execution, and across multiple workflow runs.
+    Every context instance offers two type of data storage: a global one, that's shared among all the steps within a
+    workflow, and private one, that's only accessible from a single step.
+
+    Both `set` and `get` operations on global data are governed by a lock, and considered coroutine-safe.
+    """
+
     def __init__(self, parent: Optional["Context"] = None) -> None:
-        # Global state
+        # Global data storage
         if parent:
-            self.data = parent.data
+            self._globals = parent._globals
         else:
-            self.data: Dict[str, Any] = {}
+            self._globals: Dict[str, Any] = {}
+            self._lock = asyncio.Lock()
+
+        # Local data storage
+        self._locals: Dict[str, Any] = {}
 
         # Step-specific instance
-        self.parent = parent
+        self._parent: Optional[Context] = parent
         self._events_buffer: Dict[Type[Event], List[Event]] = defaultdict(list)
 
+    async def set(self, key: str, value: Any, make_private: bool = False) -> None:
+        """Store `value` into the Context under `key`.
+
+        Args:
+            key: A unique string to identify the value stored.
+            value: The data to be stored.
+            make_private: Make the value only accessible from the step that stored it.
+
+        Raises:
+            ValueError: When make_private is True but a key already exists in the global storage.
+        """
+        if make_private:
+            if key in self._globals:
+                msg = f"A key named '{key}' already exists in the Context storage."
+                raise ValueError(msg)
+            self._locals[key] = value
+            return
+
+        await self.lock.acquire()
+        self._globals[key] = value
+        self.lock.release()
+
+    async def get(self, key: str, default: Optional[Any] = None) -> Any:
+        """Get the value corresponding to `key` from the Context.
+
+        Args:
+            key: A unique string to identify the value stored.
+            default: The value to return when `key` is missing instead of raising an exception.
+
+        Raises:
+            ValueError: When there's not value accessible corresponding to `key`.
+        """
+        if key in self._locals:
+            return self._locals[key]
+        elif key in self._globals:
+            await self.lock.acquire()
+            ret = self._globals[key]
+            self.lock.release()
+            return ret
+        elif default is not None:
+            return default
+
+        msg = f"Key '{key}' not found in Context"
+        raise ValueError(msg)
+
+    @property
+    def data(self):
+        """This property is provided for backward compatibility.
+
+        Use `globals` instead.
+        """
+        return self._globals
+
+    @property
+    def lock(self) -> asyncio.Lock:
+        """Returns a mutex to lock the Context."""
+        return self._parent._lock if self._parent else self._lock
+
     def collect_events(
         self, ev: Event, expected: List[Type[Event]]
     ) -> Optional[List[Event]]:


Code review comments:
diff_hunk: @@ -1,21 +1,94 @@
 from collections import defaultdict
+import asyncio
 from typing import Dict, Any, Optional, List, Type
 
 from .events import Event
 
 
 class Context:
+    """A global object representing a context for a given workflow run.
+
+    The Context object can be used to store data that needs to be available across iterations during a workflow
+    execution, and across multiple workflow runs.
+    Every context instance offers two type of data storage: a global one, that's shared among all the steps within a
+    workflow, and private one, that's only accessible from a single step.
+
+    Both `set` and `get` operations on global data are governed by a lock, and considered coroutine-safe.
+    """
+
     def __init__(self, parent: Optional["Context"] = None) -> None:
-        # Global state
+        # Global data storage
         if parent:
-            self.data = parent.data
+            self._globals = parent._globals
         else:
-            self.data: Dict[str, Any] = {}
+            self._globals: Dict[str, Any] = {}
+            self._lock = asyncio.Lock()
+
+        # Local data storage
+        self._locals: Dict[str, Any] = {}
 
         # Step-specific instance
-        self.parent = parent
+        self._parent: Optional[Context] = parent
         self._events_buffer: Dict[Type[Event], List[Event]] = defaultdict(list)
 
+    async def set(self, key: str, value: Any, make_private: bool = False) -> None:
+        """Store `value` into the Context under `key`.
+
+        Args:
+            key: A unique string to identify the value stored.
+            value: The data to be stored.
+            make_private: Make the value only accessible from the step that stored it.
+
+        Raises:
+            ValueError: When make_private is True but a key already exists in the global storage.
+        """
+        if make_private:
+            if key in self._globals:
+                msg = f"A key named '{key}' already exists in the Context storage."
+                raise ValueError(msg)
+            self._locals[key] = value
+            return
+
+        await self.lock.acquire()
+        self._globals[key] = value
+        self.lock.release()
+
+    async def get(self, key: str, default: Optional[Any] = None) -> Any:
+        """Get the value corresponding to `key` from the Context.
+
+        Args:
+            key: A unique string to identify the value stored.
+            default: The value to return when `key` is missing instead of raising an exception.
+
+        Raises:
+            ValueError: When there's not value accessible corresponding to `key`.
+        """
+        if key in self._locals:
+            return self._locals[key]
+        elif key in self._globals:
+            await self.lock.acquire()
path: llama-index-core/llama_index/core/workflow/context.py
text: @user1:
It might be slightly safer to do `with self.lock` instead?