Persist chat prototype

agent.py

@@ -23,19 +23,22 @@ class AgentContext:
     _counter: int = 0
 
     def __init__(
-        self, config: "AgentConfig", id: str | None = None, agent0: "Agent|None" = None
+        self, config: "AgentConfig", id: str | None = None, name: str | None = None, agent0: "Agent|None" = None,
+        log: Log.Log | None = None,
+        paused: bool = False, streaming_agent: "Agent|None" = None,
     ):
         # build context
         self.id = id or str(uuid.uuid4())
+        self.name = name
         self.config = config
-        self.log = Log.Log()
+        self.log = log or Log.Log()
         self.agent0 = agent0 or Agent(0, self.config, self)
-        self.paused = False
-        self.streaming_agent: Agent | None = None
+        self.paused = paused
+        self.streaming_agent = streaming_agent
         self.process: DeferredTask | None = None
         AgentContext._counter += 1
         self.no = AgentContext._counter
-
+        
         self._contexts[self.id] = self
 
     @staticmethod
@@ -317,9 +320,6 @@ class Agent:
                                 agent_response
                             )  # process tools requested in agent message
                             if tools_result:  # final response of message loop available
-                                await self.call_extensions(
-                                    "monologue_end", tools_result=tools_result
-                                )  # call monologue_end extensions
                                 return tools_result  # break the execution if the task is done
 
                     # exceptions inside message loop:
@@ -338,6 +338,12 @@ class Agent:
                     except Exception as e:  # Other exception kill the loop
                         self.handle_critical_exception(e)
 
+                    finally:
+                        # call message_loop_end extensions
+                        await self.call_extensions(
+                            "message_loop_end", loop_data=loop_data
+                        )
+
             # exceptions outside message loop:
             except InterventionException as e:
                 pass  # just start over
@@ -345,6 +351,8 @@ class Agent:
                 self.handle_critical_exception(e)
             finally:
                 self.context.streaming_agent = None  # unset current streamer
+                # call monologue_end extensions
+                await self.call_extensions("monologue_end", loop_data=loop_data)  # type: ignore
 
     def handle_critical_exception(self, exception: Exception):
         if isinstance(exception, HandledException):

python/extensions/message_loop_end/_90_save_chat.py

@@ -0,0 +1,8 @@
+from python.helpers.extension import Extension
+from agent import LoopData
+from python.helpers import persist_chat
+
+
+class SaveChat(Extension):
+    async def execute(self, loop_data: LoopData = LoopData(), **kwargs):
+        persist_chat.save_chat(self.agent.context)

python/helpers/files.py

@@ -1,3 +1,4 @@
+from fnmatch import fnmatch
 import os, re
 
 import re
@@ -62,6 +63,22 @@ def find_file_in_dirs(file_path, backup_dirs):
 def remove_code_fences(text):
     return re.sub(r'~~~\w*\n|~~~', '', text)
 
+def write_file(relative_path:str, content:str):
+    abs_path = get_abs_path(relative_path)
+    os.makedirs(os.path.dirname(abs_path), exist_ok=True)
+    with open(abs_path, 'w') as f:
+        f.write(content)
+
+def delete_file(relative_path:str):
+    abs_path = get_abs_path(relative_path)
+    if os.path.exists(abs_path):
+        os.remove(abs_path)
+
+def list_files(relative_path:str, filter:str="*"):
+    abs_path = get_abs_path(relative_path)
+    if not os.path.exists(abs_path):
+        return []
+    return [file for file in os.listdir(abs_path) if fnmatch(file, filter)]
 
 def get_abs_path(*relative_paths):
     return os.path.join(get_base_dir(), *relative_paths)
@@ -70,7 +87,6 @@ def exists(*relative_paths):
     path = get_abs_path(*relative_paths)
     return os.path.exists(path)
 
-
 def get_base_dir():
     # Get the base directory from the current file path
     base_dir = os.path.dirname(os.path.abspath(os.path.join(__file__,"../../")))

python/helpers/log.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 import json
-from typing import Literal, Optional, Dict
+from typing import Any, Literal, Optional, Dict
 import uuid
 from collections import OrderedDict  # Import OrderedDict
 
