mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
wakita181009andGitHub
875f20ba9f
This PR adds an aclose method to the LangGraphClient. When using the client in a FastAPI application, it's common to share a single instance across the application's lifespan. The absence of an aclose method makes it difficult to gracefully close the underlying HTTP session on application shutdown. This change enables proper resource management by allowing the client to be closed cleanly.
70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
import ast
|
|
import os
|
|
from itertools import filterfalse
|
|
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]:
|
|
return [n.name for n in node.body if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))]
|
|
|
|
|
|
def find_classes(tree: ast.AST) -> List[Tuple[str, List[str]]]:
|
|
classes = []
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.ClassDef):
|
|
methods = get_class_methods(node)
|
|
classes.append((node.name, methods))
|
|
return classes
|
|
|
|
|
|
def compare_sync_async_methods(sync_methods: List[str], async_methods: List[str]) -> List[str]:
|
|
sync_set = set(sync_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
|
|
|
|
|
|
def main():
|
|
with open(CLIENT_PATH, "r") as file:
|
|
tree = ast.parse(file.read())
|
|
|
|
classes = find_classes(tree)
|
|
|
|
def is_sync(class_spec: Tuple[str, List[str]]) -> bool:
|
|
return class_spec[0].startswith("Sync")
|
|
|
|
sync_class_name_to_methods = {class_name: class_methods for class_name, class_methods in filter(is_sync, classes)}
|
|
async_class_name_to_methods = {class_name: class_methods for class_name, class_methods in filterfalse(is_sync, classes)}
|
|
|
|
mismatches = []
|
|
|
|
for async_class_name, async_class_methods in async_class_name_to_methods.items():
|
|
sync_class_name = "Sync" + async_class_name
|
|
sync_class_methods = sync_class_name_to_methods.get(sync_class_name, [])
|
|
diff = compare_sync_async_methods(sync_class_methods, async_class_methods)
|
|
if diff:
|
|
mismatches.append((sync_class_name, async_class_name, diff))
|
|
|
|
if mismatches:
|
|
error_message = "Mismatches found between sync and async client methods:\n"
|
|
for sync_class_name, async_class_name, diff in mismatches:
|
|
error_message += f"{sync_class_name} vs {async_class_name}:\n"
|
|
for method in diff:
|
|
error_message += f" - {method}\n"
|
|
raise ValueError(error_message)
|
|
|
|
print("All sync and async client methods match.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|