Skip to content

Commit

Permalink
[SDK] Add HTTPException type
Browse files Browse the repository at this point in the history
To make it easier to raise exceptions with custom status codes in the imported file.
  • Loading branch information
hinthornw authored Dec 16, 2024
2 parents ca0ff1d + 343dc2d commit 9cbef9b
Show file tree
Hide file tree
Showing 5 changed files with 252 additions and 182 deletions.
2 changes: 1 addition & 1 deletion libs/sdk-py/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -20,5 +20,5 @@ lint lint_diff:
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE) || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)

format format_diff:
poetry run ruff format $(PYTHON_FILES)
poetry run ruff check --select I --fix $(PYTHON_FILES)
poetry run ruff format $(PYTHON_FILES)
145 changes: 74 additions & 71 deletions libs/sdk-py/langgraph_sdk/auth/__init__.py
Original file line number Diff line number Diff line change
@@ -1,23 +1,13 @@
from __future__ import annotations

import inspect
import typing
from collections.abc import Callable, Sequence
from typing import (
Any,
Generic,
Literal,
Optional,
Protocol,
TypeVar,
Union,
cast,
overload,
)

from langgraph_sdk.auth import types
from langgraph_sdk.auth import exceptions, types

TH = TypeVar("TH", bound=types.Handler)
AH = TypeVar("AH", bound=types.Authenticator)
TH = typing.TypeVar("TH", bound=types.Handler)
AH = typing.TypeVar("AH", bound=types.Authenticator)