@@ -145,7 +145,7 @@ class Log:
 
         self.updates += [item.no]
 
-    def output(self, start=None, end=None):
+    def output(self, start=None, end=None):        
         if start is None:
             start = 0
         if end is None:

python/helpers/persist_chat.py

@@ -0,0 +1,192 @@
+from collections import OrderedDict
+from typing import Any
+import uuid
+from agent import Agent, AgentConfig, AgentContext, HumanMessage, AIMessage
+from python.helpers import files
+import json
+from initialize import initialize
+
+from python.helpers.log import Log, LogItem
+
+CHATS_FOLDER = "tmp/chats"
+LOG_SIZE = 1000
+
+
+def save_chat(context: AgentContext):
+    relative_path = _get_file_path(context.id)
+    data = _serialize_context(context)
+    js = _safe_json_serialize(data, ensure_ascii=False)
+    files.write_file(relative_path, js)
+
+
+def load_chats():
+    json_files = files.list_files("tmp/chats", "*.json")
+    for file in json_files:
+        path = files.get_abs_path(CHATS_FOLDER, file)
+        js = files.read_file(path)
+        data = json.loads(js)
+        ctx = _deserialize_context(data)
+
+
+def remove_chat(ctxid):
+    files.delete_file(_get_file_path(ctxid))
+
+
+def _get_file_path(ctxid: str):
+    return f"{CHATS_FOLDER}/{ctxid}.json"
+
+
+def _serialize_context(context: AgentContext):
+    # serialize agents
+    agents = []
+    agent = context.agent0
+    while agent:
+        agents.append(_serialize_agent(agent))
+        agent = agent.data.get("subordinate", None)
+
+    return {
+        "id": context.id,
+        "agents": agents,
+        "streaming_agent": (
+            context.streaming_agent.number if context.streaming_agent else 0
+        ),
+        "log": _serialize_log(context.log),
+    }
+
+
+def _serialize_agent(agent: Agent):
+    data = {**agent.data}
+    if "superior" in data:
+        del data["superior"]
+    if "subordinate" in data:
+        del data["subordinate"]
+
+    history = []
+    for msg in agent.history:
+        history.append({"type": msg.type, "content": msg.content})
+
+    return {
+        "number": agent.number,
+        "data": data,
+        "history": history,
+    }
+
+
+def _serialize_log(log: Log):
+    return {
+        "guid": log.guid,
+        "logs": [item.output() for item in log.logs[-LOG_SIZE:]]
+,  # serialize LogItem objects
+        "progress": log.progress,
+        "progress_no": log.progress_no,
+    }
+
+
+def _deserialize_context(data):
+    config = initialize()
+    log = _deserialize_log(data.get("log", None))
+
+    context = AgentContext(
+        config=config,
+        id=data.get("id", None),
+        name=data.get("name", None),
+        log=log,
+        paused=True,
+        # agent0=agent0,
+        # streaming_agent=straming_agent,
+    )
+
+    agents = data.get("agents", [])
+    agent0 = _deserialize_agents(agents, config, context)
+    streaming_agent_no = data.get("streaming_agent", 0)
+    straming_agent = (
+        agents[streaming_agent_no] if streaming_agent_no < len(agents) else None
+    )
+
+    context.agent0 = agent0
+    context.streaming_agent = straming_agent
+
+    return context
+
+
+def _deserialize_agents(
+    agents: list[dict[str, Any]], config: AgentConfig, context: AgentContext
+) -> Agent:
+    prev: Agent | None = None
+    zero: Agent | None = None
+
+    for ag in agents:
+        current = Agent(
+            number=ag["number"],
+            config=config,
+            context=context,
+        )
+        current.data = ag.get("data", {})
+        current.history = _deserialize_history(ag.get("history", []))
+
+        if not zero:
+            zero = current
+
+        if prev:
+            prev.set_data("subordinate", current)
+            current.set_data("superior", prev)
+
+    return zero or Agent(0, config, context)
+
+
+def _deserialize_history(history: list[dict[str, Any]]):
+    result = []
+    for hist in history:
+        content = hist.get("content", "")
+        msg = (
+            HumanMessage(content=content)
+            if hist.get("type") == "human"
+            else AIMessage(content=content)
+        )
+        result.append(msg)
+    return result
+
+
+def _deserialize_log(data: dict[str, Any]) -> "Log":
+    log = Log()
+    log.guid = data.get("guid", str(uuid.uuid4()))
+    log.progress = data.get("progress", "")
+    log.progress_no = data.get("progress_no", 0)
+
+    # Deserialize the list of LogItem objects
+    i = 0
+    for item_data in data.get("logs", []):
+        log.logs.append(LogItem(
+            log=log,  # restore the log reference
+            no=item_data["no"],
+            type=item_data["type"],
+            heading=item_data.get("heading", ""),
+            content=item_data.get("content", ""),
+            kvps=OrderedDict(item_data["kvps"]) if item_data["kvps"] else None,
+            temp=item_data.get("temp", False),
+        ))
+        log.updates.append(i)
+        i += 1
+        
+    return log
+
+
+def _safe_json_serialize(obj, **kwargs):
+    def serializer(o):
+        if isinstance(o, dict):
+            return {k: v for k, v in o.items() if is_json_serializable(v)}
+        elif isinstance(o, (list, tuple)):
+            return [item for item in o if is_json_serializable(item)]
+        elif is_json_serializable(o):
+            return o
+        else:
+            return None  # Skip this property
+
+    def is_json_serializable(item):
+        try:
+            json.dumps(item)
+            return True
+        except (TypeError, OverflowError):
+            return False
+
+    return json.dumps(obj, default=serializer, **kwargs)

