Class: TurnKit::ActiveRecordStore

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

Instance Method Summary collapse

Instance Method Details

#append_message(attributes) ⇒ Object 



25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
 
# File 'lib/turnkit/stores/active_record_store.rb', line 25

def append_message(attributes)
  attrs = attributes.transform_keys(&:to_s)
  sequence = nil
  message = nil
  record = conversation_class.transaction do
    conversation_class.lock.find_by!(uid: attrs.fetch("conversation_id"))
    sequence = message_class.where(conversation_uid: attrs.fetch("conversation_id")).maximum(:sequence).to_i + 1
    message = Record.message(attrs.merge("sequence" => sequence))
    message_class.create!(
      uid: message.fetch("id"),
      conversation_uid: message.fetch("conversation_id"),
      turn_uid: message["turn_id"],
      role: message.fetch("role"),
      kind: message.fetch("kind"),
      sequence: message.fetch("sequence"),
      content: message.fetch("content"),
      text: message.fetch("text"),
      tool_execution_uid: message["tool_execution_id"],
      provider_message_id: message["provider_message_id"],
      metadata: message.fetch("metadata")
    )
  end
  message_hash(record)
end

#create_conversation(attributes) ⇒ Object 



5
6
7
8
 
# File 'lib/turnkit/stores/active_record_store.rb', line 5

def create_conversation(attributes)
  record = conversation_class.create!(record_attributes(attributes, id_key: "uid"))
  conversation_hash(record)
end

#create_tool_execution(attributes) ⇒ Object 



99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
 
# File 'lib/turnkit/stores/active_record_store.rb', line 99

def create_tool_execution(attributes)
  attrs = Record.tool_execution(attributes)
  record = tool_execution_class.create!(
    uid: attrs.fetch("id"),
    turn_uid: attrs.fetch("turn_id"),
    tool_call_id: attrs.fetch("tool_call_id"),
    tool_name: attrs.fetch("tool_name"),
    status: attrs.fetch("status"),
    arguments: attrs["arguments"] || {},
    result: attrs["result"],
    error: attrs["error"],
    started_at: attrs["started_at"],
    completed_at: attrs["completed_at"]
  )
  tool_execution_hash(record)
end

#create_turn(attributes) ⇒ Object 



58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
 
# File 'lib/turnkit/stores/active_record_store.rb', line 58

def create_turn(attributes)
  attrs = Record.turn(attributes)
  record = turn_class.create!(
    uid: attrs.fetch("id"),
    conversation_uid: attrs.fetch("conversation_id"),
    agent_name: attrs["agent_name"],
    parent_turn_uid: attrs["parent_turn_id"],
    parent_tool_execution_uid: attrs["parent_tool_execution_id"],
    root_turn_uid: attrs.fetch("root_turn_id"),
    context_message_sequence: attrs["context_message_sequence"].to_i,
    status: attrs.fetch("status"),
    model: attrs["model"],
    options: attrs["options"] || {},
    usage: attrs["usage"] || {},
    cost: attrs["cost"],
    error: attrs["error"],
    output_text: attrs["output_text"],
    started_at: attrs["started_at"],
    heartbeat_at: attrs["heartbeat_at"],
    completed_at: attrs["completed_at"]
  )
  turn_hash(record)
end

#find_stale_turns(before:) ⇒ Object 



130
131
132
 
# File 'lib/turnkit/stores/active_record_store.rb', line 130

def find_stale_turns(before:)
  turn_class.where(status: %w[pending running]).where("COALESCE(heartbeat_at, started_at, created_at) < ?", before).map { |record| turn_hash(record) }
end

#latest_message_sequence(conversation_id) ⇒ Object 



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

def latest_message_sequence(conversation_id)
  message_class.where(conversation_uid: conversation_id).maximum(:sequence).to_i
end

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



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

def list_messages(conversation_id, through_sequence: nil, turn_id: nil)
  scope = message_class.where(conversation_uid: conversation_id)
  if through_sequence
    scope = scope.where("sequence <= ? OR turn_uid = ?", through_sequence, turn_id)
  end
  scope.order(:sequence, :created_at, :uid).map { |record| message_hash(record) }
end

#list_tool_executions(turn_id:) ⇒ Object 



126
127
128
 
# File 'lib/turnkit/stores/active_record_store.rb', line 126

def list_tool_executions(turn_id:)
  tool_execution_class.where(turn_uid: turn_id).order(:created_at, :uid).map { |record| tool_execution_hash(record) }
end

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



92
93
94
95
96
97
 
# File 'lib/turnkit/stores/active_record_store.rb', line 92

def list_turns(root_turn_id: nil, conversation_id: nil)
  scope = turn_class.all
  scope = scope.where(root_turn_uid: root_turn_id) if root_turn_id
  scope = scope.where(conversation_uid: conversation_id) if conversation_id
  scope.order(:created_at, :uid).map { |record| turn_hash(record) }
end

#load_conversation(id) ⇒ Object 



10
11
12
 
# File 'lib/turnkit/stores/active_record_store.rb', line 10

def load_conversation(id)
  conversation_hash(conversation_class.find_by!(uid: id))
end

#load_tool_execution(id) ⇒ Object 



116
117
118
 
# File 'lib/turnkit/stores/active_record_store.rb', line 116

def load_tool_execution(id)
  tool_execution_hash(tool_execution_class.find_by!(uid: id))
end

#load_turn(id) ⇒ Object 



82
83
84
 
# File 'lib/turnkit/stores/active_record_store.rb', line 82

def load_turn(id)
  turn_hash(turn_class.find_by!(uid: id))
end

#next_message_sequence(conversation_id) ⇒ Object 



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

def next_message_sequence(conversation_id)
  conversation_class.transaction do
    conversation = conversation_class.lock.find_by!(uid: conversation_id)
    message_class.where(conversation_uid: conversation.uid).maximum(:sequence).to_i + 1
  end
end

#update_tool_execution(id, attributes) ⇒ Object 



120
121
122
123
124
 
# File 'lib/turnkit/stores/active_record_store.rb', line 120

def update_tool_execution(id, attributes)
  record = tool_execution_class.find_by!(uid: id)
  record.update!(Record.tool_execution_update(attributes))
  tool_execution_hash(record)
end

#update_turn(id, attributes) ⇒ Object 



86
87
88
89
90
 
# File 'lib/turnkit/stores/active_record_store.rb', line 86

def update_turn(id, attributes)
  record = turn_class.find_by!(uid: id)
  record.update!(Record.turn_update(attributes))
  turn_hash(record)
end