llm.rb 0.4.1 → 0.5.0

data/lib/llm/providers/gemini/files.rb

@@ -19,7 +19,7 @@ class LLM::Gemini
   #
   #   llm = LLM.gemini(ENV["KEY"])
   #   bot = LLM::Chat.new(llm).lazy
-  #   file = llm.files.create file: LLM::File("/audio/haiku.mp3")
+  #   file = llm.files.create file: "/audio/haiku.mp3"
   #   bot.chat(file)
   #   bot.chat("Describe the audio file I sent to you")
   #   bot.chat("The audio file is the first message I sent to you.")
@@ -30,7 +30,7 @@ class LLM::Gemini
   #
   #   llm = LLM.gemini(ENV["KEY"])
   #   bot = LLM::Chat.new(llm).lazy
-  #   file = llm.files.create file: LLM::File("/audio/haiku.mp3")
+  #   file = llm.files.create file: "/audio/haiku.mp3"
   #   bot.chat(["Describe the audio file I sent to you", file])
   #   bot.messages.select(&:assistant?).each { print "[#{_1.role}]", _1.content, "\n" }
   class Files
@@ -71,13 +71,14 @@ class LLM::Gemini
     # Create a file
     # @example
     #   llm = LLM.gemini(ENV["KEY"])
-    #   res = llm.files.create file: LLM::File("/audio/haiku.mp3"),
+    #   res = llm.files.create file: "/audio/haiku.mp3"
     # @see https://ai.google.dev/gemini-api/docs/files Gemini docs
-    # @param [File] file The file
+    # @param [String, LLM::File] file The file
     # @param [Hash] params Other parameters (see Gemini docs)
     # @raise (see LLM::Provider#request)
     # @return [LLM::Response::File]
     def create(file:, **params)
+      file = LLM.File(file)
       req = Net::HTTP::Post.new(request_upload_url(file:), {})
       req["content-length"] = file.bytesize
       req["X-Goog-Upload-Offset"] = 0

data/lib/llm/providers/gemini/format/completion_format.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module LLM::Gemini::Format
+  ##
+  # @private
+  class CompletionFormat
+    ##
+    # @param [LLM::Message, Hash] message
+    #  The message to format
+    def initialize(message)
+      @message = message
+    end
+
+    ##
+    # Formats the message for the Gemini chat completions API
+    # @return [Hash]
+    def format
+      catch(:abort) do
+        if Hash === message
+          {role: message[:role], parts: format_content(message[:content])}
+        elsif message.tool_call?
+          {role: message.role, parts: message.extra[:original_tool_calls].map { {"functionCall" => _1} }}
+        else
+          {role: message.role, parts: format_content(message.content)}
+        end
+      end
+    end
+
+    def format_content(content)
+      case content
+      when Array
+        content.empty? ? throw(:abort, nil) : content.flat_map { format_content(_1) }
+      when LLM::Response::File
+        file = content
+        [{file_data: {mime_type: file.mime_type, file_uri: file.uri}}]
+      when LLM::File
+        file = content
+        [{inline_data: {mime_type: file.mime_type, data: file.to_b64}}]
+      when String
+        [{text: content}]
+      when LLM::Message
+        format_content(content.content)
+      when LLM::Function::Return
+        [{text: content.value}]
+      else
+        raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
+                                       "is not supported by the Gemini API"
+      end
+    end
+
+    def message = @message
+    def content = message.content
+  end
+end

data/lib/llm/providers/gemini/format.rb

@@ -4,45 +4,46 @@ class LLM::Gemini
   ##
   # @private
   module Format
+    require_relative "format/completion_format"
+
     ##
     # @param [Array<LLM::Message>] messages
     #  The messages to format
     # @return [Array<Hash>]
     def format(messages)
-      messages.map do
-        if Hash === _1
-          {role: _1[:role], parts: [format_content(_1[:content])]}
-        else
-          {role: _1.role, parts: [format_content(_1.content)]}
-        end
+      messages.filter_map do |message|
+        CompletionFormat.new(message).format
       end
     end
 
     private
 
     ##
-    # @param [String, Array, LLM::Response::File, LLM::File] content
-    #  The content to format
+    # @param [JSON::Schema] schema
+    #  The schema to format
     # @return [Hash]
