Class: Llmemory::ShortTerm::Stores::PostgresStore

Inherits:
Base
  • Object
show all
Includes:
Crypto::FieldHelpers
Defined in:
lib/llmemory/short_term/stores/postgres_store.rb

Instance Method Summary collapse

Constructor Details

#initialize(database_url: nil, cipher: nil) ⇒ PostgresStore

Returns a new instance of PostgresStore.



12
13
14
15
16
17
18
 
# File 'lib/llmemory/short_term/stores/postgres_store.rb', line 12

def initialize(database_url: nil, cipher: nil)
  @database_url = database_url || Llmemory.configuration.database_url
  @connection = nil
  @cipher = cipher || Llmemory.build_cipher
  @table_ready = false
  @conn_mutex = Mutex.new
end

Instance Method Details

#delete(user_id, session_id) ⇒ Object 



48
49
50
51
52
53
54
55
56
57
 
# File 'lib/llmemory/short_term/stores/postgres_store.rb', line 48

def delete(user_id, session_id)
  with_conn do |connection|
    ensure_table!(connection)
    connection.exec_params(
      "DELETE FROM llmemory_checkpoints WHERE user_id = $1 AND session_id = $2",
      [user_id, session_id]
    )
  end
  true
end

#list_sessions(user_id:) ⇒ Object 



94
95
96
97
98
99
100
101
102
103
 
# File 'lib/llmemory/short_term/stores/postgres_store.rb', line 94

def list_sessions(user_id:)
  with_conn do |connection|
    ensure_table!(connection)
    result = connection.exec_params(
      "SELECT session_id FROM llmemory_checkpoints WHERE user_id = $1",
      [user_id]
    )
    result.map { |r| r["session_id"] }
  end
end

#list_users ⇒ Object 



86
87
88
89
90
91
92
 
# File 'lib/llmemory/short_term/stores/postgres_store.rb', line 86

def list_users
  with_conn do |connection|
    ensure_table!(connection)
    result = connection.exec("SELECT DISTINCT user_id FROM llmemory_checkpoints")
    result.map { |r| r["user_id"] }
  end
end

#load(user_id, session_id) ⇒ Object 



36
37
38
39
40
41
42
43
44
45
46
 
# File 'lib/llmemory/short_term/stores/postgres_store.rb', line 36

def load(user_id, session_id)
  with_conn do |connection|
    ensure_table!(connection)
    result = connection.exec_params(
      "SELECT state FROM llmemory_checkpoints WHERE user_id = $1 AND session_id = $2",
      [user_id, session_id]
    )
    row = result.first
    row ? deserialize(row["state"]) : nil
  end
end

#save(user_id, session_id, state) ⇒ Object 



20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
 
# File 'lib/llmemory/short_term/stores/postgres_store.rb', line 20

def save(user_id, session_id, state)
  with_conn do |connection|
    ensure_table!(connection)
    connection.exec_params(
      <<~SQL,
        INSERT INTO llmemory_checkpoints (user_id, session_id, state, updated_at)
        VALUES ($1, $2, $3, $4)
        ON CONFLICT (user_id, session_id)
        DO UPDATE SET state = $3, updated_at = $4
      SQL
      [user_id, session_id, serialize(state), Time.now.utc.iso8601]
    )
  end
  true
end

#update(user_id, session_id, &block) ⇒ Object 



59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
 
# File 'lib/llmemory/short_term/stores/postgres_store.rb', line 59

def update(user_id, session_id, &block)
  with_conn do |connection|
    ensure_table!(connection)
    connection.transaction do
      result = connection.exec_params(
        "SELECT state FROM llmemory_checkpoints WHERE user_id = $1 AND session_id = $2 FOR UPDATE",
        [user_id, session_id]
      )
      row = result.first
      current = row ? deserialize(row["state"]) : nil
      new_state = yield(current)
      next new_state if new_state.nil?

      connection.exec_params(
        <<~SQL,
          INSERT INTO llmemory_checkpoints (user_id, session_id, state, updated_at)
          VALUES ($1, $2, $3, $4)
          ON CONFLICT (user_id, session_id)
          DO UPDATE SET state = $3, updated_at = $4
        SQL
        [user_id, session_id, serialize(new_state), Time.now.utc.iso8601]
      )
      new_state
    end
  end
end