crystil 0.5.0

data/lib/crystil/wrappers/geminiai.rb

@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+require_relative "constants"
+
+module Crystil
+  module Wrappers
+    # Wrapper for Google GenerativeAI Ruby client (gemini-ai gem version 4.3.0)
+    class GeminiAI
+      def initialize(config, collector, sentinel = nil)
+        @config = config
+        @collector = collector
+        @sentinel = sentinel
+      end
+
+      def register(client)
+        validate_client!(client)
+
+        # Prevent double registration
+        return client if client.instance_variable_defined?(:@crystil_registered)
+
+        # Store references in client instance
+        client.instance_variable_set(:@crystil_config, @config)
+        client.instance_variable_set(:@crystil_collector, @collector)
+        client.instance_variable_set(:@crystil_sentinel, @sentinel)
+        client.instance_variable_set(:@crystil_registered, true)
+
+        # Wrap the generate_content and stream_generate_content methods
+        wrap_generate_content_method(client)
+        wrap_stream_generate_content_method(client)
+
+        client
+      end
+
+      private
+
+      def validate_client!(client)
+        return if client.respond_to?(:generate_content) && client.respond_to?(:stream_generate_content)
+
+        raise RegistrationError,
+              "Client does not appear to be a valid Gemini AI client " \
+              "(missing generate_content or stream_generate_content method)"
+      end
+
+      def wrap_generate_content_method(client)
+        client.singleton_class.class_eval do
+          include Base
+
+          alias_method :original_generate_content, :generate_content
+
+          define_method(:generate_content) do |*args, **kwargs, &block|
+            # Handle both positional hash and keyword arguments
+            parameters = if args.first.is_a?(Hash)
+                           args.first
+                         elsif kwargs.any?
+                           kwargs
+                         else
+                           {}
+                         end
+
+            start_time = Time.now
+            version = defined?(::Gemini::GEM) && ::Gemini::GEM.is_a?(Hash) ? ::Gemini::GEM[:version] : nil
+
+            sentinel = instance_variable_get(:@crystil_sentinel)
+            sentinel&.raise_if_irrelevant!(
+              title: GOOGLE_CLIENT_TITLE,
+              request: parameters,
+              version: version
+            )
+
+            # Call original method
+            response = original_generate_content(*args, **kwargs, &block)
+
+            # Submit analytics
+            crystil_submit_analytics(
+              method: :generate_content,
+              args: [],
+              kwargs: parameters,
+              response: response,
+              start_time: start_time,
+              end_time: Time.now,
+              title: GOOGLE_CLIENT_TITLE,
+              version: version
+            )
+
+            response
+          rescue CrystilRequestInterceptedError => e
+            # We don't want to send intercepts to collector
+            raise e
+          rescue StandardError => e
+            crystil_submit_error_analytics(
+              method: :generate_content,
+              args: [],
+              kwargs: parameters,
+              error: e,
+              start_time: start_time,
+              end_time: Time.now,
+              title: GOOGLE_CLIENT_TITLE,
+              version: version
+            )
+
+            raise e
+          end
+        end
+      end
+
+      # Handles the stream_generate_content from the client.  Note, that
+      # the Crystil backend does not expect the payload to be merged like
+      # the OpenAI payload is.  So, all the results are just sent as an
+      # array.
+      def wrap_stream_generate_content_method(client)
+        client.singleton_class.class_eval do
+          include Base
+
+          alias_method :original_stream_generate_content, :stream_generate_content
+
+          define_method(:stream_generate_content) do |*args, **kwargs, &block|
+            # Handle both positional hash and keyword arguments
+            parameters = if args.first.is_a?(Hash)
+                           args.first
+                         elsif kwargs.any?
+                           kwargs
+                         else
+                           {}
+                         end
+
+            start_time = Time.now
+            version = defined?(::Gemini::GEM) && ::Gemini::GEM.is_a?(Hash) ? ::Gemini::GEM[:version] : nil
+
+            sentinel = instance_variable_get(:@crystil_sentinel)
+            sentinel&.raise_if_irrelevant!(
+              title: GOOGLE_CLIENT_TITLE,
+              request: parameters,
+              version: version
+            )
+
+            # Call original method with wrapped block
+            response = original_stream_generate_content(*args, **kwargs, &block)
+
+            # Submit analytics with accumulated response
+            crystil_submit_analytics(
+              method: :stream_generate_content,
+              args: [],
+              kwargs: parameters,
+              response: response,
+              start_time: start_time,
+              end_time: Time.now,
+              title: GOOGLE_CLIENT_TITLE,
+              version: version
+            )
+
+            response
+          rescue CrystilRequestInterceptedError => e
+            # We don't want to send intercepts to collector
+            raise e
+          rescue StandardError => e
+            crystil_submit_error_analytics(
+              method: :stream_generate_content,
+              args: [],
+              kwargs: parameters,
+              error: e,
+              start_time: start_time,
+              end_time: Time.now,
+              title: GOOGLE_CLIENT_TITLE,
+              version: version
+            )
+
+            raise e
+          end
+        end
+      end
+    end
+  end
+end