-    #  The formatted content
-    def format_content(content)
-      case content
-      when Array
-        content.map { format_content(_1) }
-      when LLM::Response::File
-        file = content
-        {file_data: {mime_type: file.mime_type, file_uri: file.uri}}
-      when LLM::File
-        file = content
-        {inline_data: {mime_type: file.mime_type, data: file.to_b64}}
-      when String
-        {text: content}
-      when LLM::Message
-        format_content(content.content)
-      else
-        raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
-                                       "is not supported by the Gemini API"
-      end
+    def format_schema(schema)
+      return {} unless schema
+      {
+        "generationConfig" => {
+          "response_mime_type" => "application/json",
+          "response_schema" => schema
+        }
+      }
+    end
+
+    ##
+    # @param [Array<LLM::Function>] tools
+    #  The tools to format
+    # @return [Hash]
+    def format_tools(tools)
+      return {} unless tools
+      functions = tools.grep(LLM::Function)
+      {
+        "tools" => {
+          "functionDeclarations" => functions.map { _1.format(self) }
+        }
+      }
     end
   end
 end

data/lib/llm/providers/gemini/images.rb

@@ -58,16 +58,17 @@ class LLM::Gemini
     #   res = llm.images.edit image: LLM::File("cat.png"), prompt: "Add a hat to the cat"
     #   IO.copy_stream res.images[0], "hatoncat.png"
     # @see https://ai.google.dev/gemini-api/docs/image-generation Gemini docs
-    # @param [LLM::File] image The image to edit
+    # @param [String, LLM::File] image The image to edit
     # @param [String] prompt The prompt
     # @param [Hash] params Other parameters (see Gemini docs)
     # @raise (see LLM::Provider#request)
     # @note (see LLM::Gemini::Images#create)
     # @return [LLM::Response::Image]
     def edit(image:, prompt:, model: "gemini-2.0-flash-exp-image-generation", **params)
