Cleanup resources on Gradio state deletion

app.py

@@ -34,6 +34,10 @@ def load():
 
     return ChatState(database, messenger, chat)
 
+def eject(state):
+    state.database.cleanup()
+    stat.chat.cleanup()
+
 def upload(data, state):
     return state.database(data)
 
@@ -53,7 +57,10 @@ def prompt(message, history, state):
 #
 #
 with gr.Blocks() as demo:
-    state = gr.State(load)
+    state = gr.State(
+        value=load,
+        delete_callback=eject,
+    )
     with gr.Row():
         with gr.Column():
             repository = gr.Textbox(
@@ -67,7 +74,10 @@ with gr.Blocks() as demo:
             )
             data.upload(
                 fn=upload,
-                inputs=[data, state],
+                inputs=[
+                    data,
+                    state,
+                ],
                 outputs=repository,
             )
 
@@ -76,7 +86,11 @@ with gr.Blocks() as demo:
             interaction = gr.Textbox()
             interaction.submit(
                 fn=prompt,
-                inputs=[interaction, chatbot, state],
+                inputs=[
+                    interaction,
+                    chatbot,
+                    state,
+                ],
                 outputs=chatbot,
             )
 

mylib/__init__.py

@@ -1,5 +1,5 @@
 from ._chat import ChatController
-from ._files import FileManager
+from ._files import FileManager, VectorStoreManager
 from ._errors import ErrorLogger
 from ._logging import Logger
 from ._message import MessageHandler

mylib/_chat.py

@@ -61,6 +61,11 @@ class ChatController:
 
         return self.send(prompt)
 
+    def cleanup(self):
+        self.client.beta.threads.delete(self.thread.id)
+        self.client.beta.assistants.delete(self.assistant.id)
+        self.attached = False
+
     def send(self, content):
         self.client.beta.threads.messages.create(
             self.thread.id,

mylib/_files.py

@@ -50,39 +50,44 @@ class FileStream:
             s.close()
         self.streams.clear()
 
-class FileManager:
-    def __init__(self, client, prefix, batch_size=20):
+class VectorStoreManager:
+    def __init__(self, client, vector_store_id):
         self.client = client
+        self.vector_store_id = vector_store_id
+
+    def __iter__(self):
+        while True:
+            vs_files = self.client.beta.vector_stores.files.list(
+                vector_store_id=self.vector_store_id,
+                **kwargs,
+            )
+            for f in vs_files.data:
+                yield f.id
+
+            if not vs_files.has_more:
+                break
+
+    def cleanup(self):
+        for i in self:
+            self.client.files.delete(i)
+        self.client.beta.vector_stores.delete(self.vector_store_id)
+
+class FileManager(VectorStoreManager):
+    def __init__(self, client, prefix, batch_size=20):
+        super().__init__(client, None)
+
         self.prefix = prefix
         self.batch_size = batch_size
-
         self.storage = set()
-        self.vector_store_id = None
 
     def __bool__(self):
         return self.vector_store_id is not None
 
-    def __iter__(self):
-        if self:
-            kwargs = {}
-            while True:
-                vs_files = self.client.beta.vector_stores.files.list(
-                    vector_store_id=self.vector_store_id,
-                    **kwargs,
-                )
-                for f in vs_files.data:
-                    result = self.client.files.retrieve(f.id)
-                    yield result.filename
-
-                if not vs_files.has_more:
-                    break
-                kwargs['after'] = vs_files.after
-
     def __call__(self, paths):
         files = []
         self.test_and_setup()
 
-        for p in self.ls(paths):
+        for p in self.each(paths):
             with FileStream(p) as stream:
                 for s in stream:
                     if s.checksum not in self.storage:
@@ -92,7 +97,19 @@ class FileManager:
                     self.put(files)
                     files.clear()
 
-        return '\n'.join(self)
+        return '\n'.join(self.ls())
+
+    def ls(self):
+        for i in self:
+            self.client.files.retrieve(i)
+            yield i.filename
+
+    def cleanup(self):
+        if self.storage:
+            assert self.vector_store_id is not None
+            super().cleanup()
+            self.storage.clear()
+        self.vector_store_id = None
 
     def test_and_setup(self):
         if self:
@@ -105,7 +122,7 @@ class FileManager:
             )
             self.vector_store_id = vector_store.id
 
-    def ls(self, paths):
+    def each(self, paths):
         left = 0
         while left < len(paths):
             right = left + self.batch_size

tools/cleanup.py

@@ -0,0 +1,53 @@
+import time
+from argparse imAAport ArgumentParser
+
+from scipy import constants
+from openai import OpenAI
+
+from mylib import Logger, VectorStoreManager
+
+class AgeCheck:
+    def __init__(self):
+        self.now = time.time()
+
+    def __contains__(self, other):
+        raise NotImplementedError()
+
+class NoAgeCheck(AgeCheck):
+    def	__contains__(self, other):
+        return False
+
+class HourAgeCheck(AgeCheck):
+    def	__init__(self, hours):
+        super().__init__()
+        self.limit = hours
+
+    def __contains__(self, other):
+        age = (other - self.now) * constants.hour
+        return age < self.limit
+
+if __name__ == '__main__':
+    arguments = ArgumentParser()
+    arguments.add_argument('--max-age-hours', type=int)
+    args = arguments.parse_args()
+
+    client = OpenAI()
+    if args.max_age_hours is None:
+        acheck = NoAgeCheck()
+    else:
+        acheck = HourAgeCheck(args.max_age_hours)
+
+    while True:
+        assistants = client.beta.assistants.list()
+        for a in assistants:
+            if a.created_at in acheck:
+                continue
+            if a.tool_resources.file_search is not None:
+                for i in a.tool_resources.file_search.vector_store_ids:
+                    Logger.warning(f'{a.id} {i}')
+                    vsm = VectorStoreManager(i)
+                    vsm.cleanup()
+            client.beta.assistants.delete(a.id)
+
+        if not assistants.has_more:
+            break