class Auth:
Expand Down Expand Up @@ -77,13 +67,19 @@ async def authorize_thread_create(params: Auth.on.threads.create.value):
Provides access to all type definitions used in the auth system,
like ThreadsCreate, AssistantsRead, etc."""

exceptions = exceptions
"""Reference to auth exception definitions.
Provides access to all exception definitions used in the auth system,
like HTTPException, etc."""

def __init__(self) -> None:
self.on = _On(self)
# These are accessed by the API. Changes to their names or types is
# will be considered a breaking change.
self._handlers: dict[tuple[str, str], list[types.Handler]] = {}
self._global_handlers: list[types.Handler] = []
self._authenticate_handler: Optional[types.Authenticator] = None
self._authenticate_handler: typing.Optional[types.Authenticator] = None
self._handler_cache: dict[tuple[str, str], types.Handler] = {}

def authenticate(self, fn: AH) -> AH:
Expand Down Expand Up @@ -145,24 +141,26 @@ async def authenticate(

## Helper types & utilities

V = TypeVar("V", contravariant=True)
V = typing.TypeVar("V", contravariant=True)


class _ActionHandler(Protocol[V]):
class _ActionHandler(typing.Protocol[V]):
async def __call__(
self, *, ctx: types.AuthContext, value: V
) -> types.HandlerResult: ...


T = TypeVar("T", covariant=True)
T = typing.TypeVar("T", covariant=True)


class _ResourceActionOn(Generic[T]):
class _ResourceActionOn(typing.Generic[T]):
def __init__(
self,
auth: Auth,
resource: Literal["threads", "crons", "assistants"],
action: Literal["create", "read", "update", "delete", "search", "create_run"],
resource: typing.Literal["threads", "crons", "assistants"],
action: typing.Literal[
"create", "read", "update", "delete", "search", "create_run"
],
value: type[T],
) -> None:
self.auth = auth
Expand All @@ -176,19 +174,19 @@ def __call__(self, fn: _ActionHandler[T]) -> _ActionHandler[T]:
return fn


VCreate = TypeVar("VCreate", covariant=True)
VUpdate = TypeVar("VUpdate", covariant=True)
VRead = TypeVar("VRead", covariant=True)
VDelete = TypeVar("VDelete", covariant=True)
VSearch = TypeVar("VSearch", covariant=True)
VCreate = typing.TypeVar("VCreate", covariant=True)
VUpdate = typing.TypeVar("VUpdate", covariant=True)
VRead = typing.TypeVar("VRead", covariant=True)
VDelete = typing.TypeVar("VDelete", covariant=True)
VSearch = typing.TypeVar("VSearch", covariant=True)


class _ResourceOn(Generic[VCreate, VRead, VUpdate, VDelete, VSearch]):
class _ResourceOn(typing.Generic[VCreate, VRead, VUpdate, VDelete, VSearch]):
"""
Generic base class for resource-specific handlers.
"""

value: type[Union[VCreate, VUpdate, VRead, VDelete, VSearch]]
value: type[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]]

Create: type[VCreate]
Read: type[VRead]
Expand All @@ -199,7 +197,7 @@ class _ResourceOn(Generic[VCreate, VRead, VUpdate, VDelete, VSearch]):
def __init__(
self,
auth: Auth,
resource: Literal["threads", "crons", "assistants"],
resource: typing.Literal["threads", "crons", "assistants"],
) -> None:
self.auth = auth
self.resource = resource
Expand All @@ -219,56 +217,58 @@ def __init__(
auth, resource, "search", self.Search
)

@overload
@typing.overload
def __call__(
self,
fn: Union[
_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
_ActionHandler[dict[str, Any]],
fn: typing.Union[
_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
_ActionHandler[dict[str, typing.Any]],
],
) -> _ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]]: ...
) -> _ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]]: ...

@overload
@typing.overload
def __call__(
self,
*,
resources: Union[str, Sequence[str]],
actions: Optional[Union[str, Sequence[str]]] = None,
resources: typing.Union[str, Sequence[str]],
actions: typing.Optional[typing.Union[str, Sequence[str]]] = None,
) -> Callable[
[_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]]],
_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
[_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]]],
_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
]: ...

def __call__(
self,
fn: Union[
_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
_ActionHandler[dict[str, Any]],
fn: typing.Union[
_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
_ActionHandler[dict[str, typing.Any]],
None,
] = None,
*,
resources: Union[str, Sequence[str], None] = None,
actions: Optional[Union[str, Sequence[str]]] = None,
) -> Union[
_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
resources: typing.Union[str, Sequence[str], None] = None,
actions: typing.Optional[typing.Union[str, Sequence[str]]] = None,
) -> typing.Union[
_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
Callable[
[_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]]],
_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
[_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]]],
_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
],
]:
if fn is not None:
_validate_handler(fn)
return cast(
_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
return typing.cast(
_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
_register_handler(self.auth, self.resource, "*", fn),
)

def decorator(
handler: _ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
) -> _ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]]:
handler: _ActionHandler[
typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]
],
) -> _ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]]:
_validate_handler(handler)
return cast(
_ActionHandler[Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
return typing.cast(
_ActionHandler[typing.Union[VCreate, VUpdate, VRead, VDelete, VSearch]],
_register_handler(self.auth, self.resource, "*", handler),
)

Expand All @@ -284,7 +284,7 @@ class _AssistantsOn(
types.AssistantsSearch,
]
):
value = Union[
value = typing.Union[
types.AssistantsCreate,
types.AssistantsRead,
types.AssistantsUpdate,
Expand All @@ -307,7 +307,7 @@ class _ThreadsOn(
types.ThreadsSearch,
]
):
value = Union[
value = typing.Union[
type[types.ThreadsCreate],
type[types.ThreadsRead],
type[types.ThreadsUpdate],
Expand All @@ -325,7 +325,7 @@ class _ThreadsOn(
def __init__(
self,
auth: Auth,
resource: Literal["threads", "crons", "assistants"],
resource: typing.Literal["threads", "crons", "assistants"],
) -> None:
super().__init__(auth, resource)
self.create_run: _ResourceActionOn[types.RunsCreate] = _ResourceActionOn(
Expand All @@ -343,7 +343,7 @@ class _CronsOn(
]
):
value = type[
Union[
typing.Union[
types.CronsCreate,
types.CronsRead,
types.CronsUpdate,
Expand All @@ -359,7 +359,7 @@ class _CronsOn(
Search = types.CronsSearch


AHO = TypeVar("AHO", bound=_ActionHandler[dict[str, Any]])
AHO = typing.TypeVar("AHO", bound=_ActionHandler[dict[str, typing.Any]])


class _On:
Expand All @@ -381,26 +381,26 @@ def __init__(self, auth: Auth) -> None:
self.assistants = _AssistantsOn(auth, "assistants")
self.threads = _ThreadsOn(auth, "threads")
self.crons = _CronsOn(auth, "crons")
self.value = dict[str, Any]
self.value = dict[str, typing.Any]

@overload
@typing.overload
def __call__(
self,
*,
resources: Union[str, Sequence[str]],
actions: Optional[Union[str, Sequence[str]]] = None,
resources: typing.Union[str, Sequence[str]],
actions: typing.Optional[typing.Union[str, Sequence[str]]] = None,
) -> Callable[[AHO], AHO]: ...

@overload
@typing.overload
def __call__(self, fn: AHO) -> AHO: ...

def __call__(
self,
fn: Optional[AHO] = None,
fn: typing.Optional[AHO] = None,
*,
resources: Union[str, Sequence[str], None] = None,
actions: Optional[Union[str, Sequence[str]]] = None,
) -> Union[AHO, Callable[[AHO], AHO]]:
resources: typing.Union[str, Sequence[str], None] = None,
actions: typing.Optional[typing.Union[str, Sequence[str]]] = 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:
Expand Down Expand Up @@ -439,7 +439,10 @@ def decorator(handler: AHO) -> AHO:


def _register_handler(
auth: Auth, resource: Optional[str], action: Optional[str], fn: types.Handler
auth: Auth,
resource: typing.Optional[str],
action: typing.Optional[str],
fn: types.Handler,
) -> types.Handler:
_validate_handler(fn)
resource = resource or "*"
Expand All @@ -457,7 +460,7 @@ def _register_handler(
return fn


def _validate_handler(fn: Callable[..., Any]) -> None:
def _validate_handler(fn: Callable[..., typing.Any]) -> None:
"""Validates that an auth handler function meets the required signature.
Auth handlers must:
Expand Down Expand Up @@ -486,4 +489,4 @@ def _validate_handler(fn: Callable[..., Any]) -> None:
)


