diff --git a/.github/scripts/check_sdk_methods.py b/.github/scripts/check_sdk_methods.py index 6a7d18bf0..5daeed7e5 100644 --- a/.github/scripts/check_sdk_methods.py +++ b/.github/scripts/check_sdk_methods.py @@ -1,10 +1,15 @@ import ast import os from itertools import filterfalse -from typing import List, Tuple +from typing import Dict, List, Tuple ROOT_PATH = os.path.abspath(os.path.join(__file__, "..", "..", "..")) CLIENT_PATH = os.path.join(ROOT_PATH, "libs", "sdk-py", "langgraph_sdk", "client.py") +ASYNC_TO_SYNC_METHOD_MAP: Dict[str, str] = { + "aclose": "close", + "__aenter__": "__enter__", + "__aexit__": "__exit__", +} def get_class_methods(node: ast.ClassDef) -> List[str]: @@ -22,7 +27,7 @@ def find_classes(tree: ast.AST) -> List[Tuple[str, List[str]]]: def compare_sync_async_methods(sync_methods: List[str], async_methods: List[str]) -> List[str]: sync_set = set(sync_methods) - async_set = set(async_methods) + async_set = {ASYNC_TO_SYNC_METHOD_MAP.get(async_method, async_method) for async_method in async_methods} missing_in_sync = list(async_set - sync_set) missing_in_async = list(sync_set - async_set) return missing_in_sync + missing_in_async diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index b481f5119..a03e6390b 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -16,6 +16,7 @@ import os import re import sys from collections.abc import AsyncIterator, Iterator, Sequence +from types import TracebackType from typing import ( Any, Callable, @@ -236,6 +237,24 @@ class LangGraphClient: self.crons = CronClient(self.http) self.store = StoreClient(self.http) + async def __aenter__(self) -> LangGraphClient: + """Enter the async context manager.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Exit the async context manager.""" + await self.aclose() + + async def aclose(self) -> None: + """Close the underlying HTTP client.""" + if hasattr(self, "http"): + await self.http.client.aclose() + class HttpClient: """Handle async requests to the LangGraph API. @@ -3212,6 +3231,24 @@ class SyncLangGraphClient: self.crons = SyncCronClient(self.http) self.store = SyncStoreClient(self.http) + def __enter__(self) -> SyncLangGraphClient: + """Enter the sync context manager.""" + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Exit the sync context manager.""" + self.close() + + def close(self) -> None: + """Close the underlying HTTP client.""" + if hasattr(self, "http"): + self.http.client.close() + class SyncHttpClient: """Handle synchronous requests to the LangGraph API.