Module: RubyLLM::Protocols::Cohere::Chat

Defined in:
lib/ruby_llm/protocols/cohere/chat.rb

Overview

Chat methods for the Cohere v2 API implementation

Constant Summary collapse

FINISH_REASONS = 
{
  'COMPLETE' => :stop, 'STOP_SEQUENCE' => :stop, 'MAX_TOKENS' => :max_tokens, 'TOOL_CALL' => :tool_calls
}.freeze

Class Method Summary collapse

Class Method Details

.add_optional_fields(payload, messages, temperature:, max_output_tokens:, citations:, schema:) ⇒ Object 



43
44
45
46
47
48
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 43

def add_optional_fields(payload, messages, temperature:, max_output_tokens:, citations:, schema:)
  payload[:temperature] = temperature unless temperature.nil?
  payload[:max_tokens] = max_output_tokens unless max_output_tokens.nil?
  payload[:documents] = Media.format_documents(messages) if citations && Media.documents?(messages)
  payload[:response_format] = build_response_format(schema) if schema
end

.add_thinking(payload, thinking) ⇒ Object

Reasoning is on by default for models that support it, so an explicit disable is as meaningful as an explicit enable.



77
78
79
80
81
82
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 77

def add_thinking(payload, thinking)
  return unless thinking&.enabled?
  return payload[:thinking] = { type: 'disabled' } if thinking.disabled?

  payload[:thinking] = { type: 'enabled', token_budget: thinking.budget }.compact
end

.add_tools(payload, tools, tool_prefs) ⇒ Object 



57
58
59
60
61
62
63
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 57

def add_tools(payload, tools, tool_prefs)
  return if tools.empty?

  payload[:tools] = tools.values.map { |tool| Tools.function_for(tool) }
  tool_choice = Tools.build_tool_choice(tool_prefs[:choice])
  payload[:tool_choice] = tool_choice if tool_choice
end

.build_response_format(schema) ⇒ Object

Cohere takes a bare JSON Schema under json_schema, with no name or strict wrapper.



67
68
69
70
71
72
73
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 67

def build_response_format(schema)
  normalized = RubyLLM::Support::Utils.deep_dup(schema[:schema])
  normalized.delete(:strict)
  normalized.delete('strict')

  { type: 'json_object', json_schema: normalized }
end

.citation_span(data, offsets) ⇒ Object 



190
191
192
193
194
195
196
197
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 190

def citation_span(data, offsets)
  return [nil, nil] unless text_citation?(data)

  offset = offsets[data['content_index'] || 0]
  return [nil, nil] unless offset

  [data['start'] && (offset + data['start']), data['end'] && (offset + data['end'])]
end

.citation_url(document) ⇒ Object 



210
211
212
213
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 210

