Message history (#35)

* edit message history

* change to update

* bump package ver

js/examples/editMessageHistory.ts

@@ -0,0 +1,23 @@
+import ProgramState from '../src/index.js';
+
+(async () => {
+    const s = await new ProgramState().fromSGL('http://localhost:30000');
+    await s
+        .add(s.system`You are a helpful assistant`, { id: 'a' })
+        .add(s.user`Tell me a joke`, { id: 'b' })
+        .add(s.assistant`${s.gen('joke')}`, { id: 'c' });
+    // LLM Response
+    console.log(s.get('joke'));
+    // Full message object
+    console.log(s.get('a'));
+    // Reset
+    await s
+        .update('b', s.user`Tell me a short story`, undefined, {
+            deleteMessagesAfter: true,
+        })
+        .add(
+            s.assistant`${s.gen('story', { sampling_params: { max_new_tokens: 24 } })}`,
+            { id: 'd' },
+        );
+    console.log('---RESET---\n', s.get('story'), s.get('d'));
+})();

js/package.json

@@ -1,6 +1,6 @@
 {
     "name": "enochian-js",
-    "version": "0.0.8",
+    "version": "0.0.9",
     "description": "An ergonmic LLM programming library",
     "main": "dist/index.js",
     "type": "module",

js/src/programState.ts

@@ -461,28 +461,99 @@ export default class ProgramState {
         key: string,
         options: { from: 'tools'; tools: T },
     ): (ToolResponse<T> | { toolUsed: 'respondToUser'; response: string })[];
+    get(key: string, options: { from: 'messages' }): Message;
     get<Z extends z.ZodType, T extends ToolUseParams>(
         key: string,
-        options?: { from: 'zod'; schema: Z } | { from: 'tools'; tools: T },
+        options?:
+            | { from: 'zod'; schema: Z }
+            | { from: 'tools'; tools: T }
+            | { from: 'messages' },
     ):
         | string
         | z.infer<Z>
+        | Message
         | (ToolResponse<T> | { toolUsed: 'respondToUser'; response: string })[]
         | undefined {
-        const value = this._answers[key]?.text;
-        if (!value) return undefined;
+        if (options?.from === 'messages') {
+            return this._messages.find((m) => m.id === key);
+        } else {
+            const value = this._answers[key]?.text;
+            if (!value) {
+                // If you can't find it in the answers try in the message histroy
+                const message = this._messages.find((m) => m.id === key);
+                if (message) {
+                    return message.content;
+                } else {
+                    return undefined;
+                }
+            }
 
-        if (!options) {
-            return value as string;
+            if (!options) {
+                return value as string;
+            }
+            if (options.from === 'zod') {
+                return options.schema.parse(JSON.parse(value));
+            }
+            if (options.from === 'tools') {
+                return JSON.parse(value) as
+                    | ToolResponse<T>
+                    | { toolUsed: 'respondToUser'; response: string };
+            }
+        }
+    }
+
+    update(
+        id: string,
+        newMessage: Message,
+        metadata?: { [key: string]: unknown },
+        opts?: { deleteMessagesAfter?: boolean },
+    ): ProgramState;
+    update(
+        id: string,
+        newMessage: Promise<Message>,
+        metadata?: { [key: string]: unknown },
+        opts?: { deleteMessagesAfter?: boolean },
+    ): Promise<ProgramState>;
+    update(
+        id: string,
+        newMessage: AsyncGenerator<Message, Message, undefined>,
+        metadata?: { [key: string]: unknown },
+        opts?: { deleteMessagesAfter?: boolean },
+    ): AsyncGenerator<Message>;
+    update(
+        id: string,
+        newMessage:
+            | Message
+            | Promise<Message>
+            | AsyncGenerator<Message, Message, undefined>,
+        metadata?: { [key: string]: unknown },
+        opts?: { deleteMessagesAfter?: boolean },
+    ): ProgramState | Promise<ProgramState> | AsyncGenerator<Message> {
+        const messageIdx = this._messages.findIndex((m) => m.id === id);
+        if (!messageIdx) {
+            console.warn(`No message with id ${id}`);
+            return this;
         }
-        if (options.from === 'zod') {
-            return options.schema.parse(JSON.parse(value));
+        if (opts?.deleteMessagesAfter) {
+            this._messages.splice(
+                messageIdx,
+                this._messages.length - messageIdx,
+            );
+        } else {
+            this._messages.splice(messageIdx, 1);
         }
-        if (options.from === 'tools') {
-            return JSON.parse(value) as
-                | ToolResponse<T>
-                | { toolUsed: 'respondToUser'; response: string };
+        // @ts-expect-error TS can't infer the overloaded type of `add`
+        return this.add(newMessage, metadata);
+    }
+
+    delete(id: string): ProgramState {
+        const messageIdx = this._messages.findIndex((m) => m.id === id);
+        if (!messageIdx) {
+            console.warn(`No message with id ${id}`);
+            return this;
         }
+        this._messages.splice(messageIdx, 1);
+        return this;
     }
 
     getPrompt(): string {