Class: TurnKit::MemoryStore

Inherits:
Store
  • Object
show all
Defined in:
lib/turnkit/memory_store.rb

Instance Method Summary collapse

Constructor Details

#initialize ⇒ MemoryStore

Returns a new instance of MemoryStore.



5
6
7
8
9
10
11
12
 
# File 'lib/turnkit/memory_store.rb', line 5

def initialize
  @mutex = Mutex.new
  @conversations = {}
  @turns = {}
  @messages = {}
  @tool_executions = {}
  @message_sequences = Hash.new(0)
end

Instance Method Details

#append_message(attributes) ⇒ Object 



35
36
37
38
39
40
41
 
# File 'lib/turnkit/memory_store.rb', line 35

def append_message(attributes)
  attrs = stringify(attributes)
  attrs["sequence"] ||= next_message_sequence(attrs.fetch("conversation_id"))
  message = Record.message(attrs)
  @mutex.synchronize { @messages[message.fetch("id")] = message }
  duplicate(message)
end

#create_conversation(attributes) ⇒ Object 



14
15
16
17
18
19
 
# File 'lib/turnkit/memory_store.rb', line 14

def create_conversation(attributes)
  record = Record.conversation(attributes)

  @mutex.synchronize { @conversations[record.fetch("id")] = record }
  record.dup
end

#create_tool_execution(attributes) ⇒ Object 



80
81
82
83
84
85
 
# File 'lib/turnkit/memory_store.rb', line 80

def create_tool_execution(attributes)
  record = Record.tool_execution(attributes)

  @mutex.synchronize { @tool_executions[record.fetch("id")] = record }
  duplicate(record)
end

#create_turn(attributes) ⇒ Object 



51
52
53
54
55
56
 
# File 'lib/turnkit/memory_store.rb', line 51

def create_turn(attributes)
  record = Record.turn(attributes)

  @mutex.synchronize { @turns[record.fetch("id")] = record }
  duplicate(record)
end

#find_stale_turns(before:) ⇒ Object 



109
110
111
112
113
114
115
 
# File 'lib/turnkit/memory_store.rb', line 109

def find_stale_turns(before:)
  @mutex.synchronize do
    @turns.values.select do |turn|
      %w[pending running].include?(turn["status"]) && stale_anchor(turn) && stale_anchor(turn) < before
    end.map { |turn| duplicate(turn) }
  end
end

#latest_message_sequence(conversation_id) ⇒ Object 



31
32
33
 
# File 'lib/turnkit/memory_store.rb', line 31

def latest_message_sequence(conversation_id)
  @mutex.synchronize { @message_sequences[conversation_id].to_i }
end

#list_messages(conversation_id, through_sequence: nil, turn_id: nil) ⇒ Object 



43
44
45
46
47
48
49
 
# File 'lib/turnkit/memory_store.rb', line 43

def list_messages(conversation_id, through_sequence: nil, turn_id: nil)
  @mutex.synchronize do
    rows = @messages.values.select { |message| message["conversation_id"] == conversation_id }
    rows = rows.select { |message| message["sequence"].to_i <= through_sequence.to_i || message["turn_id"] == turn_id } if through_sequence
    rows.sort_by { |message| [ message["sequence"].to_i, message["created_at"].to_f, message["id"] ] }.map { |message| duplicate(message) }
  end
end

#list_tool_executions(turn_id:) ⇒ Object 



100
101
102
103
104
105
106
107
 
# File 'lib/turnkit/memory_store.rb', line 100

def list_tool_executions(turn_id:)
  @mutex.synchronize do
    @tool_executions.values
      .select { |execution| execution["turn_id"] == turn_id }
      .sort_by { |execution| [ execution["created_at"].to_f, execution["id"] ] }
      .map { |execution| duplicate(execution) }
  end
end

#list_turns(root_turn_id: nil, conversation_id: nil) ⇒ Object 



71
72
73
74
75
76
77
78
 
# File 'lib/turnkit/memory_store.rb', line 71

def list_turns(root_turn_id: nil, conversation_id: nil)
  @mutex.synchronize do
    rows = @turns.values
    rows = rows.select { |turn| turn["root_turn_id"] == root_turn_id } if root_turn_id
    rows = rows.select { |turn| turn["conversation_id"] == conversation_id } if conversation_id
    rows.sort_by { |turn| [ turn["created_at"].to_f, turn["id"] ] }.map { |turn| duplicate(turn) }
  end
end

#load_conversation(id) ⇒ Object 



21
22
23
 
# File 'lib/turnkit/memory_store.rb', line 21

def load_conversation(id)
  @mutex.synchronize { duplicate(@conversations.fetch(id)) }
end

#load_tool_execution(id) ⇒ Object 



87
88
89
 
# File 'lib/turnkit/memory_store.rb', line 87

def load_tool_execution(id)
  @mutex.synchronize { duplicate(@tool_executions.fetch(id)) }
end

#load_turn(id) ⇒ Object 



58
59
60
 
# File 'lib/turnkit/memory_store.rb', line 58

def load_turn(id)
  @mutex.synchronize { duplicate(@turns.fetch(id)) }
end

#next_message_sequence(conversation_id) ⇒ Object 



25
26
27
28
29
 
# File 'lib/turnkit/memory_store.rb', line 25

def next_message_sequence(conversation_id)
  @mutex.synchronize do
    @message_sequences[conversation_id] += 1
  end
end

#update_tool_execution(id, attributes) ⇒ Object 



91
92
93
94
95
96
97
98
 
# File 'lib/turnkit/memory_store.rb', line 91

def update_tool_execution(id, attributes)
  attrs = Record.tool_execution_update(attributes)
  @mutex.synchronize do
    record = @tool_executions.fetch(id)
    record.merge!(attrs.merge("updated_at" => Clock.now))
    duplicate(record)
  end
end

#update_turn(id, attributes) ⇒ Object 



62
63
64
65
66
67
68
69
 
# File 'lib/turnkit/memory_store.rb', line 62

def update_turn(id, attributes)
  attrs = Record.turn_update(attributes)
  @mutex.synchronize do
    record = @turns.fetch(id)
    record.merge!(attrs.merge("updated_at" => Clock.now))
    duplicate(record)
  end
end