Class: PromptObjects::Repositories::RunRepository

Inherits:
BaseRepository show all
Defined in:
lib/prompt_objects/repositories/run_repository.rb

Instance Method Summary collapse

Methods inherited from BaseRepository

#initialize

Constructor Details

This class inherits a constructor from PromptObjects::Repositories::BaseRepository

Instance Method Details

#active_children(parent_run_id:) ⇒ Object



155
156
157
158
159
160
161
162
163
# File 'lib/prompt_objects/repositories/run_repository.rb', line 155

def active_children(parent_run_id:)
  read do |db|
    db.execute(<<~SQL, [parent_run_id]).map { |row| parse(row) }
      SELECT * FROM runs
      WHERE parent_run_id = ? AND status IN ('queued', 'running', 'waiting')
      ORDER BY enqueued_at ASC
    SQL
  end
end

#active_count(workspace_session_id:, po_name: nil) ⇒ Object



140
141
142
143
144
145
146
147
148
149
150
151
152
153
# File 'lib/prompt_objects/repositories/run_repository.rb', line 140

def active_count(workspace_session_id:, po_name: nil)
  conditions = ["workspace_session_id = ?", "status IN ('queued', 'running', 'waiting')"]
  params = [workspace_session_id]
  if po_name
    conditions << "po_name = ?"
    params << po_name
  end
  read do |db|
    db.get_first_value(
      "SELECT COUNT(*) FROM runs WHERE #{conditions.join(' AND ')}",
      params
    ).to_i
  end
end

#active_for_source(workspace_session_id:, po_name:, source_type:, source_name:) ⇒ Object



116
117
118
119
120
121
122
123
124
125
126
127
# File 'lib/prompt_objects/repositories/run_repository.rb', line 116

def active_for_source(workspace_session_id:, po_name:, source_type:, source_name:)
  row = read do |db|
    db.get_first_row(<<~SQL, [workspace_session_id, po_name, source_type, source_name])
      SELECT * FROM runs
      WHERE workspace_session_id = ? AND po_name = ? AND source_type = ?
        AND source_name = ? AND status IN ('queued', 'running', 'waiting')
      ORDER BY turn_ordinal ASC
      LIMIT 1
    SQL
  end
  row && parse(row)
end

#descendant_count(root_run_id:) ⇒ Object



165
166
167
168
169
170
171
172
# File 'lib/prompt_objects/repositories/run_repository.rb', line 165

def descendant_count(root_run_id:)
  read do |db|
    db.get_first_value(
      "SELECT COUNT(*) FROM runs WHERE root_run_id = ? AND id != ?",
      [root_run_id, root_run_id]
    ).to_i
  end
end

#enqueue(workspace_session_id:, thread_id:, po_name:, input_content:, definition_version:, source_type:, source_name: nil, client_request_id: nil, provider: nil, model: nil, capability_policy_version: nil, parent_run_id: nil, priority: 0, metadata: {}, id: SecureRandom.uuid) ⇒ Object



6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# File 'lib/prompt_objects/repositories/run_repository.rb', line 6

def enqueue(workspace_session_id:, thread_id:, po_name:, input_content:,
            definition_version:, source_type:, source_name: nil,
            client_request_id: nil, provider: nil, model: nil,
            capability_policy_version: nil, parent_run_id: nil,
            priority: 0, metadata: {}, id: SecureRandom.uuid)
  enqueue_with_result(
    workspace_session_id: workspace_session_id,
    thread_id: thread_id,
    po_name: po_name,
    input_content: input_content,
    definition_version: definition_version,
    source_type: source_type,
    source_name: source_name,
    client_request_id: client_request_id,
    provider: provider,
    model: model,
    capability_policy_version: capability_policy_version,
    parent_run_id: parent_run_id,
    priority: priority,
    metadata: ,
    id: id
  ).first
end

#enqueue_with_result(workspace_session_id:, thread_id:, po_name:, input_content:, definition_version:, source_type:, source_name: nil, client_request_id: nil, provider: nil, model: nil, capability_policy_version: nil, parent_run_id: nil, priority: 0, metadata: {}, id: SecureRandom.uuid) ⇒ Object



30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
# File 'lib/prompt_objects/repositories/run_repository.rb', line 30

def enqueue_with_result(workspace_session_id:, thread_id:, po_name:, input_content:,
                        definition_version:, source_type:, source_name: nil,
                        client_request_id: nil, provider: nil, model: nil,
                        capability_policy_version: nil, parent_run_id: nil,
                        priority: 0, metadata: {}, id: SecureRandom.uuid)
  transaction do |db|
    if client_request_id
      existing = find_idempotent(db, workspace_session_id, source_type, client_request_id)
      next [parse(existing), false] if existing
    end

    turn_ordinal = db.get_first_value(
      "SELECT COALESCE(MAX(turn_ordinal), 0) + 1 FROM runs WHERE thread_id = ?",
      [thread_id]
    )
    root_run_id = parent_run_id ? root_for(db, parent_run_id) : id
    now = Time.now.utc.iso8601
    params = [
      id, client_request_id, workspace_session_id, thread_id, po_name,
      turn_ordinal, input_content, definition_version, provider, model,
      capability_policy_version, source_type, source_name, parent_run_id,
      root_run_id, priority, now, JSON.generate()
    ]
    db.execute(<<~SQL, params)
      INSERT INTO runs (
        id, client_request_id, workspace_session_id, thread_id, po_name,
        turn_ordinal, input_content, definition_version, provider, model,
        capability_policy_version, source_type, source_name, parent_run_id,
        root_run_id, status, priority, enqueued_at, metadata
      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'queued', ?, ?, ?)
    SQL
    [parse(db.get_first_row("SELECT * FROM runs WHERE id = ?", [id])), true]
  end
