diff --git a/src/openjarvis/sdk.py b/src/openjarvis/sdk.py index dfb9e0fa..95d3f497 100644 --- a/src/openjarvis/sdk.py +++ b/src/openjarvis/sdk.py @@ -108,6 +108,12 @@ class MemoryHandle: self._backend.close() self._backend = None + def __enter__(self) -> MemoryHandle: + return self + + def __exit__(self, *exc: Any) -> None: + self.close() + class Jarvis: """High-level OpenJarvis SDK. @@ -116,9 +122,13 @@ class Jarvis: from openjarvis import Jarvis + with Jarvis() as j: + response = j.ask("Hello, what can you do?") + print(response) + + # Or without context manager: j = Jarvis() - response = j.ask("Hello, what can you do?") - print(response) + response = j.ask("Hello") j.close() """ @@ -479,5 +489,11 @@ class Jarvis: self._audit_logger = None self._engine = None + def __enter__(self) -> Jarvis: + return self + + def __exit__(self, *exc: Any) -> None: + self.close() + __all__ = ["Jarvis", "JarvisSystem", "MemoryHandle", "SystemBuilder"] diff --git a/src/openjarvis/system.py b/src/openjarvis/system.py index 3046ec7a..aab2fc7c 100644 --- a/src/openjarvis/system.py +++ b/src/openjarvis/system.py @@ -274,6 +274,12 @@ class JarvisSystem: if self.trace_store and hasattr(self.trace_store, "close"): self.trace_store.close() + def __enter__(self) -> JarvisSystem: + return self + + def __exit__(self, *exc: Any) -> None: + self.close() + class SystemBuilder: """Config-driven fluent builder for JarvisSystem."""