From 26104b4989ed21d0df136fe2e4050c2590bd1611 Mon Sep 17 00:00:00 2001
From: Shreya Shankar <ss.shankar505@gmail.com>
Date: Wed, 2 Oct 2024 08:39:35 -0700
Subject: [PATCH] fix: add context manager cache for clearing util

---
 docetl/operations/utils.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/docetl/operations/utils.py b/docetl/operations/utils.py
index e7fbd917..d3b5e3da 100644
--- a/docetl/operations/utils.py
+++ b/docetl/operations/utils.py
@@ -35,6 +35,7 @@
 cache = Cache(LLM_CACHE_DIR)
 cache.close()
 
+
 def freezeargs(func):
     """
     Decorator to convert mutable dictionary arguments into immutable.
@@ -135,8 +136,8 @@ def clear_cache(console: Console = Console()):
     """
     console.log("[bold yellow]Clearing LLM cache...[/bold yellow]")
     try:
-        cache.clear()
-        cache.close()
+        with cache as c:
+            c.clear()
         # Remove all files in the cache directory
         cache_dir = CACHE_DIR
         if not os.path.exists(cache_dir):
@@ -413,6 +414,7 @@ def call_llm(
                 # TODO: HITL
                 return {}
 
+
 class InvalidOutputError(Exception):
     """
     Custom exception raised when the LLM output is invalid or cannot be parsed.
@@ -455,6 +457,7 @@ def target():
 
     return decorator
 
+
 def call_llm_with_cache(
     model: str,
     op_type: str,