From 52bd5b13a7cae5e73e250bd3c339d5595006771a Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Thu, 9 Jan 2025 08:03:30 -0800 Subject: [PATCH] Add store auth types --- libs/sdk-py/langgraph_sdk/auth/__init__.py | 69 +++++++++ libs/sdk-py/langgraph_sdk/auth/types.py | 157 ++++++++++++++++++++- libs/sdk-py/pyproject.toml | 2 +- 3 files changed, 222 insertions(+), 6 deletions(-) diff --git a/libs/sdk-py/langgraph_sdk/auth/__init__.py b/libs/sdk-py/langgraph_sdk/auth/__init__.py index 317c89e9f..e2ab57416 100644 --- a/libs/sdk-py/langgraph_sdk/auth/__init__.py +++ b/libs/sdk-py/langgraph_sdk/auth/__init__.py @@ -461,6 +461,73 @@ class _CronsOn( Search = types.CronsSearch +class _StoreOn: + def __init__(self, auth: Auth) -> None: + self._auth = auth + + @typing.overload + def __call__( + self, + *, + actions: typing.Optional[ + typing.Union[ + typing.Literal["put", "get", "search", "list_namespaces", "delete"], + Sequence[ + typing.Literal["put", "get", "search", "list_namespaces", "delete"] + ], + ] + ] = None, + ) -> Callable[[AHO], AHO]: ... + + @typing.overload + def __call__(self, fn: AHO) -> AHO: ... + + def __call__( + self, + fn: typing.Optional[AHO] = None, + *, + actions: typing.Optional[ + typing.Union[ + typing.Literal["put", "get", "search", "list_namespaces", "delete"], + Sequence[ + typing.Literal["put", "get", "search", "list_namespaces", "delete"] + ], + ] + ] = None, + ) -> typing.Union[AHO, Callable[[AHO], AHO]]: + """Register a handler for specific resources and actions. + + Can be used as a decorator or with explicit resource/action parameters: + + @auth.on.store + async def handler(): ... # Handle all store ops + + @auth.on.store(actions=("put", "get", "search", "delete")) + async def handler(): ... # Handle specific store ops + + @auth.on.store.put + async def handler(): ... # Handle store.put ops + """ + if fn is not None: + # Used as a plain decorator + _register_handler(self._auth, None, None, fn) + return fn + + # Used with parameters, return a decorator + def decorator( + handler: AHO, + ) -> AHO: + if isinstance(actions, str): + action_list = [actions] + else: + action_list = list(actions) if actions is not None else ["*"] + for action in action_list: + _register_handler(self._auth, "store", action, handler) + return handler + + return decorator + + AHO = typing.TypeVar("AHO", bound=_ActionHandler[dict[str, typing.Any]]) @@ -524,6 +591,7 @@ async def rate_limit_writes(ctx: AuthContext, value: Any) -> bool: "threads", "runs", "crons", + "store", "value", ) @@ -532,6 +600,7 @@ def __init__(self, auth: Auth) -> None: self.assistants = _AssistantsOn(auth, "assistants") self.threads = _ThreadsOn(auth, "threads") self.crons = _CronsOn(auth, "crons") + self.store = _StoreOn(auth) self.value = dict[str, typing.Any] @typing.overload diff --git a/libs/sdk-py/langgraph_sdk/auth/types.py b/libs/sdk-py/langgraph_sdk/auth/types.py index b057aeb84..4a4777339 100644 --- a/libs/sdk-py/langgraph_sdk/auth/types.py +++ b/libs/sdk-py/langgraph_sdk/auth/types.py @@ -5,7 +5,7 @@ and typed dictionaries for various API operations. Note: - All typing.TypedDict classes use total=False to make all fields optional by default. + All typing.TypedDict classes use total=False to make all fields typing.Optional by default. """ import functools @@ -157,7 +157,7 @@ class MinimalUserDict(typing.TypedDict, total=False): identity: typing_extensions.Required[str] """The required unique identifier for the user.""" display_name: str - """The optional display name for the user.""" + """The typing.Optional display name for the user.""" is_authenticated: bool """Whether the user is authenticated. Defaults to True.""" permissions: Sequence[str] @@ -358,11 +358,34 @@ class AuthContext(BaseAuthContext): allowing for fine-grained access control decisions. """ - resource: typing.Literal["runs", "threads", "crons", "assistants"] + resource: typing.Literal["runs", "threads", "crons", "assistants", "store"] """The resource being accessed.""" - action: typing.Literal["create", "read", "update", "delete", "search", "create_run"] - """The action being performed on the resource.""" + action: typing.Literal[ + "create", + "read", + "update", + "delete", + "search", + "create_run", + "put", + "get", + "list_namespaces", + ] + """The action being performed on the resource. + + Most resources support the following actions: + - create: Create a new resource + - read: Read information about a resource + - update: Update an existing resource + - delete: Delete a resource + - search: Search for resources + + The store supports the following actions: + - put: Add or update a document in the store + - get: Get a document from the store + - list_namespaces: List the namespaces in the store + """ class ThreadsCreate(typing.TypedDict, total=False): @@ -759,6 +782,93 @@ class CronsSearch(typing.TypedDict, total=False): """Offset for pagination.""" +class StoreGet(typing.TypedDict): + """Operation to retrieve a specific item by its namespace and key.""" + + namespace: tuple[str, ...] + """Hierarchical path that uniquely identifies the item's location.""" + + key: str + """Unique identifier for the item within its specific namespace.""" + + +class StoreSearch(typing.TypedDict): + """Operation to search for items within a specified namespace hierarchy.""" + + namespace_prefix: tuple[str, ...] + """Hierarchical path prefix defining the search scope. + + ???+ example "Examples" + + ```python + () # Search entire store + ("documents",) # Search all documents + ("users", "content") # Search within user content + ``` + """ + + filter: typing.Optional[dict[str, typing.Any]] + """Key-value pairs for filtering results based on exact matches or comparison operators.""" + + limit: int + """Maximum number of items to return in the search results.""" + + offset: int + """Number of matching items to skip for pagination.""" + + query: typing.Optional[str] + """Naturalj language search query for semantic search capabilities.""" + + +class StoreListNamespaces(typing.TypedDict): + """Operation to list and filter namespaces in the store.""" + + prefix: typing.Optional[tuple[str, ...]] + """Optional conditions for filtering namespaces.""" + + suffix: typing.Optional[tuple[str, ...]] + """Optional conditions for filtering namespaces.""" + + max_depth: typing.Optional[int] + """Maximum depth of namespace hierarchy to return. + + Note: + Namespaces deeper than this level will be truncated. + """ + + limit: int + """Maximum number of namespaces to return.""" + + offset: int + """Number of namespaces to skip for pagination.""" + + +class StorePut(typing.TypedDict): + """Operation to store, update, or delete an item in the store.""" + + namespace: tuple[str, ...] + """Hierarchical path that identifies the location of the item.""" + + key: str + """Unique identifier for the item within its namespace.""" + + value: typing.Optional[dict[str, typing.Any]] + """The data to store, or None to mark the item for deletion.""" + + index: typing.Optional[typing.Union[typing.Literal[False], list[str]]] + """Optional index configuration for full-text search.""" + + +class StoreDelete(typing.TypedDict): + """Operation to delete an item from the store.""" + + namespace: tuple[str, ...] + """Hierarchical path that uniquely identifies the item's location.""" + + key: str + """Unique identifier for the item within its specific namespace.""" + + class on: """Namespace for type definitions of different API operations. @@ -894,6 +1004,38 @@ class search: value = CronsSearch + class store: + """Types for store-related operations.""" + + value = typing.Union[ + StoreGet, StoreSearch, StoreListNamespaces, StorePut, StoreDelete + ] + + class put: + """Type for store put parameters.""" + + value = StorePut + + class get: + """Type for store get parameters.""" + + value = StoreGet + + class search: + """Type for store search parameters.""" + + value = StoreSearch + + class delete: + """Type for store delete parameters.""" + + value = StoreDelete + + class list_namespaces: + """Type for store list namespaces parameters.""" + + value = StoreListNamespaces + __all__ = [ "on", @@ -909,4 +1051,9 @@ class search: "AssistantsUpdate", "AssistantsDelete", "AssistantsSearch", + "StoreGet", + "StoreSearch", + "StoreListNamespaces", + "StorePut", + "StoreDelete", ] diff --git a/libs/sdk-py/pyproject.toml b/libs/sdk-py/pyproject.toml index 5613ded15..b60851127 100644 --- a/libs/sdk-py/pyproject.toml +++ b/libs/sdk-py/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-sdk" -version = "0.1.49" +version = "0.1.50" description = "SDK for interacting with LangGraph API" authors = [] license = "MIT"