def citation_url(document)
  url = document['url']
  url if url.is_a?(String) && url.match?(%r{\Ahttps?://}i)
end

.completion_url ⇒ Object 



22
23
24
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 22

def completion_url
  'v2/chat'
end

.extract_text(blocks) ⇒ Object

Returns the joined text of the response along with the offset each content block starts at, so citation spans resolve against the content string RubyLLM exposes.



145
146
147
148
149
150
151
152
153
154
155
156
157
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 145

def extract_text(blocks)
  text = +''
  offsets = {}

  blocks.each_with_index do |block, index|
    next unless block['type'] == 'text'

    offsets[index] = text.length
    text << block['text'].to_s
  end

  [text, offsets]
end

.extract_thinking(blocks, message_data) ⇒ Object

The tool plan is the model's reasoning for a tool-calling turn, and is the only reasoning Cohere returns when thinking blocks are absent.



161
162
163
164
165
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 161

def extract_thinking(blocks, message_data)
  thoughts = blocks.select { |block| block['type'] == 'thinking' }
                   .map { |block| block['thinking'] }.join
  thoughts.empty? ? message_data['tool_plan'] : thoughts
end

.finish_reasons ⇒ Object 



14
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 14

def finish_reasons = FINISH_REASONS

.format_assistant_message(msg) ⇒ Object

Cohere returns tool_plan alongside tool calls, and RubyLLM surfaces it as thinking, but the newer models reject it on the way back in. It is optional in a request, so the plan stays out of the history.



101
102
103
104
105
106
107
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 101

def format_assistant_message(msg)
  message = { role: 'assistant' }
  content = Media.format_content(msg.content, msg.attachments)
  message[:content] = content unless content.empty?
  message[:tool_calls] = Tools.format_tool_calls(msg.tool_calls) if msg.tool_call?
  message
end

.format_message(msg, citations: false) ⇒ Object 



88
89
90
91
92
93
94
95
96
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 88

def format_message(msg, citations: false)
  return Tools.format_tool_result(msg) if msg.tool_result?
  return format_assistant_message(msg) if msg.role == :assistant

  {
    role: msg.role.to_s,
    content: Media.format_content(msg.content, msg.attachments, citations: citations)
  }
end

.format_messages(messages, citations: false) ⇒ Object 



84
85
86
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 84

def format_messages(messages, citations: false)
  messages.map { |msg| format_message(msg, citations: citations) }
end

.normalize_finish_reason(reason) ⇒ Object 



16
17
18
19
20
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 16

def normalize_finish_reason(reason)
  return nil if reason.nil?

  finish_reasons.fetch(reason.to_s) { reason.to_s.to_sym }
end

.parse_citation(data, offsets = {}) ⇒ Object

Only citations of the response text carry offsets into #content. Citations of the thinking blocks or the tool plan point into text RubyLLM exposes elsewhere, so they keep their snippet without a span.



174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 174

def parse_citation(data, offsets = {})
  source = Array(data['sources']).first
  document = source_document(source)
  start_index, end_index = citation_span(data, offsets)

  Citation.new(
    url: citation_url(document),
    title: document['title'] || document['id'],
    cited_text: document['text'] || document['snippet'],
    text: data['text'],
    start_index: start_index,
    end_index: end_index,
    source_index: source_index(source)
  )
end

.parse_citations(citations, offsets) ⇒ Object 



167
168
169
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 167

def parse_citations(citations, offsets)
  Array(citations).map { |citation| parse_citation(citation, offsets) }
end

.parse_completion_body(data, raw:) ⇒ Object 



109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 109

def parse_completion_body(data, raw:)
  message_data = data['message'] || {}
  blocks = Array(message_data['content'])
  content, offsets = extract_text(blocks)
  finish_reason = normalize_finish_reason(data['finish_reason'])

  Message.new(
    role: :assistant,
    content: content,
    citations: parse_citations(message_data['citations'], offsets),
    thinking: Thinking.build(text: extract_thinking(blocks, message_data)),
    tool_calls: Tools.parse_tool_calls(message_data['tool_calls'], response: raw, finish_reason:),
    finish_reason: finish_reason,
    model: model&.id,
    raw: raw,
    **usage_tokens(data['usage'] || {})
  )
end

.render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil, schema: nil, thinking: nil, citations: false, caching: nil, tool_prefs: nil) ⇒ Object

rubocop:disable-next Lint/UnusedMethodArgument



27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 27

def render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil,
                   schema: nil, thinking: nil, citations: false, caching: nil, tool_prefs: nil)
  warn_unsupported_citations(model) if citations && !model.supports?(:citations)

  payload = {
    model: model.id,
    messages: format_messages(messages, citations: citations),
    stream: stream
  }

  add_optional_fields(payload, messages, temperature:, max_output_tokens:, citations:, schema:)
  add_tools(payload, tools, tool_prefs || {})
  add_thinking(payload, thinking)
  payload
end

.source_document(source) ⇒ Object 



204
205
206
207
208
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 204

def source_document(source)
  return {} unless source

  source['document'] || source['tool_output'] || {}
end

.source_index(source) ⇒ Object

Documents RubyLLM sends, and those Cohere numbers itself, are identified as doc:N where N is the document's position.



217
218
219
220
221
222
223
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 217

def source_index(source)
  id = source && source['id']
  return unless id.is_a?(String)

  match = id.match(/\Adoc:(\d+)\z/)
  match && match[1].to_i
end

.text_citation?(data) ⇒ Boolean

Returns:



199
200
201
202
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 199

def text_citation?(data)
  type = data['type']
  type.nil? || type == 'TEXT_CONTENT'
end

.usage_tokens(usage) ⇒ Object

Cohere reports what the model processed under tokens and what it charges for under billed_units; the two differ because Cohere does not bill its own preamble.



131
132
133
134
135
136
137
138
139
140
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 131

def usage_tokens(usage)
  tokens = usage['tokens'] || {}
  billed = usage['billed_units'] || {}

  {
    input_tokens: tokens['input_tokens'] || billed['input_tokens'],
    output_tokens: tokens['output_tokens'] || billed['output_tokens'],
    cache_read_tokens: usage['cached_tokens']
  }
end

.warn_unsupported_citations(model) ⇒ Object 



50
51
52
53
54
55
 
# File 'lib/ruby_llm/protocols/cohere/chat.rb', line 50

def warn_unsupported_citations(model)
  RubyLLM.logger.warn(
    "#{model.id} does not support citations according to the model registry. " \
    'with_citations may have no effect.'
  )
end