diff --git a/lib/llm/message.rb b/lib/llm/message.rb index 8b41cd1b..1668ecd8 100644 --- a/lib/llm/message.rb +++ b/lib/llm/message.rb @@ -4,9 +4,14 @@ module LLM class Message attr_accessor :role, :content - def initialize(role, content) + def initialize(role, content, extra = {}) @role = role @content = content + @extra = extra + end + + def logprobs + @extra[:logprobs] end def to_h diff --git a/lib/llm/providers/openai/response_parser.rb b/lib/llm/providers/openai/response_parser.rb index 79be36bb..0f6952e4 100644 --- a/lib/llm/providers/openai/response_parser.rb +++ b/lib/llm/providers/openai/response_parser.rb @@ -19,7 +19,7 @@ def parse_completion(raw) { model: raw["model"], choices: raw["choices"].map do - LLM::Message.new(*_1["message"].values_at("role", "content")) + LLM::Message.new(*_1["message"].values_at("role", "content"), {completion: self, logprobs: _1["logprobs"]}) end, prompt_tokens: raw.dig("usage", "prompt_tokens"), completion_tokens: raw.dig("usage", "completion_tokens"),