data/lib/crystil/wrappers/google.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require_relative "constants"
+
+module Crystil
+  module Wrappers
+    # Wrapper for Google GenerativeAI Ruby client (google-genai library version 0.1)
+    class Google
+      def initialize(config, collector, sentinel = nil)
+        @config = config
+        @collector = collector
+        @sentinel = sentinel
+      end
+
+      def register(client)
+        validate_client!(client)
+
+        # Prevent double registration
+        return client if client.instance_variable_defined?(:@crystil_registered)
+
+        # Patch response class after client is validated (ensures gem is loaded)
+        patch_response_class!
+
+        # Store references in client instance
+        client.instance_variable_set(:@crystil_config, @config)
+        client.instance_variable_set(:@crystil_collector, @collector)
+        client.instance_variable_set(:@crystil_sentinel, @sentinel)
+        client.instance_variable_set(:@crystil_registered, true)
+
+        # Wrap the generate_content method
+        wrap_generate_content_method(client)
+
+        client
+      end
+
+      private
+
+      def validate_client!(client)
+        return if client.respond_to?(:models)
+
+        raise RegistrationError,
+              "Client does not appear to be a valid Google GenAI client (missing models method)"
+      end
+
+      # Monkey-patch the GenerateContentResponse class to include missing fields
+      def patch_response_class!
+        # Check if the constant exists and hasn't been patched yet
+        return unless defined?(::Google::Genai::Types::GenerateContentResponse)
+        return if ::Google::Genai::Types::GenerateContentResponse.instance_variable_defined?(:@_crystil_patched)
+
+        ::Google::Genai::Types::GenerateContentResponse.class_eval do
+          # camelCase intentional — these names must match Google's API response
+          # keys verbatim so the patched accessors survive serialization through
+          # `Base#extract_response`.
+          attr_accessor :modelVersion, :usageMetadata, :responseId, :createTime # rubocop:disable Naming/MethodName
+        end
+
+        ::Google::Genai::Types::GenerateContentResponse.instance_variable_set(:@_crystil_patched, true)
+      end
+
+      def wrap_generate_content_method(client)
+        # Get the models resource
+        models_resource = client.models
+
+        # Store references on the models resource (needed for Base module)
+        models_resource.instance_variable_set(:@crystil_config, client.instance_variable_get(:@crystil_config))
+        models_resource.instance_variable_set(:@crystil_collector, client.instance_variable_get(:@crystil_collector))
+        models_resource.instance_variable_set(:@crystil_sentinel, client.instance_variable_get(:@crystil_sentinel))
+
+        # Store the original generate_content method
+        original_generate_content = models_resource.method(:generate_content)
+
+        # Wrap the generate_content method
+        models_resource.define_singleton_method(:generate_content) do |*args, **kwargs, &block|
+          # Include Base module methods
+          extend Base unless singleton_class.include?(Base)
+
+          start_time = Time.now
+          version = defined?(::Google::Genai::VERSION) ? ::Google::Genai::VERSION : nil
+
+          # Extract parameters for analytics
+          params = kwargs.any? ? kwargs : (args.first || {})
+
+          sentinel = instance_variable_get(:@crystil_sentinel)
+          sentinel&.raise_if_irrelevant!(
+            title: GOOGLE_CLIENT_TITLE,
+            request: params,
+            version: version
+          )
+
+          # Call original method
+          response = if kwargs.any?
+                       original_generate_content.call(**kwargs, &block)
+                     else
+                       original_generate_content.call(*args, &block)
+                     end
+
+          # Submit analytics
+          crystil_submit_analytics(
+            method: :generate_content,
+            args: args,
+            kwargs: params,
+            response: response,
+            start_time: start_time,
+            end_time: Time.now,
+            title: GOOGLE_CLIENT_TITLE,
+            version: version
+          )
+
+          response
+        rescue CrystilRequestInterceptedError => e
+          # We don't want to send intercepts to collector
+          raise e
+        rescue StandardError => e
+          params = kwargs.any? ? kwargs : (args.first || {})
+
+          crystil_submit_error_analytics(
+            method: :generate_content,
+            args: args,
+            kwargs: params,
+            error: e,
+            start_time: start_time,
+            end_time: Time.now,
+            title: GOOGLE_CLIENT_TITLE,
+            version: version
+          )
+
+          raise e
+        end
+      end
+    end
+  end
+end