rescue SQLite3::ConstraintException
  existing = client_request_id && get_by_client_request(
    workspace_session_id: workspace_session_id,
    source_type: source_type,
    client_request_id: client_request_id
  )
  return [existing, false] if existing

  raise
end

#get(id) ⇒ Object



74
75
76
77
# File 'lib/prompt_objects/repositories/run_repository.rb', line 74

def get(id)
  row = read { |db| db.get_first_row("SELECT * FROM runs WHERE id = ?", [id]) }
  row && parse(row)
end

#get_by_client_request(workspace_session_id:, source_type:, client_request_id:) ⇒ Object



79
80
81
82
83
84
# File 'lib/prompt_objects/repositories/run_repository.rb', line 79

def get_by_client_request(workspace_session_id:, source_type:, client_request_id:)
  row = read do |db|
    find_idempotent(db, workspace_session_id, source_type, client_request_id)
  end
  row && parse(row)
end

#list(workspace_session_id:, thread_id: nil, statuses: nil, limit: 100) ⇒ Object



86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
# File 'lib/prompt_objects/repositories/run_repository.rb', line 86

def list(workspace_session_id:, thread_id: nil, statuses: nil, limit: 100)
  conditions = ["workspace_session_id = ?"]
  params = [workspace_session_id]
  if thread_id
    conditions << "thread_id = ?"
    params << thread_id
  end
  if statuses&.any?
    unknown = statuses.map(&:to_s) - Execution::Contracts::RUN_STATUSES
    raise ArgumentError, "unknown run status: #{unknown.first}" if unknown.any?

    conditions << "status IN (#{Array.new(statuses.size, '?').join(', ')})"
    params.concat(statuses.map(&:to_s))
  end
  params << limit
  sql = "SELECT * FROM runs WHERE #{conditions.join(' AND ')} ORDER BY enqueued_at DESC LIMIT ?"
  read { |db| db.execute(sql, params) }.map { |row| parse(row) }
end

#thread_head?(run) ⇒ Boolean

Returns:

  • (Boolean)


129
130
131
132
133
134
135
136
137
138
# File 'lib/prompt_objects/repositories/run_repository.rb', line 129

def thread_head?(run)
  read do |db|
    count = db.get_first_value(<<~SQL, [run.thread_id, run.turn_ordinal, run.id])
      SELECT COUNT(*) FROM runs
      WHERE thread_id = ? AND turn_ordinal < ? AND id != ?
        AND status IN ('queued', 'running', 'waiting')
    SQL
    count.to_i.zero?
  end
end

#transition(id, status:, final_response: nil, cancellation_reason: nil, usage: nil, error: nil, history_cutoff_message_id: nil, input_message_id: nil) ⇒ Object



174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
# File 'lib/prompt_objects/repositories/run_repository.rb', line 174

def transition(id, status:, final_response: nil, cancellation_reason: nil,
               usage: nil, error: nil, history_cutoff_message_id: nil,
               input_message_id: nil)
  validate_status!(status)
  now = Time.now.utc.iso8601
  transaction do |db|
    current = db.get_first_row("SELECT * FROM runs WHERE id = ?", [id])
    next nil unless current

    started_at = status.to_s == "running" ? (current["started_at"] || now) : current["started_at"]
    finished_at = terminal_status?(status) ? now : current["finished_at"]
    params = [
      status.to_s, started_at, finished_at, final_response,
      cancellation_reason, usage && JSON.generate(usage),
      error&.class&.name, error&.message, history_cutoff_message_id,
      input_message_id, id
    ]
    db.execute(<<~SQL, params)
      UPDATE runs SET
        status = ?, started_at = ?, finished_at = ?, final_response = ?,
        cancellation_reason = ?, usage = ?, error_class = ?, error_message = ?,
        history_cutoff_message_id = COALESCE(?, history_cutoff_message_id),
        input_message_id = COALESCE(?, input_message_id)
      WHERE id = ?
    SQL
    db.execute(
      "UPDATE sessions SET last_run_status = ?, last_activity_at = ?, updated_at = ? WHERE id = ?",
      [status.to_s, now, now, current["thread_id"]]
    )
    parse(db.get_first_row("SELECT * FROM runs WHERE id = ?", [id]))
  end
end

#tree(root_run_id:) ⇒ Object



105
106
107
108
109
110
111
112
113
114
# File 'lib/prompt_objects/repositories/run_repository.rb', line 105

def tree(root_run_id:)
  read do |db|
    db.execute(<<~SQL, [root_run_id, root_run_id, root_run_id]).map { |row| parse(row) }
      SELECT * FROM runs
      WHERE id = ? OR root_run_id = ?
      ORDER BY CASE WHEN id = ? THEN 0 ELSE 1 END,
        enqueued_at ASC, turn_ordinal ASC, id ASC
    SQL
  end
end