-      req  = Net::HTTP::Post.new("/v1beta/models/#{model}:generateContent?key=#{secret}", headers)
-      body = JSON.dump({
-        contents: [{parts: [{text: prompt}, format_content(image)]}],
+      req   = Net::HTTP::Post.new("/v1beta/models/#{model}:generateContent?key=#{secret}", headers)
+      image = LLM.File(image)
+      body  = JSON.dump({
+        contents: [{parts: [{text: prompt}, format.format_content(image)]}],
         generationConfig: {responseModalities: ["TEXT", "IMAGE"]}
       }.merge!(params)).b
       set_body_stream(req, StringIO.new(body))
@@ -84,6 +85,10 @@ class LLM::Gemini
 
     private
 
+    def format
+      @format ||= CompletionFormat.new(nil)
+    end
+
     def secret
       @provider.instance_variable_get(:@secret)
     end

data/lib/llm/providers/gemini/response_parser/completion_parser.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module LLM::Gemini::ResponseParser
+  class CompletionParser
+    def initialize(body)
+      @body = OpenStruct.from_hash(body)
+    end
+
+    def format(response)
+      {
+        model:,
+        prompt_tokens:,
+        completion_tokens:,
+        total_tokens:,
+        choices: format_choices(response)
+      }
+    end
+
+    private
+
+    def format_choices(response)
+      candidates.map.with_index do |choice, index|
+        content = choice.content
+        role, parts = content.role, content.parts
+        text  = parts.filter_map { _1["text"] }.join
+        tools = parts.filter_map { _1["functionCall"] }
+        extra = {index:, response:, tool_calls: format_tool_calls(tools), original_tool_calls: tools}
+        LLM::Message.new(role, text, extra)
+      end
+    end
+
+    def format_tool_calls(tools)
+      (tools || []).map do |tool|
+        function = {name: tool.name, arguments: tool.args}
+        OpenStruct.new(function)
+      end
+    end
+
+    def body = @body
+    def model = body.modelVersion
+    def prompt_tokens = body.usageMetadata.promptTokenCount
+    def completion_tokens = body.usageMetadata.candidatesTokenCount
+    def total_tokens = body.usageMetadata.totalTokenCount
+    def candidates = body.candidates
+  end
+end

data/lib/llm/providers/gemini/response_parser.rb

@@ -4,33 +4,24 @@ class LLM::Gemini
   ##
   # @private
   module ResponseParser
+    require_relative "response_parser/completion_parser"
+
     ##
     # @param [Hash] body
     #  The response body from the LLM provider
     # @return [Hash]
-    def parse_embedding(body)
-      {
-        model: "text-embedding-004",
-        embeddings: body.dig("embedding", "values")
-      }
+    def parse_completion(body)
+      CompletionParser.new(body).format(self)
     end
 
     ##
     # @param [Hash] body
     #  The response body from the LLM provider
     # @return [Hash]
-    def parse_completion(body)
+    def parse_embedding(body)
       {
-        model: body["modelVersion"],
-        choices: body["candidates"].map do
-          LLM::Message.new(
-            _1.dig("content", "role"),
-            _1.dig("content", "parts", 0, "text"),
-            {response: self}
-          )
-        end,
-        prompt_tokens: body.dig("usageMetadata", "promptTokenCount"),
-        completion_tokens: body.dig("usageMetadata", "candidatesTokenCount")
+        model: "text-embedding-004",
+        embeddings: body.dig("embedding", "values")
       }
     end
 
@@ -41,10 +32,12 @@ class LLM::Gemini
     def parse_image(body)
       {
         urls: [],
-        images: body["candidates"].flat_map do |candidate|
-          candidate["content"]["parts"].filter_map do
-            next unless _1.dig("inlineData", "data")
-            StringIO.new(_1["inlineData"]["data"].unpack1("m0"))
+        images: body["candidates"].flat_map do |c|
+          parts = c["content"]["parts"]
+          parts.filter_map do
+            data = _1.dig("inlineData", "data")
+            next unless data
+            StringIO.new(data.unpack1("m0"))
           end
         end
       }

data/lib/llm/providers/gemini.rb

@@ -74,12 +74,13 @@ module LLM
     # @raise [LLM::Error::PromptError]
     #  When given an object a provider does not understand
     # @return (see LLM::Provider#complete)
-    def complete(prompt, role = :user, model: default_model, schema: nil, **params)
+    def complete(prompt, role = :user, model: default_model, schema: nil, tools: nil, **params)
+      params = [format_schema(schema), format_tools(tools), params].inject({}, &:merge!).compact
       model.respond_to?(:id) ? model.id : model
       path = ["/v1beta/models/#{model}", "generateContent?key=#{@secret}"].join(":")
       req  = Net::HTTP::Post.new(path, headers)
       messages = [*(params.delete(:messages) || []), LLM::Message.new(role, prompt)]
-      body = JSON.dump({contents: format(messages)}.merge!(expand_schema(schema)))
+      body = JSON.dump({contents: format(messages)}.merge!(params))
       set_body_stream(req, StringIO.new(body))
       res = request(@http, req)
       Response::Completion.new(res).extend(response_parser)
@@ -136,16 +137,6 @@ module LLM
       }
     end
 
-    def expand_schema(schema)
-      return {} unless schema
-      {
-        "generationConfig" => {
-          "response_mime_type" => "application/json",
-          "response_schema" => schema
-        }
-      }
-    end
-
     def response_parser
       LLM::Gemini::ResponseParser
     end

data/lib/llm/providers/ollama/format/completion_format.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module LLM::Ollama::Format
+  ##
+  # @private
+  class CompletionFormat
+    ##
+    # @param [LLM::Message] message
+    #  The message to format
+    def initialize(message)
+      @message = message
+    end
+
+    ##
+    # Returns the message for the Ollama chat completions API
+    # @return [Hash]
+    def format
+      catch(:abort) do
+        if Hash === message
+          {role: message[:role]}.merge(format_content(message[:content]))
+        else
+          format_message
+        end
+      end
+    end
+
+    private
+
+    def format_content(content)
+      case content
+      when LLM::File
+        if content.image?
+          {content: "This message has an image associated with it", images: [content.to_b64]}
+        else
+          raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
+                                         "is not an image, and therefore not supported by the " \
+                                         "Ollama API"
+        end
+      when String
+        {content:}
+      when LLM::Message
+        format_content(content.content)
+      else
+        raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
+                                       "is not supported by the Ollama API"
+      end
+    end
+
+    def format_message
+      case content
+      when Array
+        format_array
+      else
+        {role: message.role}.merge(format_content(content))
+      end
+    end
+
+    def format_array
+      if content.empty?
+        nil
+      elsif returns.any?
+        returns.map { {role: "tool", tool_call_id: _1.id, content: JSON.dump(_1.value)} }
+      else
+        [{role: message.role, content: content.flat_map { format_content(_1) }}]
+      end
+    end
+
+    def message = @message
+    def content = message.content
+    def returns = content.grep(LLM::Function::Return)
+  end
+end

data/lib/llm/providers/ollama/format.rb

@@ -4,47 +4,27 @@ class LLM::Ollama
   ##
   # @private
   module Format
+    require_relative "format/completion_format"
+
     ##
     # @param [Array<LLM::Message>] messages
     #  The messages to format
     # @return [Array<Hash>]
     def format(messages)
-      messages.map do
-        if Hash === _1
-          {role: _1[:role]}
-            .merge!(_1)
-            .merge!(format_content(_1[:content]))
-        else
-          {role: _1.role}.merge! format_content(_1.content)
-        end
+      messages.filter_map do |message|
+        CompletionFormat.new(message).format
       end
     end
 
     private
 
     ##
-    # @param [String, URI] content
-    #  The content to format
-    # @return [String, Hash]
-    #  The formatted content
-    def format_content(content)
-      case content
-      when LLM::File
-        if content.image?
-          {content: "This message has an image associated with it", images: [content.to_b64]}
-        else
-          raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
-                                         "is not an image, and therefore not supported by the " \
-                                         "Ollama API"
-        end
-      when String
-        {content:}
-      when LLM::Message
-        format_content(content.content)
-      else
-        raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
-                                       "is not supported by the Ollama API"
-      end
+    # @param [Array<LLM::Function>] tools
+    #  The tools to format
+    # @return [Hash]
+    def format_tools(tools)
+      return {} unless tools
+      {tools: tools.map { _1.format(self) }}
     end
   end
 end

data/lib/llm/providers/ollama/response_parser/completion_parser.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module LLM::Ollama::ResponseParser
+  ##
+  # @private
+  class CompletionParser
+    def initialize(body)
+      @body = OpenStruct.from_hash(body)
+    end
+
+    def format(response)
+      {
+        model:,
+        choices: [format_choices(response)],
+        prompt_tokens:,
+        completion_tokens:
+      }
+    end
+
+    private
+
+    def format_choices(response)
+      role, content, calls = message.to_h.values_at(:role, :content, :tool_calls)
+      extra = {response:, tool_calls: format_tool_calls(calls)}
+      LLM::Message.new(role, content, extra)
+    end
+
+    def format_tool_calls(tools)
+      return [] unless tools
+      tools.filter_map do |tool|
+        next unless tool["function"]
+        OpenStruct.new(tool["function"])
+      end
+    end
+
+    def body = @body
+    def model = body.model
+    def prompt_tokens = body.prompt_eval_count
+    def completion_tokens = body.eval_count
+    def message = body.message
+  end
+end

data/lib/llm/providers/ollama/response_parser.rb

@@ -4,29 +4,26 @@ class LLM::Ollama
   ##
   # @private
   module ResponseParser
+    require_relative "response_parser/completion_parser"
+
     ##
     # @param [Hash] body
     #  The response body from the LLM provider
     # @return [Hash]
-    def parse_embedding(body)
-      {
-        model: body["model"],
-        embeddings: body["data"].map { _1["embedding"] },
-        prompt_tokens: body.dig("usage", "prompt_tokens"),
-        total_tokens: body.dig("usage", "total_tokens")
-      }
+    def parse_completion(body)
+      CompletionParser.new(body).format(self)
     end
 
     ##
     # @param [Hash] body
     #  The response body from the LLM provider
     # @return [Hash]
-    def parse_completion(body)
+    def parse_embedding(body)
       {
         model: body["model"],
-        choices: [LLM::Message.new(*body["message"].values_at("role", "content"), {response: self})],
-        prompt_tokens: body.dig("prompt_eval_count"),
-        completion_tokens: body.dig("eval_count")
+        embeddings: body["data"].map { _1["embedding"] },
+        prompt_tokens: body.dig("usage", "prompt_tokens"),
+        total_tokens: body.dig("usage", "total_tokens")
       }
     end
   end

data/lib/llm/providers/ollama.rb

@@ -60,14 +60,11 @@ module LLM
     # @raise [LLM::Error::PromptError]
     #  When given an object a provider does not understand
     # @return (see LLM::Provider#complete)
-    def complete(prompt, role = :user, model: default_model, schema: nil, **params)
-      params = {model:, stream: false}
-                 .merge!(expand_schema(schema))
-                 .merge!(params)
-                 .compact
+    def complete(prompt, role = :user, model: default_model, schema: nil, tools: nil, **params)
+      params = [{model:, stream: false, format: schema}, format_tools(tools), params].inject({}, &:merge!).compact
       req = Net::HTTP::Post.new("/api/chat", headers)
       messages = [*(params.delete(:messages) || []), LLM::Message.new(role, prompt)]
-      body = JSON.dump({messages: format(messages)}.merge!(params))
+      body = JSON.dump({messages: [format(messages)].flatten}.merge!(params))
       set_body_stream(req, StringIO.new(body))
       res = request(@http, req)
       Response::Completion.new(res).extend(response_parser)
@@ -104,11 +101,6 @@ module LLM
       }
     end
 
