Merge branch 'f/switch-to-pg' into f/add-user-queries

memory-store/migrations/000015_entries.down.sql

@@ -1,5 +1,9 @@
 BEGIN;
 
+DROP TRIGGER IF EXISTS trg_optimized_update_token_count_after ON entries;
+
+DROP FUNCTION IF EXISTS optimized_update_token_count_after;
+
 -- Drop foreign key constraint if it exists
 ALTER TABLE IF EXISTS entries
 DROP CONSTRAINT IF EXISTS fk_entries_session;

memory-store/migrations/000015_entries.up.sql

@@ -14,8 +14,8 @@ CREATE TABLE IF NOT EXISTS entries (
     content JSONB[] NOT NULL,
     tool_call_id TEXT DEFAULT NULL,
     tool_calls JSONB[] NOT NULL DEFAULT '{}',
-    token_count INTEGER NOT NULL,
     model TEXT NOT NULL,
+    token_count INTEGER DEFAULT NULL,
     created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
     timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
     CONSTRAINT pk_entries PRIMARY KEY (session_id, entry_id, created_at)
@@ -52,4 +52,37 @@ BEGIN
     END IF;
 END $$;
 
+-- TODO: We should consider using a timescale background job to update the token count
+-- instead of a trigger.
+-- https://docs.timescale.com/use-timescale/latest/user-defined-actions/create-and-register/
+CREATE
+OR REPLACE FUNCTION optimized_update_token_count_after () RETURNS TRIGGER AS $$
+DECLARE
+    token_count INTEGER;
+BEGIN
+    -- Compute token_count outside the UPDATE statement for clarity and potential optimization
+    token_count := cardinality(
+        ai.openai_tokenize(
+            'gpt-4o', -- FIXME: Use `NEW.model`
+            array_to_string(NEW.content::TEXT[], ' ')
+        )
+    );
+
+    -- Perform the update only if token_count differs
+    IF token_count <> NEW.token_count THEN
+        UPDATE entries
+        SET token_count = token_count
+        WHERE entry_id = NEW.entry_id;
+    END IF;
+
+    RETURN NULL;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE TRIGGER trg_optimized_update_token_count_after
+AFTER INSERT
+OR
+UPDATE ON entries FOR EACH ROW
+EXECUTE FUNCTION optimized_update_token_count_after ();
+
 COMMIT;