__all__ = ["Auth", "types"]
__all__ = ["Auth", "types", "exceptions"]
73 changes: 73 additions & 0 deletions libs/sdk-py/langgraph_sdk/auth/exceptions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
"""Exceptions used in the auth system."""

import http
import typing


class HTTPException(Exception):
"""HTTP exception that you can raise to return a specific HTTP error response.
Since this is defined in the auth module, we default to a 401 status code.
Args:
status_code (int, optional): HTTP status code for the error. Defaults to 401 "Unauthorized".
detail (str | None, optional): Detailed error message. If None, uses a default
message based on the status code.
headers (typing.Mapping[str, str] | None, optional): Additional HTTP headers to
include in the error response.
Attributes:
status_code (int): The HTTP status code of the error
detail (str): The error message or description
headers (typing.Mapping[str, str] | None): Additional HTTP headers
Example:
Default:
```python
raise HTTPException()
# HTTPException(status_code=401, detail='Unauthorized')
```
Add headers:
```python
raise HTTPException(headers={"X-Custom-Header": "Custom Value"})
# HTTPException(status_code=401, detail='Unauthorized', headers={"WWW-Authenticate": "Bearer"})
```
Custom error:
```python
raise HTTPException(status_code=404, detail="Not found")
```
"""

def __init__(
self,
status_code: int = 401,
detail: typing.Optional[str] = None,
headers: typing.Optional[typing.Mapping[str, str]] = None,
) -> None:
if detail is None:
detail = http.HTTPStatus(status_code).phrase
self.status_code = status_code
self.detail = detail
self.headers = headers

def __str__(self) -> str:
"""Return a string representation of the HTTP exception.
Returns:
str: A string in the format 'status_code: detail'
"""
return f"{self.status_code}: {self.detail}"

def __repr__(self) -> str:
"""Return a detailed string representation of the HTTP exception.
Returns:
str: A string representation showing the class name and all attributes
"""
class_name = self.__class__.__name__
return f"{class_name}(status_code={self.status_code!r}, detail={self.detail!r})"


__all__ = ["HTTPException"]
Loading

0 comments on commit 9cbef9b

Please sign in to comment.