Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add store auth types #2971

Merged
merged 2 commits into from
Jan 9, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 69 additions & 0 deletions libs/sdk-py/langgraph_sdk/auth/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from __future__ import annotations

Check notice on line 1 in libs/sdk-py/langgraph_sdk/auth/__init__.py

View workflow job for this annotation

GitHub Actions / benchmark

Benchmark results

......................................... fanout_to_subgraph_10x: Mean +- std dev: 63.3 ms +- 1.6 ms ......................................... fanout_to_subgraph_10x_sync: Mean +- std dev: 55.0 ms +- 1.0 ms ......................................... fanout_to_subgraph_10x_checkpoint: Mean +- std dev: 76.5 ms +- 1.7 ms ......................................... fanout_to_subgraph_10x_checkpoint_sync: Mean +- std dev: 96.8 ms +- 2.0 ms ......................................... fanout_to_subgraph_100x: Mean +- std dev: 662 ms +- 39 ms ......................................... fanout_to_subgraph_100x_sync: Mean +- std dev: 535 ms +- 19 ms ......................................... fanout_to_subgraph_100x_checkpoint: Mean +- std dev: 774 ms +- 26 ms ......................................... fanout_to_subgraph_100x_checkpoint_sync: Mean +- std dev: 968 ms +- 20 ms ......................................... react_agent_10x: Mean +- std dev: 31.1 ms +- 0.7 ms ......................................... react_agent_10x_sync: Mean +- std dev: 23.0 ms +- 0.2 ms ......................................... react_agent_10x_checkpoint: Mean +- std dev: 38.3 ms +- 0.8 ms ......................................... react_agent_10x_checkpoint_sync: Mean +- std dev: 36.7 ms +- 0.7 ms ......................................... react_agent_100x: Mean +- std dev: 343 ms +- 7 ms ......................................... react_agent_100x_sync: Mean +- std dev: 272 ms +- 3 ms ......................................... react_agent_100x_checkpoint: Mean +- std dev: 638 ms +- 10 ms ......................................... react_agent_100x_checkpoint_sync: Mean +- std dev: 619 ms +- 10 ms ......................................... wide_state_25x300: Mean +- std dev: 23.5 ms +- 0.6 ms ......................................... wide_state_25x300_sync: Mean +- std dev: 15.4 ms +- 0.2 ms ......................................... wide_state_25x300_checkpoint: Mean +- std dev: 250 ms +- 14 ms ......................................... wide_state_25x300_checkpoint_sync: Mean +- std dev: 245 ms +- 13 ms ......................................... wide_state_15x600: Mean +- std dev: 27.5 ms +- 0.6 ms ......................................... wide_state_15x600_sync: Mean +- std dev: 17.8 ms +- 0.1 ms ......................................... wide_state_15x600_checkpoint: Mean +- std dev: 430 ms +- 14 ms ......................................... wide_state_15x600_checkpoint_sync: Mean +- std dev: 426 ms +- 14 ms ......................................... wide_state_9x1200: Mean +- std dev: 27.4 ms +- 0.5 ms ......................................... wide_state_9x1200_sync: Mean +- std dev: 17.9 ms +- 0.3 ms ......................................... wide_state_9x1200_checkpoint: Mean +- std dev: 284 ms +- 15 ms ......................................... wide_state_9x1200_checkpoint_sync: Mean +- std dev: 282 ms +- 17 ms

Check notice on line 1 in libs/sdk-py/langgraph_sdk/auth/__init__.py

View workflow job for this annotation

GitHub Actions / benchmark

Comparison against main

+-----------------------------------------+---------+-----------------------+ | Benchmark | main | changes | +=========================================+=========+=======================+ | fanout_to_subgraph_100x_checkpoint | 819 ms | 774 ms: 1.06x faster | +-----------------------------------------+---------+-----------------------+ | react_agent_100x_checkpoint | 657 ms | 638 ms: 1.03x faster | +-----------------------------------------+---------+-----------------------+ | react_agent_100x_checkpoint_sync | 637 ms | 619 ms: 1.03x faster | +-----------------------------------------+---------+-----------------------+ | fanout_to_subgraph_100x_checkpoint_sync | 980 ms | 968 ms: 1.01x faster | +-----------------------------------------+---------+-----------------------+ | react_agent_10x_checkpoint_sync | 37.2 ms | 36.7 ms: 1.01x faster | +-----------------------------------------+---------+-----------------------+ | react_agent_100x_sync | 275 ms | 272 ms: 1.01x faster | +-----------------------------------------+---------+-----------------------+ | wide_state_15x600_sync | 17.9 ms | 17.8 ms: 1.01x faster | +-----------------------------------------+---------+-----------------------+ | wide_state_9x1200_sync | 17.8 ms | 17.9 ms: 1.00x slower | +-----------------------------------------+---------+-----------------------+ | fanout_to_subgraph_10x_checkpoint | 75.7 ms | 76.5 ms: 1.01x slower | +-----------------------------------------+---------+-----------------------+ | fanout_to_subgraph_10x | 62.6 ms | 63.3 ms: 1.01x slower | +-----------------------------------------+---------+-----------------------+ | fanout_to_subgraph_100x_sync | 528 ms | 535 ms: 1.01x slower | +-----------------------------------------+---------+-----------------------+ | fanout_to_subgraph_10x_sync | 54.0 ms | 55.0 ms: 1.02x slower | +-----------------------------------------+---------+-----------------------+ | fanout_to_subgraph_100x | 648 ms | 662 ms: 1.02x slower | +-----------------------------------------+---------+-----------------------+ | wide_state_9x1200_checkpoint_sync | 276 ms | 282 ms: 1.02x slower | +-----------------------------------------+---------+-----------------------+ | Geometric mean | (ref) | 1.00x faster | +-----------------------------------------+---------+-----------------------+ Benchmark hidden because not significant (14): react_agent_10x_checkpoint, wide_state_9x1200, wide_state_15x600_checkpoint_sync, wide_state_25x300_checkpoint_sync, wide_state_25x300, wide_state_15x600_checkpoint, react_agent_10x_sync, wide_state_25x300_sync, react_agent_100x, react_agent_10x, wide_state_25x300_checkpoint, fanout_to_subgraph_10x_checkpoint_sync, wide_state_15x600, wide_state_9x1200_checkpoint

import inspect
import typing
Expand Down Expand Up @@ -461,6 +461,73 @@
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]])


Expand Down Expand Up @@ -524,6 +591,7 @@
"threads",
"runs",
"crons",
"store",
"value",
)

Expand All @@ -532,6 +600,7 @@
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
Expand Down
157 changes: 152 additions & 5 deletions libs/sdk-py/langgraph_sdk/auth/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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",
Expand All @@ -909,4 +1051,9 @@ class search:
"AssistantsUpdate",
"AssistantsDelete",
"AssistantsSearch",
"StoreGet",
"StoreSearch",
"StoreListNamespaces",
"StorePut",
"StoreDelete",
]
2 changes: 1 addition & 1 deletion libs/sdk-py/pyproject.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
Loading