reverted changes to db connection

Author: codinglabsong

api/llm/agent.py

@@ -1,3 +1,4 @@
+import atexit
 import os
 from functools import lru_cache
 from typing import List, Optional
@@ -13,15 +14,17 @@
 _agent_executor = None
 
 
-@lru_cache(maxsize=1)
+@lru_cache
 def get_checkpointer():
     DATABASE_URL = os.environ.get("DATABASE_URL")
     if not DATABASE_URL:
         raise RuntimeError("DATABASE_URL is not set. Point it to your Neon connection string.")
 
-    checkpointer = PostgresSaver.from_conn_string(DATABASE_URL)
-    checkpointer.setup()  # create tables on first run; no-op afterward
-    return checkpointer
+    cm = PostgresSaver.from_conn_string(DATABASE_URL)
+    saver = cm.__enter__()  # enter the context manager once
+    atexit.register(lambda: cm.__exit__(None, None, None))  # clean shutdown
+    saver.setup()  # create tables on first run; no-op afterward
+    return saver
 
 
 def get_agent():