data/lib/crystil/wrappers/groq.rb

@@ -0,0 +1,295 @@
+# frozen_string_literal: true
+
+require_relative "constants"
+
+module Crystil
+  module Wrappers
+    # Wrapper for the Groq Ruby client (`groq` gem, drnic/groq-ruby).
+    #
+    # Groq is an inference framework that hosts third-party models (Meta Llama,
+    # OpenAI gpt-oss, Qwen, compound systems), so it sits in
+    # `conversation.client.provider = "groq"` rather than `title`. `title` is
+    # derived per-call from the model-ID prefix — e.g.
+    # `meta-llama/llama-4-scout-17b-16e-instruct` → `"meta-llama"`,
+    # `openai/gpt-oss-20b` → `"openai"`. Legacy un-prefixed IDs
+    # (`llama-3.1-8b-instant`, `allam-2-7b`) fall back to the first
+    # alphanumeric run (`"llama"`, `"allam"`); anything unparseable falls
+    # back to `GROQ_PROVIDER` (`"groq"`). `title` is never nil.
+    #
+    # Unlike the JS / Python groq-sdk, the Ruby `groq` gem's `Client#chat`
+    # returns only the assistant message hash (`response.body.dig("choices", 0,
+    # "message")`), discarding `usage`, `model`, and the rest of the chat
+    # completion envelope. To preserve the wire shape the backend extractor
+    # expects, we patch the lower-level `Client#post(path:, body:)` and filter
+    # by path — every chat call goes through `/openai/v1/chat/completions`, and
+    # `body`/`response.body` at that layer carry the full chat-completion
+    # request and response.
+    class Groq
+      # HTTP path for Groq's chat-completion endpoint (Groq's API is OpenAI-compatible,
+      # namespaced under `/openai/v1/`). The `post` wrapper filters on this so non-chat
+      # requests pass through without analytics or sentinel.
+      CHAT_COMPLETIONS_PATH = "/openai/v1/chat/completions"
+
+      def initialize(config, collector, sentinel = nil)
+        @config = config
+        @collector = collector
+        @sentinel = sentinel
+      end
+
+      def register(client)
+        validate_client!(client)
+
+        # Prevent double registration
+        return client if client.instance_variable_defined?(:@crystil_registered)
+
+        # Patch the gem's streaming JSON parser. Class-level patch on
+        # ::Groq::Client, idempotent via @_crystil_stream_patched — runs once
+        # per process. Placed after the per-client guard so repeat register
+        # calls on the same client are a true no-op. See patch_stream_handler!
+        # for the bug being worked around.
+        patch_stream_handler!
+
+        # Store references in client instance
+        client.instance_variable_set(:@crystil_config, @config)
+        client.instance_variable_set(:@crystil_collector, @collector)
+        client.instance_variable_set(:@crystil_sentinel, @sentinel)
+        client.instance_variable_set(:@crystil_registered, true)
+
+        wrap_post_method(client)
+
+        client
+      end
+
+      # Derive a telemetry `title` from a Groq model identifier. Mirrors
+      # `extractModelTitle` in the JS SDK (`javascript-sdk/src/utils.ts`).
+      #
+      # Rules:
+      #   - String with `/`: return everything before the first `/`
+      #     (`meta-llama/llama-4-...` → `"meta-llama"`, `openai/gpt-oss-20b`
+      #     → `"openai"`).
+      #   - String without `/`: return the first run of alphanumeric characters
+      #     (`llama-3.1-8b-instant` → `"llama"`, `allam-2-7b` → `"allam"`,
+      #     `gpt-4o` → `"gpt"`).
+      #   - Anything else (non-string, empty, unrecognized shape): return
+      #     `fallback`.
+      #
+      # Telemetry invariant: `conversation.client.title` is never nil — the
+      # caller always supplies a sensible string fallback (`GROQ_PROVIDER` for
+      # Groq calls).
+      #
+      # Public because the wrapped-method closure (inside
+      # singleton_class.class_eval) needs to call it.
+      def self.extract_model_title(model, fallback)
+        return fallback unless model.is_a?(String)
+
+        slash = model.index("/")
+        return model[0...slash] || fallback if slash&.positive?
+
+        model[/\A[A-Za-z0-9]+/] || fallback
+      end
+
+      # Build the `response` hash sent on a failed call. For streaming requests
+      # that errored after one or more chunks were already merged, the partial
+      # `accumulated` response is preserved and the error info is folded in —
+      # so the backend can record what was generated before the failure.
+      # Non-streaming and pre-chunk failures get the original error-only shape.
+      #
+      # Public for the same reason as `extract_model_title` — invoked from the
+      # wrapped-method closure inside `singleton_class.class_eval`.
+      def self.build_error_response(streaming:, accumulated:, error:)
+        if streaming && accumulated.is_a?(Hash) && !accumulated.empty?
+          accumulated.merge("error" => error.message, "error_class" => error.class.name)
+        else
+          { error: error.message, class: error.class.name }
+        end
+      end
+
+      private
+
+      def validate_client!(client)
+        return if defined?(::Groq::Client) && client.is_a?(::Groq::Client)
+        # Fallback for mock objects in tests — must have both methods.
+        return if client.respond_to?(:chat) && client.respond_to?(:post)
+
+        raise RegistrationError,
+              "Client does not appear to be a valid Groq client (missing chat method)"
+      end
+
+      # The `groq` gem (drnic/groq-ruby v0.3.2) parses each SSE chunk inside
+      # `Client#to_json_stream` with:
+      #
+      #   delta = chunk.dig("choices", 0, "delta")
+      #   content = delta.dig("content")
+      #
+      # That second line crashes with `NoMethodError: undefined method 'dig' for
+      # nil:NilClass` when `delta` is nil — which happens on the terminal usage
+      # chunk Groq sends when `stream_options.include_usage = true` (the chunk
+      # has `choices: []`, so the first `dig` returns nil). Since we inject
+      # `include_usage: true` for JS/Python parity (token counts), this fires
+      # on every real streaming call. Replace the method with a copy that uses
+      # `delta&.dig("content")`. Idempotent via `@_crystil_stream_patched`.
+      def patch_stream_handler!
+        return unless defined?(::Groq::Client)
+        return if ::Groq::Client.instance_variable_defined?(:@_crystil_stream_patched)
+
+        # The gem lazy-loads event_stream_parser inside Client#chat. Our patched
+        # to_json_stream needs the constant available at the patch's class_eval
+        # site too, so eagerly require it now.
+        require "event_stream_parser"
+
+        ::Groq::Client.class_eval do
+          private
+
+          def to_json_stream(user_proc:)
+            parser = ::EventStreamParser::Parser.new
+
+            proc do |chunk, _bytes, env|
+              if env && env.status != 200
+                raise_error = Faraday::Response::RaiseError.new
+                raise_error.on_complete(env.merge(body: try_parse_json(chunk)))
+              end
+
+              parser.feed(chunk) do |_type, data|
+                next if data == "[DONE]"
+
+                chunk = JSON.parse(data)
+                delta = chunk.dig("choices", 0, "delta")
+                content = delta&.dig("content")
+
+                arity = user_proc.is_a?(Proc) ? user_proc.arity : user_proc.method(:call).arity
+                if arity == 1
+                  user_proc.call(content)
+                else
+                  user_proc.call(content, chunk)
+                end
+              end
+            end
+          end
+        end
+
+        ::Groq::Client.instance_variable_set(:@_crystil_stream_patched, true)
+      end
+
+      def wrap_post_method(client)
+        chat_path = CHAT_COMPLETIONS_PATH
+
+        client.singleton_class.class_eval do
+          include Base
+
+          alias_method :original_post, :post
+
+          define_method(:post) do |path:, body:|
+            # Non-chat-completions paths pass through untouched — no analytics,
+            # no sentinel. (At time of writing the gem only ever uses post for
+            # chat completions, but be defensive about future endpoints.)
+            return original_post(path: path, body: body) unless path == chat_path
+
+            start_time = Time.now
+            version = defined?(::Groq::VERSION) ? ::Groq::VERSION : nil
+            title = Crystil::Wrappers::Groq.extract_model_title(body[:model], GROQ_PROVIDER)
+
+            sentinel = instance_variable_get(:@crystil_sentinel)
+            sentinel&.raise_if_irrelevant!(
+              title: title,
+              request: body,
+              provider: GROQ_PROVIDER,
+              version: version
+            )
+
+            # Dup before mutating so the caller's body hash is unchanged.
+            body = body.dup
+            streaming = body[:stream_chunk].respond_to?(:call)
+            accumulated_response = streaming ? {} : nil
+
+            if streaming
+              # Match the JS/Python behavior: ask Groq to include usage on the
+              # terminal chunk so the merged response carries token counts.
+              # A caller-set value (true or false) wins — we only force `true`
+              # when the key is absent.
+              body[:stream_options] = { include_usage: true }.merge(body[:stream_options] || {})
+
+              user_callback = body[:stream_chunk]
+              body[:stream_chunk] = proc do |content, chunk|
+                if chunk.is_a?(Hash)
+                  normalized = crystil_normalize_openai_chunk(chunk)
+                  crystil_merge_streaming_chunk(accumulated_response, normalized)
+                end
+
+                # The gem inspects user_proc.arity to decide between
+                # call(content) and call(content, chunk); forward with the
+                # caller's intended signature.
+                arity = user_callback.is_a?(Proc) ? user_callback.arity : user_callback.method(:call).arity
+                if arity == 1
+                  user_callback.call(content)
+                else
+                  user_callback.call(content, chunk)
+                end
+              end
+            end
+
+            response = original_post(path: path, body: body)
+
+            final_response = if streaming
+                               accumulated_response
+                             elsif response.respond_to?(:body)
+                               response.body
+                             else
+                               response
+                             end
+
+            # Sanitize body for analytics: stream_chunk is a Proc and the
+            # collector would fail to JSON-encode it. Mirror the OpenAI
+            # wrapper's `:stream` → `true` collapse.
+            analytics_body = body.dup
+            analytics_body[:stream_chunk] = true if analytics_body[:stream_chunk].respond_to?(:call)
+
+            crystil_submit_analytics(
+              method: :post,
+              args: [],
+              kwargs: analytics_body,
+              response: final_response,
+              start_time: start_time,
+              end_time: Time.now,
+              provider: GROQ_PROVIDER,
+              title: title,
+              version: version
+            )
+
+            response
+          rescue CrystilRequestInterceptedError => e
+            # We don't want to send intercepts to collector
+            raise e
+          rescue StandardError => e
+            analytics_body = body.dup
+            analytics_body[:stream_chunk] = true if analytics_body[:stream_chunk].respond_to?(:call)
+
+            # On streaming failure after one or more chunks merged, pass the
+            # accumulated response (plus error info) through so the backend
+            # extractor can pull partial assistant content / token usage from
+            # what was received before the error. Non-streaming and
+            # pre-chunk failures fall back to the helper's default
+            # `{error, class}` shape.
+            error_response = Crystil::Wrappers::Groq.build_error_response(
+              streaming: streaming, accumulated: accumulated_response, error: e
+            )
+
+            crystil_submit_error_analytics(
+              method: :post,
+              args: [],
+              kwargs: analytics_body,
+              error: e,
+              start_time: start_time,
+              end_time: Time.now,
+              provider: GROQ_PROVIDER,
+              title: title,
+              version: version,
+              response: error_response
+            )
+
+            raise e
+          end
+        end
+      end
+    end
+  end
+end