-    def expand_schema(schema)
-      return {} unless schema
-      {format: schema}
-    end
-
     def response_parser
       LLM::Ollama::ResponseParser
     end

data/lib/llm/providers/openai/audio.rb

@@ -7,7 +7,7 @@ class LLM::OpenAI
   # @example
   #   llm = LLM.openai(ENV["KEY"])
   #   res = llm.audio.create_speech(input: "A dog on a rocket to the moon")
-  #   File.binwrite("rocket.mp3", res.audio.string)
+  #   IO.copy_stream res.audio, "rocket.mp3"
   class Audio
     ##
     # Returns a new Audio object
@@ -43,16 +43,16 @@ class LLM::OpenAI
     # Create an audio transcription
     # @example
     #   llm = LLM.openai(ENV["KEY"])
-    #   res = llm.audio.create_transcription(file: LLM::File("/rocket.mp3"))
+    #   res = llm.audio.create_transcription(file: "/audio/rocket.mp3")
     #   res.text # => "A dog on a rocket to the moon"
     # @see https://platform.openai.com/docs/api-reference/audio/createTranscription OpenAI docs
-    # @param [LLM::File] file The input audio
+    # @param [String, LLM::File] file The input audio
     # @param [String] model The model to use
     # @param [Hash] params Other parameters (see OpenAI docs)
     # @raise (see LLM::Provider#request)
     # @return [LLM::Response::AudioTranscription]
     def create_transcription(file:, model: "whisper-1", **params)
