diff --git a/python/signal_bot/llm_client.py b/python/signal_bot/llm_client.py index dbccc76..90e6f81 100644 --- a/python/signal_bot/llm_client.py +++ b/python/signal_bot/llm_client.py @@ -4,7 +4,7 @@ from __future__ import annotations import base64 import logging -from typing import Any +from typing import Any, Self import httpx @@ -66,6 +66,14 @@ class LLMClient: response.raise_for_status() return [m["name"] for m in response.json().get("models", [])] + def __enter__(self) -> Self: + """Enter the context manager.""" + return self + + def __exit__(self, *args: object) -> None: + """Close the HTTP client on exit.""" + self.close() + def close(self) -> None: """Close the HTTP client.""" self._client.close() diff --git a/python/signal_bot/main.py b/python/signal_bot/main.py index 5ecc7ab..c6ce0bd 100644 --- a/python/signal_bot/main.py +++ b/python/signal_bot/main.py @@ -165,15 +165,13 @@ def main( inventory_file=inventory_file, ) - signal = SignalClient(config.signal_api_url, config.phone_number) - llm = LLMClient(model=llm_model, host=llm_host, port=llm_port) - registry = DeviceRegistry(signal, Path(registry_file)) - try: + with ( + SignalClient(config.signal_api_url, config.phone_number) as signal, + LLMClient(model=llm_model, host=llm_host, port=llm_port) as llm, + ): + registry = DeviceRegistry(signal, Path(registry_file)) run_loop(config, signal, llm, registry) - finally: - signal.close() - llm.close() if __name__ == "__main__": diff --git a/python/signal_bot/signal_client.py b/python/signal_bot/signal_client.py index ea18e4f..cb8d91d 100644 --- a/python/signal_bot/signal_client.py +++ b/python/signal_bot/signal_client.py @@ -4,7 +4,7 @@ from __future__ import annotations import json import logging -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Self import httpx import websockets.sync.client @@ -121,6 +121,14 @@ class SignalClient: else: self.send(message.source, text) + def __enter__(self) -> Self: + """Enter the context manager.""" + return self + + def __exit__(self, *args: object) -> None: + """Close the HTTP client on exit.""" + self.close() + def close(self) -> None: """Close the HTTP client.""" self._client.close()