run_ui.py

@@ -8,9 +8,11 @@ from flask import Flask, request, jsonify, Response
 from flask_basicauth import BasicAuth
 from agent import AgentContext
 from initialize import initialize
+from python.helpers import files
 from python.helpers.files import get_abs_path
 from python.helpers.print_style import PrintStyle
 from python.helpers.dotenv import load_dotenv
+from python.helpers import persist_chat
 
 
 # initialize the internal Flask server
@@ -119,6 +121,7 @@ async def handle_message(sync: bool):
             response = {
                 "ok": True,
                 "message": result,
+                "context": context.id,
             }
         else:
 
@@ -126,6 +129,7 @@ async def handle_message(sync: bool):
             response = {
                 "ok": True,
                 "message": "Message received.",
+                "context": context.id,
             }
 
     except Exception as e:
@@ -183,6 +187,7 @@ async def reset():
         # context instance - get or create
         context = get_context(ctxid)
         context.reset()
+        persist_chat.save_chat(context)
 
         response = {
             "ok": True,
@@ -211,6 +216,7 @@ async def remove():
 
         # context instance - get or create
         AgentContext.remove(ctxid)
+        persist_chat.remove_chat(ctxid)
 
         response = {
             "ok": True,
@@ -235,7 +241,7 @@ async def poll():
 
         # data sent to the server
         input = request.get_json()
-        ctxid = input.get("context", uuid.uuid4())
+        ctxid = input.get("context", None)
         from_no = input.get("log_from", 0)
 
         # context instance - get or create
@@ -286,6 +292,9 @@ def run():
 
     #load env vars
     load_dotenv()
+
+    # initialize contexts from persisted chats
+    persist_chat.load_chats()
     
     # Suppress only request logs but keep the startup messages
     from werkzeug.serving import WSGIRequestHandler

webui/index.js

@@ -73,6 +73,8 @@ async function sendMessage() {
                 } else {
                     toast("Undefined error.", "error")
                 }
+            } else {
+                setContext(response.context)
             }
 
             //setMessage('user', message);
@@ -186,7 +188,8 @@ async function poll() {
 
         if (response.ok) {
 
-            setContext(response.context)
+            if (!context) setContext(response.context)
+            if (response.context != context) return //skip late polls after context change
 
             if (lastLogGuid != response.log_guid) {
                 chatHistory.innerHTML = ""
@@ -408,7 +411,6 @@ function scrollChanged(isAtBottom) {
     const inputAS = Alpine.$data(autoScrollSwitch);
     inputAS.autoScroll = isAtBottom
     // autoScrollSwitch.checked = isAtBottom
-    console.log(isAtBottom)
 }
 
 function updateAfterScroll() {