-      multi = LLM::Multipart.new(params.merge!(file:, model:))
+      multi = LLM::Multipart.new(params.merge!(file: LLM.File(file), model:))
       req = Net::HTTP::Post.new("/v1/audio/transcriptions", headers)
       req["content-type"] = multi.content_type
       set_body_stream(req, multi.body)
@@ -65,7 +65,7 @@ class LLM::OpenAI
     # @example
     #   # Arabic => English
     #   llm = LLM.openai(ENV["KEY"])
-    #   res = llm.audio.create_translation(file: LLM::File("/bismillah.mp3"))
+    #   res = llm.audio.create_translation(file: "/audio/bismillah.mp3")
     #   res.text # => "In the name of Allah, the Beneficent, the Merciful."
     # @see https://platform.openai.com/docs/api-reference/audio/createTranslation OpenAI docs
     # @param [LLM::File] file The input audio
@@ -74,7 +74,7 @@ class LLM::OpenAI
     # @raise (see LLM::Provider#request)
     # @return [LLM::Response::AudioTranslation]
     def create_translation(file:, model: "whisper-1", **params)
-      multi = LLM::Multipart.new(params.merge!(file:, model:))
+      multi = LLM::Multipart.new(params.merge!(file: LLM.File(file), model:))
       req = Net::HTTP::Post.new("/v1/audio/translations", headers)
       req["content-type"] = multi.content_type
       set_body_stream(req, multi.body)

data/lib/llm/providers/openai/files.rb

@@ -14,7 +14,7 @@ class LLM::OpenAI
   #
   #   llm = LLM.openai(ENV["KEY"])
   #   bot = LLM::Chat.new(llm).lazy
-  #   file = llm.files.create file: LLM::File("/documents/freebsd.pdf")
+  #   file = llm.files.create file: "/documents/freebsd.pdf"
   #   bot.chat(file)
   #   bot.chat("Describe the document")
   #   bot.messages.select(&:assistant?).each { print "[#{_1.role}]", _1.content, "\n" }
@@ -24,7 +24,7 @@ class LLM::OpenAI
   #
   #   llm = LLM.openai(ENV["KEY"])
   #   bot = LLM::Chat.new(llm).lazy
-  #   file = llm.files.create file: LLM::File("/documents/openbsd.pdf")
+  #   file = llm.files.create file: "/documents/openbsd.pdf"
   #   bot.chat(["Describe the document I sent to you", file])
   #   bot.messages.select(&:assistant?).each { print "[#{_1.role}]", _1.content, "\n" }
   class Files
@@ -62,7 +62,7 @@ class LLM::OpenAI
     # Create a file
     # @example
     #   llm = LLM.openai(ENV["KEY"])
-    #   res = llm.files.create file: LLM::File("/documents/haiku.txt"),
+    #   res = llm.files.create file: "/documents/haiku.txt"
     # @see https://platform.openai.com/docs/api-reference/files/create OpenAI docs
     # @param [File] file The file
     # @param [String] purpose The purpose of the file (see OpenAI docs)