ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/providers/zai/chat.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class ZAI
+      # Chat methods of the Z.ai API integration
+      # https://docs.z.ai/api-reference/llm/chat-completion
+      module Chat
+        FINISH_REASONS = Protocols::ChatCompletions::Chat::FINISH_REASONS.merge(
+          'sensitive' => :content_filter,
+          'model_context_window_exceeded' => :max_tokens
+        ).freeze
+
+        # Search results come back only when the tool sets search_result.
+        FILE_MODEL_PREFIX = 'glm-5.3-flash'
+
+        SERVER_TOOL_ALIASES = {
+          web_search: lambda { |options|
+            settings = { enable: true, search_result: true }.merge(Support::Utils.deep_symbolize_keys(options))
+            { tool: { type: 'web_search', web_search: settings } }
+          }
+        }.freeze
+
+        # The GLM-5.3-Flash series is the one that reads files.
+        def self.reads_files?(model)
+          model&.id.to_s.start_with?(FILE_MODEL_PREFIX)
+        end
+
+        module_function
+
+        def finish_reasons = FINISH_REASONS
+
+        def render_role(role)
+          role.to_s
+        end
+
+        def render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil,
+                           schema: nil, thinking: nil, citations: false, caching: nil, tool_prefs: nil)
+          payload = super
+          raise ArgumentError, 'Z.ai does not accept calls: :one' if payload.delete(:parallel_tool_calls) == false
+
+          payload[:thinking] = { type: thinking.disabled? ? 'disabled' : 'enabled' } if thinking&.enabled?
+          payload.delete(:reasoning_effort) if thinking&.disabled?
+          payload
+        end
+
+        def apply_end_user(payload, identifier)
+          payload.merge(user_id: identifier)
+        end
+
+        # Sources from the web_search tool arrive in a top-level web_search
+        # array, on the last chunk when streaming.
+        def parse_root_citations(data)
+          results = data['web_search']
+          return super unless results.is_a?(Array) && results.any?
+
+          results.each_with_index.filter_map do |result, index|
+            next unless result.is_a?(Hash)
+
+            Citation.new(url: result['link'], title: result['title'], cited_text: result['content'],
+                         source_index: index)
+          end
+        end
+
+        # Z.ai caches prompt prefixes by itself and takes no cache controls.
+        def openai_prompt_caching?
+          false
+        end
+
+        def render_thinking(msg)
+          text = msg.thinking&.text if msg.role == :assistant
+          text ? { reasoning_content: text } : {}
+        end
+
+        def supported_attachment?(attachment)
+          return true if remote_pdf?(attachment) || inline_video?(attachment)
+
+          Protocols::ChatCompletions::Media.supported_attachment?(
+            attachment, document_attachments: document_attachments(attachment), audio_attachments: false
+          )
+        end
+
+        def render_content(content, attachments = [])
+          Protocols::ChatCompletions::Media.render_parts(content, attachments) do |attachment|
+            next { type: 'file', file: { file_url: attachment.source.to_s } } if remote_pdf?(attachment)
+            next render_video(attachment) if inline_video?(attachment)
+
+            Protocols::ChatCompletions::Media.render_attachment(
+              attachment, document_attachments: document_attachments(attachment),
+                          image_attachments: true, audio_attachments: false
+            )
+          end
+        end
+
+        def inline_video?(attachment)
+          attachment.video? && !attachment.provider_file?
+        end
+
+        def render_video(video)
+          { type: 'video_url', video_url: { url: video.url_or_data_uri } }
+        end
+
+        # The API fetches a remote PDF itself.
+        def remote_pdf?(attachment)
+          attachment.url? && attachment.pdf? && Chat.reads_files?(@model)
+        end
+
+        # Files uploaded to Z.ai are referenced by id; other documents are not inlined.
+        def document_attachments(attachment)
+          return :none unless attachment.provider_file? && attachment.source.provider == 'zai'
+
+          attachment.audio? ? :none : :pdf
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/zai/ocr.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class ZAI
+      # Layout parsing with GLM-OCR. Files go up as URLs or data URIs, and
+      # each page comes back as a list of layout elements.
+      module OCR
+        def ocr_url
+          'layout_parsing'
+        end
+
+        def render_ocr_payload(file, model:, pages: nil, provider_options: {})
+          attachment = file.is_a?(Attachment) ? file : Attachment.new(file, config: @config)
+
+          payload = { model:, file: attachment.url_or_data_uri }.merge(ocr_page_range(pages))
+                                                                .merge(provider_options.transform_keys(&:to_sym))
+          @ocr_first_page = payload[:start_page_id].to_i.pred.clamp(0..)
+          payload
+        end
+
+        def parse_ocr_response(response, model:)
+          data = response.body
+          RubyLLM::OCR.new(pages: parse_ocr_pages(data), model:, usage: data['usage'], raw: data)
+        end
+
+        private
+
+        def ocr_page_range(pages)
+          return {} if pages.nil?
+
+          first, last = pages.minmax
+          continuous = first && pages.sort.each_cons(2).all? { |page, following| following == page + 1 }
+          raise ArgumentError, 'Z.ai reads one continuous page range' unless continuous
+
+          { start_page_id: first + 1, end_page_id: last + 1 }
+        end
+
+        def parse_ocr_pages(data)
+          details = data['layout_details']
+          if details.nil? || details.empty?
+            return [RubyLLM::OCR::Page.new(index: @ocr_first_page.to_i, markdown: data['md_results'], raw: data)]
+          end
+
+          details.each_with_index.map { |elements, position| parse_ocr_page(elements, position) }
+        end
+
+        def ocr_element_markdown(element)
+          content = element['content']
+          content && element['label'] == 'image' ? "![](#{content})" : content
+        end
+
+        def parse_ocr_page(elements, position)
+          RubyLLM::OCR::Page.new(
+            index: @ocr_first_page.to_i + position,
+            markdown: elements.filter_map { |element| ocr_element_markdown(element) }.join("\n\n"),
+            images: elements.select { |element| element['label'] == 'image' },
+            tables: elements.select { |element| element['label'] == 'table' },
+            raw: { 'layout_details' => elements }
+          )
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/zai/token_counting.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class ZAI
+      # Counts a request through Z.ai's tokenizer, which takes the chat
+      # messages and tools and leaves response format and thinking out.
+      module TokenCounting
+        COUNT_TOKENS_KEYS = %i[model messages tools].freeze
+
+        def count_tokens_url
+          'tokenizer'
+        end
+
+        def render_count_tokens_payload(messages, model:, **)
+          render_payload(messages, model: model, temperature: nil, **).slice(*COUNT_TOKENS_KEYS)
+        end
+
+        def parse_count_tokens_response(response)
+          response.body.fetch('usage').fetch('total_tokens')
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/zai/videos.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class ZAI
+      # CogVideoX jobs. Reference images are the first frame, or the first
+      # and last frames, as URLs or data URIs. Job states come from the
+      # async-result endpoint, and finished videos are hosted by Z.ai.
+      module Videos
+        MAX_VIDEO_IMAGES = 2
+        private_constant :MAX_VIDEO_IMAGES
+
+        def video_url
+          'videos/generations'
+        end
+
+        def render_video_payload(prompt, model:, with: [], provider_options: {})
+          payload = { model:, prompt: }.compact
+          payload[:image_url] = with.map(&:url_or_data_uri) if with.any?
+          payload.merge(provider_options)
+        end
+
+        def parse_video_job(response, model:, **request)
+          id = response.body['id']
+          raise Error.new('Z.ai did not return a video job id', response:) unless id
+
+          VideoJob.new(id:, protocol: self, model:, raw: response.body, **request)
+        end
+
+        def video_request_settings(request)
+          { duration: request[:duration], resolution: request[:size] }
+        end
+
+        def video_job_url(job)
+          "async-result/#{job.id}"
+        end
+
+        def parse_video_job_status(response, **)
+          body = response.body
+          case body['task_status']
+          when 'SUCCESS' then { status: :completed, raw: body }
+          when 'FAIL' then { status: :failed, raw: body, error: 'Z.ai could not generate the video' }
+          else { status: :pending, raw: body }
+          end
+        end
+
+        def download_video(job)
+          url = job.raw.dig('video_result', 0, 'url')
+          raise Error, 'Z.ai video job has no video' unless url
+
+          Video.new(url:, mime_type: 'video/mp4', model: job.raw['model'] || job.model, raw: job.raw)
+        end
+
+        private
+
+        def validate_animate_inputs!(with:)
+          raise Error, 'Z.ai takes at most a first and a last frame' if with.size > MAX_VIDEO_IMAGES
+
+          with.each do |attachment|
+            next if attachment.url?
+            raise UnsupportedAttachmentError, attachment.mime_type if attachment.provider_file? || !attachment.image?
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/zai.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    # Z.ai API integration for Zhipu's GLM models.
+    class ZAI < Provider
+      # Z.ai's dialect of the Chat Completions API.
+      class ChatCompletions < Protocols::ChatCompletions
+        include ZAI::Chat
+        include ZAI::OCR
+        include ZAI::TokenCounting
+        include ZAI::Videos
+
+        def server_tool_aliases
+          ZAI::Chat::SERVER_TOOL_ALIASES
+        end
+
+        private
+
+        def provider_file_references?
+          true
+        end
+
+        # Z.ai reads a PDF only through its Files API, so any inline PDF is
+        # uploaded, not just a large one.
+        def default_large_file_upload_threshold
+          0
+        end
+
+        # Only the GLM-5.3-Flash series accepts a file_id.
+        def provider_file_attachable?(attachment)
+          attachment.pdf? && ZAI::Chat.reads_files?(@model)
+        end
+      end
+
+      protocol :chat_completions, ChatCompletions
+      protocol :files, Protocols::ZAI::Files
+
+      def api_base
+        configured_api_base(:zai_api_base, 'https://api.z.ai/api/paas/v4', endpoints: Protocols::ChatCompletions::ENDPOINT_PATHS)
+      end
+
+      def headers
+        { 'Authorization' => "Bearer #{@config.zai_api_key}" }
+      end
+
+      def account_identity
+        account_digest(api_base, @config.zai_api_key)
+      end
+
+      class << self
+        def models_dev_provider_ids # :nodoc:
+          %w[zai]
+        end
+
+        def configuration_options
+          %i[zai_api_key zai_api_base]
+        end
+
+        def configuration_requirements
+          %i[zai_api_key]
+        end
+      end
+    end
+  end
+end