Class: Raif::StreamingResponses::Bedrock

Inherits:
Object
  • Object
show all
Defined in:
app/models/raif/streaming_responses/bedrock.rb

Instance Method Summary collapse

Instance Method Details

#current_response ⇒ Object 



85
86
87
 
# File 'app/models/raif/streaming_responses/bedrock.rb', line 85

def current_response
  @response
end

#initialize_new_message ⇒ Object 



5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
 
# File 'app/models/raif/streaming_responses/bedrock.rb', line 5

def initialize_new_message
  # Initialize empty AWS response object
  @message = Aws::BedrockRuntime::Types::Message.new(
    role: "assistant",
    content: []
  )

  @output = Aws::BedrockRuntime::Types::ConverseOutput::Message.new(message: @message)

  @usage = Aws::BedrockRuntime::Types::TokenUsage.new(
    input_tokens: 0,
    output_tokens: 0,
    total_tokens: 0
  )

  @response = Aws::BedrockRuntime::Types::ConverseResponse.new(
    output: @output,
    usage: @usage,
    stop_reason: nil
  )
end

#process_streaming_event(event_type, event) ⇒ Object 



27
28
29
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
73
74
75
76
77
78
79
80
81
82
83
 
# File 'app/models/raif/streaming_responses/bedrock.rb', line 27

def process_streaming_event(event_type, event)
  delta = nil

  case event.event_type
  when :message_start
    initialize_new_message
  when :content_block_start
    index = event.content_block_index

    if event.start.is_a?(Aws::BedrockRuntime::Types::ContentBlockStart::ToolUse)
      tool_use = event.start.tool_use
      @message.content[index] = Aws::BedrockRuntime::Types::ContentBlock.new(
        tool_use: Aws::BedrockRuntime::Types::ToolUseBlock.new(
          tool_use_id: tool_use.tool_use_id,
          name: tool_use.name,
          input: ""
        )
      )
    else
      @message.content[index] = Aws::BedrockRuntime::Types::ContentBlock::Text.new(text: "")
    end
  when :content_block_delta
    index = event.content_block_index

    if event.delta.is_a?(Aws::BedrockRuntime::Types::ContentBlockDelta::Text)
      @message.content[index] ||= Aws::BedrockRuntime::Types::ContentBlock::Text.new(text: "")
      delta = event.delta.text
      @message.content[index].text += delta
    elsif event.delta.is_a?(Aws::BedrockRuntime::Types::ContentBlockDelta::ToolUse)
      tool_use = event.delta.tool_use
      @message.content[index] ||= Aws::BedrockRuntime::Types::ContentBlock.new
      @message.content[index].tool_use ||= Aws::BedrockRuntime::Types::ToolUseBlock.new(
        tool_use_id: tool_use.tool_use_id,
        name: tool_use.name,
        input: ""
      )

      @message.content[index].tool_use.input += event.delta.tool_use.input
    end
  when :content_block_stop
    content_block = @message.content[event.content_block_index]

    if content_block&.tool_use&.input.is_a?(String)
      begin
        content_block.tool_use.input = JSON.parse(content_block.tool_use.input)
      rescue JSON::ParserError
        # If parsing fails, leave as a string
      end
    end
  when :message_stop
    @response.stop_reason = event.stop_reason
  when :metadata
    @response.usage = event.usage if event.respond_to?(:usage)
  end

  [delta, @response.stop_reason]
end