From 2d5f6a2797e6c7ca0b1edc38e18d187681bd1d85 Mon Sep 17 00:00:00 2001
From: James Ketrenos <james_git@ketrenos.com>
Date: Mon, 31 Mar 2025 15:26:41 -0700
Subject: [PATCH] Save and restore contexts

---
 src/.gitignore |  2 ++
 src/server.py  | 61 ++++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 src/.gitignore

diff --git a/src/.gitignore b/src/.gitignore
new file mode 100644
index 0000000..f936fe9
--- /dev/null
+++ b/src/.gitignore
@@ -0,0 +1,2 @@
+sessions
+
diff --git a/src/server.py b/src/server.py
index 3c34b9d..68a7491 100644
--- a/src/server.py
+++ b/src/server.py
@@ -331,8 +331,7 @@ def is_valid_uuid(value):
         uuid_obj = uuid.UUID(value, version=4)
         return str(uuid_obj) == value
     except (ValueError, TypeError):
-        return False
-    
+        return False    
 
 def default_tools(tools):
     return [{**tool, "enabled": True} for tool in tools]
@@ -394,7 +393,8 @@ class WebServer:
                 if not response:
                     return JSONResponse({ "error": "Usage: { reset: rag|tools|history|system-prompt}"})
                 else:
-                    return JSONResponse(response);
+                    self.save_context(context_id)
+                    return JSONResponse(response)
         
             except:
                 return JSONResponse({ "error": "Usage: { reset: rag|tools|history|system-prompt}"})
@@ -411,6 +411,7 @@ class WebServer:
             if not system_prompt:
                 return JSONResponse({ "status": "error", "message": "System prompt can not be empty." })
             context["system"] =  [{"role": "system", "content": system_prompt}]
+            self.save_context(context_id)
             return JSONResponse({ "system-prompt": system_prompt })
         
         @self.app.get('/api/system-prompt/{context_id}')
@@ -429,6 +430,8 @@ class WebServer:
                 async for message in self.chat(context=context, content=data['content']):
                     # Convert to JSON and add newline
                     yield json.dumps(message) + "\n"
+                    # Save the history as its generated
+                    self.save_context(context_id)
                     # Explicitly flush after each yield
                     await asyncio.sleep(0)  # Allow the event loop to process the write
             
@@ -472,6 +475,7 @@ class WebServer:
                 for tool in context["tools"]:
                     if modify == tool["function"]["name"]:
                         tool["enabled"] = enabled
+                        self.save_context(context_id)
                         return JSONResponse(context["tools"])
                 return JSONResponse({ "status": f"{modify} not found in tools." }), 404
             except:
@@ -494,7 +498,56 @@ class WebServer:
                 return FileResponse(full_path)
             self.logging.info(f"Serve index.html for {path}")
             return FileResponse('/opt/airc/src/ketr-chat/build/index.html')
+
+    def save_context(self, session_id):
+        """
+        Serialize a Python dictionary to a file in the sessions directory.
         
+        Args:
+            data: Dictionary containing the session data
+            session_id: UUID string for the context. If it doesn't exist, it is created
+            
+        Returns:
+            The session_id used for the file
+        """
+        context = self.upsert_context(session_id)
+
+        # Create sessions directory if it doesn't exist
+        if not os.path.exists("sessions"):
+            os.makedirs("sessions")
+        
+        # Create the full file path
+        file_path = os.path.join("sessions", session_id)
+        
+        # Serialize the data to JSON and write to file
+        with open(file_path, 'w') as f:
+            json.dump(context, f)
+            
+        return session_id
+
+    def load_context(self, session_id):
+        """
+        Load a serialized Python dictionary from a file in the sessions directory.
+        
+        Args:
+            session_id: UUID string for the filename
+            
+        Returns:
+            The deserialized dictionary, or a new context if it doesn't exist on disk.
+        """
+        file_path = os.path.join("sessions", session_id)
+        
+        # Check if the file exists
+        if not os.path.exists(file_path):
+            return self.create_context(session_id)
+        
+        # Read and deserialize the data
+        with open(file_path, 'r') as f:
+            self.contexts[session_id] = json.load(f)
+            
+        return self.contexts[session_id]
+
+
     def create_context(self, context_id = None):
         if not context_id:
             context_id = str(uuid.uuid4())
@@ -517,7 +570,7 @@ class WebServer:
             logging.info(f"Context {context_id} found.")
             return self.contexts[context_id]
         logging.info(f"Context {context_id} not found. Creating new context.")
-        return self.create_context(context_id)
+        return self.load_context(context_id)
 
     async def chat(self, context, content):
         content = content.strip()