setup context manger for SignalClient and LLMClient

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()

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__":

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()