From 47519108ab3a9d678e5bfedf94742ca577e665a9 Mon Sep 17 00:00:00 2001
From: Diwank Singh Tomer <diwank.singh@gmail.com>
Date: Sat, 14 Dec 2024 22:58:23 +0530
Subject: [PATCH] feat(memory-store): Add entry tables

Signed-off-by: Diwank Singh Tomer <diwank.singh@gmail.com>
---
 .../migrations/000015_entries.down.sql        | 16 ++++++
 memory-store/migrations/000015_entries.up.sql | 55 +++++++++++++++++++
 .../000016_entry_relations.down.sql           | 12 ++++
 .../migrations/000016_entry_relations.up.sql  | 55 +++++++++++++++++++
 4 files changed, 138 insertions(+)
 create mode 100644 memory-store/migrations/000015_entries.down.sql
 create mode 100644 memory-store/migrations/000015_entries.up.sql
 create mode 100644 memory-store/migrations/000016_entry_relations.down.sql
 create mode 100644 memory-store/migrations/000016_entry_relations.up.sql

diff --git a/memory-store/migrations/000015_entries.down.sql b/memory-store/migrations/000015_entries.down.sql
new file mode 100644
index 000000000..36ec58280
--- /dev/null
+++ b/memory-store/migrations/000015_entries.down.sql
@@ -0,0 +1,16 @@
+BEGIN;
+
+-- Drop foreign key constraint if it exists
+ALTER TABLE IF EXISTS entries
+DROP CONSTRAINT IF EXISTS fk_entries_session;
+
+-- Drop indexes
+DROP INDEX IF EXISTS idx_entries_by_session;
+
+-- Drop the hypertable (this will also drop the table)
+DROP TABLE IF EXISTS entries;
+
+-- Drop the enum type
+DROP TYPE IF EXISTS chat_role;
+
+COMMIT;
diff --git a/memory-store/migrations/000015_entries.up.sql b/memory-store/migrations/000015_entries.up.sql
new file mode 100644
index 000000000..e03573464
--- /dev/null
+++ b/memory-store/migrations/000015_entries.up.sql
@@ -0,0 +1,55 @@
+BEGIN;
+
+-- Create chat_role enum
+CREATE TYPE chat_role AS ENUM('user', 'assistant', 'tool', 'system');
+
+-- Create entries table
+CREATE TABLE IF NOT EXISTS entries (
+    session_id UUID NOT NULL,
+    entry_id UUID NOT NULL,
+    source TEXT NOT NULL,
+    role chat_role NOT NULL,
+    event_type TEXT NOT NULL DEFAULT 'message.create',
+    name TEXT,
+    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,
+    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)
+);
+
+-- Convert to hypertable if not already
+SELECT
+    create_hypertable (
+        'entries',
+        by_range ('created_at', INTERVAL '1 day'),
+        if_not_exists => TRUE
+    );
+
+SELECT
+    add_dimension (
+        'entries',
+        by_hash ('session_id', 2),
+        if_not_exists => TRUE
+    );
+
+-- Create indexes for efficient querying
+CREATE INDEX IF NOT EXISTS idx_entries_by_session ON entries (session_id DESC, entry_id DESC);
+
+-- Add foreign key constraint to sessions table
+DO $$
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1 FROM pg_constraint WHERE conname = 'fk_entries_session'
+    ) THEN
+        ALTER TABLE entries
+        ADD CONSTRAINT fk_entries_session
+        FOREIGN KEY (session_id)
+        REFERENCES sessions(session_id);
+    END IF;
+END $$;
+
+COMMIT;
\ No newline at end of file
diff --git a/memory-store/migrations/000016_entry_relations.down.sql b/memory-store/migrations/000016_entry_relations.down.sql
new file mode 100644
index 000000000..6d54b0c08
--- /dev/null
+++ b/memory-store/migrations/000016_entry_relations.down.sql
@@ -0,0 +1,12 @@
+BEGIN;
+
+-- Drop trigger first
+DROP TRIGGER IF EXISTS trg_enforce_leaf_nodes ON entry_relations;
+
+-- Drop function
+DROP FUNCTION IF EXISTS enforce_leaf_nodes ();
+
+-- Drop the table and its constraints
+DROP TABLE IF EXISTS entry_relations CASCADE;
+
+COMMIT;
\ No newline at end of file
diff --git a/memory-store/migrations/000016_entry_relations.up.sql b/memory-store/migrations/000016_entry_relations.up.sql
new file mode 100644
index 000000000..c61c7cd24
--- /dev/null
+++ b/memory-store/migrations/000016_entry_relations.up.sql
@@ -0,0 +1,55 @@
+BEGIN;
+
+-- Create citext extension if not exists
+CREATE EXTENSION IF NOT EXISTS citext;
+
+-- Create entry_relations table
+CREATE TABLE IF NOT EXISTS entry_relations (
+    session_id UUID NOT NULL,
+    head UUID NOT NULL,
+    relation CITEXT NOT NULL,
+    tail UUID NOT NULL,
+    is_leaf BOOLEAN NOT NULL DEFAULT FALSE,
+    CONSTRAINT pk_entry_relations PRIMARY KEY (session_id, head, relation, tail)
+);
+
+-- Add foreign key constraint to sessions table
+DO $$
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1 FROM pg_constraint WHERE conname = 'fk_entry_relations_session'
+    ) THEN
+        ALTER TABLE entry_relations
+        ADD CONSTRAINT fk_entry_relations_session
+        FOREIGN KEY (session_id)
+        REFERENCES sessions(session_id);
+    END IF;
+END $$;
+
+-- Create indexes for efficient querying
+CREATE INDEX idx_entry_relations_components ON entry_relations (session_id, head, relation, tail);
+
+CREATE INDEX idx_entry_relations_leaf ON entry_relations (session_id, relation, is_leaf);
+
+CREATE
+OR REPLACE FUNCTION enforce_leaf_nodes () RETURNS TRIGGER AS $$
+BEGIN
+    IF NEW.is_leaf THEN
+        -- Ensure no other relations point to this leaf node as a head
+        IF EXISTS (
+            SELECT 1 FROM entry_relations 
+            WHERE tail = NEW.head AND session_id = NEW.session_id
+        ) THEN
+            RAISE EXCEPTION 'Cannot assign relations to a leaf node.';
+        END IF;
+    END IF;
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE TRIGGER trg_enforce_leaf_nodes BEFORE INSERT
+OR
+UPDATE ON entry_relations FOR EACH ROW
+EXECUTE FUNCTION enforce_leaf_nodes ();
+
+COMMIT;
\ No newline at end of file