axn-ruby_llm 0.2.0 → 0.3.0

data/lib/axn/ruby_llm/ask.rb

@@ -6,7 +6,6 @@ module Axn
       include Axn
 
       expects :prompt
-      expects :json, type: :boolean, default: false
       expects :schema, optional: true
       expects :model, optional: true
       expects :system_prompt, optional: true
@@ -24,22 +23,36 @@ module Axn
       exposes :cost_breakdown, allow_nil: true
       exposes :stubbed, type: :boolean, default: false
 
-      StubMessage = Data.define(:content, :input_tokens, :output_tokens, :cache_read_tokens, :cache_write_tokens, :model_id)
+      # Shape-compatible with a real ::RubyLLM::Message on the disabled path: `.content` is the raw
+      # text (JSON when `schema:` is set, matching 2.0's read-only String #content), `.tokens` is a
+      # real ::RubyLLM::Tokens (so `.input`/`.output`/`.cache_read`/`.cache_write` all resolve), and
+      # `.parsed` mirrors Message#parsed (memoized JSON.parse over #content).
+      StubMessage = Data.define(:content, :tokens, :model) do
+        def parsed
+          return if content.nil? || content.empty?
+
+          JSON.parse(content)
+        end
+      end
 
       # RubyLLM wraps HTTP-response-level provider errors (4xx/5xx) under RubyLLM::Error, but its
-      # non-HTTP errors (bad config, missing model/prompt/role, unsupported attachment) subclass
-      # StandardError directly -- so RubyLLM::Error alone misses them. Connection-level failures
-      # (timeout, DNS, refused) never reach RubyLLM at all and surface as raw Faraday errors. All
-      # three are "known" failure shapes safe to surface verbatim; anything outside this is a bug
-      # and must not leak its message into a user-facing result.
+      # non-HTTP errors (bad config, missing model/prompt/role, unsupported attachment, a stale model
+      # registry, an unresolved pending-tool-call/approval loop state) subclass StandardError
+      # directly -- so RubyLLM::Error alone misses them. Connection-level failures (timeout, DNS,
+      # refused) never reach RubyLLM at all and surface as raw Faraday errors. All these are "known"
+      # failure shapes safe to surface verbatim; anything outside this is a bug and must not leak its
+      # message into a user-facing result.
       KNOWN_ERROR_CLASSES = [
         ::RubyLLM::Error,
         ::Faraday::Error,
         ::RubyLLM::ConfigurationError,
         ::RubyLLM::ModelNotFoundError,
+        ::RubyLLM::ModelRegistryError,
         ::RubyLLM::PromptNotFoundError,
         ::RubyLLM::InvalidRoleError,
         ::RubyLLM::InvalidToolChoiceError,
+        ::RubyLLM::PendingToolCallsError,
+        ::RubyLLM::CancelledError,
         ::RubyLLM::UnsupportedAttachmentError,
       ].freeze
       KNOWN_ERROR = ->(exception:) { KNOWN_ERROR_CLASSES.any? { |k| exception.is_a?(k) } }
@@ -79,20 +92,20 @@ module Axn
         expose(
           response: parsed_response,
           raw_message: llm_response,
-          input_tokens: sum_across(:input_tokens),
-          output_tokens: sum_across(:output_tokens),
-          cache_read_tokens: sum_across(:cache_read_tokens),
-          cache_write_tokens: sum_across(:cache_write_tokens),
+          input_tokens: token_usage.input,
+          output_tokens: token_usage.output,
+          cache_read_tokens: token_usage.cache_read,
+          cache_write_tokens: token_usage.cache_write,
           prompt_tokens: total_input_tokens,
           cost_breakdown:,
           cost: cost_breakdown&.total,
           stubbed: false,
         )
         record_otel_attributes!(
-          input_tokens: sum_across(:input_tokens),
-          output_tokens: sum_across(:output_tokens),
+          input_tokens: token_usage.input,
+          output_tokens: token_usage.output,
           cost: cost_breakdown&.total,
-          response_model: response_message&.model_id,
+          response_model: llm_response&.model,
           stubbed: false,
         )
       rescue ::RubyLLM::RateLimitError => e
@@ -104,10 +117,12 @@ module Axn
       def disabled? = !Axn::RubyLLM.enabled?
 
       def stubbed_exposures
-        content = schema || json ? { "stubbed" => true } : "stubbed response value"
+        parsed_content = schema ? { "stubbed" => true } : nil
+        content = parsed_content ? parsed_content.to_json : "stubbed response value"
+        zero_tokens = ::RubyLLM::Tokens.new(input: 0, output: 0, cache_read: 0, cache_write: 0)
         {
-          response: content,
-          raw_message: StubMessage.new(content:, input_tokens: 0, output_tokens: 0, cache_read_tokens: 0, cache_write_tokens: 0, model_id: "stubbed"),
+          response: parsed_content || content,
+          raw_message: StubMessage.new(content:, tokens: zero_tokens, model: "stubbed"),
           input_tokens: 0,
           output_tokens: 0,
           cache_read_tokens: 0,
@@ -120,99 +135,38 @@ module Axn
       end
 
       def parsed_response
-        return halted_response if halted?
-
-        if schema
-          # with_schema makes RubyLLM parse the response into a Hash on success
-          return llm_response.content if llm_response.content.is_a?(Hash)
+        return llm_response.content unless schema
 
-          fail! "Schema response was not valid JSON"
-        end
-        json ? JSON.parse(llm_response.content) : llm_response.content
-      end
+        # with_schema makes RubyLLM parse the response into JSON text on success; #parsed memoizes
+        # JSON.parse over #content and raises JSON::ParserError on malformed JSON (caught by the
+        # declared `error "Response was not valid JSON", if: JSON::ParserError` handler above).
+        parsed = llm_response.parsed
+        return parsed if parsed.is_a?(Hash)
 
-      # A halted tool (halt_after:) short-circuits the model's final turn, so with_schema/json never
-      # parsed a model response — the "response" is the tool's own payload (Halt#content). For a
-      # :structured tool that's JSON text, so parse it to honor the Hash contract a schema:/json:
-      # caller expects; fall back to the raw string for a :message tool or an unparseable payload.
-      def halted_response
-        content = llm_response.content
-        return content unless (schema || json) && content.is_a?(String)
-
-        JSON.parse(content)
-      rescue JSON::ParserError
-        content
+        fail! "Schema response was not valid JSON"
       end
 
-      # A tool call makes multiple model round-trips inside one `ask`; every assistant turn is
-      # accumulated on the chat and reports its OWN usage, so sum across them for the true per-call
-      # totals rather than just the final turn's. Non-response messages (the user prompt, tool
-      # results) carry no tokens — they contribute 0 to the token sums, and RubyLLM::Cost.aggregate
-      # ignores them (no `tokens?`) — so summing over every message is correct, and a plain (no-tool)
-      # ask (one assistant turn) is a no-op.
-      def usage_messages
-        chat.messages
-      end
-
-      # nil only when NO turn reported the field (preserving the "nil if the provider didn't return it"
-      # contract); otherwise the summed count, treating a missing turn as 0.
-      def sum_across(field)
-        values = usage_messages.map(&field)
-        values.all?(&:nil?) ? nil : values.sum(&:to_i)
-      end
+      # Every provider attempt this chat has made -- including retries and fallback attempts that
+      # produced no message -- aggregated by RubyLLM itself (Chat#tokens / Chat#cost), rather than
+      # summed by hand across chat.messages. A tool call makes multiple model round-trips inside one
+      # `ask`; this still reflects the whole call, not just the final response.
+      memo def token_usage = chat.tokens
+      memo def cost_breakdown = chat.cost
 
+      # nil only when NO turn reported the field (preserving the "nil if the provider didn't return
+      # it" contract); otherwise the summed count, treating a missing component as 0.
       def total_input_tokens
-        vals = usage_messages.flat_map { |m| [m.input_tokens, m.cache_read_tokens, m.cache_write_tokens] }
+        vals = [token_usage.input, token_usage.cache_read, token_usage.cache_write]
         vals.all?(&:nil?) ? nil : vals.sum(&:to_i)
       end
 
-      memo def cost_breakdown
-        return nil unless model_info
-
-        # chat.messages always includes the user prompt (chat.ask appends it before completing) and,
-        # in a tool loop, the tool-result messages -- none of which carry token usage. Keep only the
-        # token-bearing (billable) costs BEFORE the one?-vs-aggregate decision: a normal single-turn
-        # call then preserves the response's OWN Cost (with its tokens/model) via costs.one?, instead
-        # of being forced through aggregate -- which returns a Cost with nil tokens/model -- by the
-        # ever-present user message. `select(&:tokens?)` mirrors Cost.aggregate's own billable filter,
-        # so the multi-turn total is unchanged.
-        costs = usage_messages.map { |message| message.cost(model: model_info) }.select(&:tokens?)
-        return nil if costs.empty?
-
-        # One billable turn → its own Cost (identical to the pre-tool-loop non-tool call). Multiple →
-        # RubyLLM::Cost.aggregate sums the per-tier costs into a single breakdown.
-        costs.one? ? costs.first : ::RubyLLM::Cost.aggregate(costs)
-      end
-
-      memo def model_info
-        return nil unless response_message&.model_id
-
-        ::RubyLLM.models.find(response_message.model_id)
-      rescue ::RubyLLM::ModelNotFoundError
-        nil
-      end
-
       memo def llm_response = chat.ask(prompt)
 
-      # When a wrapped tool halts the loop (halt_after:), chat.ask returns a ::RubyLLM::Tool::Halt
-      # carrying the tool payload as #content, not a Message — and a Halt has no #model_id. Read the
-      # model (for cost lookup + OTel) from the last assistant turn accumulated on the chat in that
-      # case; for a normal response, llm_response IS that final message. Token/cost SUMS already read
-      # chat.messages, so only the model-id reads needed this indirection.
-      def response_message
-        return llm_response unless halted?
-
-        chat.messages.reverse.find { |message| message.role == :assistant }
-      end
-
-      def halted? = llm_response.is_a?(::RubyLLM::Tool::Halt)
-
       memo def chat
         ::RubyLLM.chat(model: resolved_model).tap do |c|
           c.with_instructions(system_prompt) if system_prompt
-          c.with_schema(schema) if schema
-          c.with_params(response_format: { type: "json_object" }) if json && !schema
-          c.with_params(temperature:) if temperature
+          c.with_schema(resolved_schema) if schema
+          c.with_temperature(temperature) if temperature
           c.with_tools(*resolved_tools) if resolved_tools.any?
         end
       end
@@ -221,6 +175,78 @@ module Axn
         model || Axn::RubyLLM.config.default_model
       end
 
+      # `schema:` accepts a raw JSON Schema Hash (passed through unchanged -- Chat#with_schema
+      # already normalizes a bare Hash), a Schematist::Schema class/instance (likewise passed
+      # through -- with_schema itself checks for #to_json_schema), or an Axn class: the same
+      # reflection the tool adapter already uses for input (`input_schema`), mirrored here for
+      # output. `output_schema` is axn's own public JSON Schema Hash for its `exposes` contract.
+      #
+      # Adjustments axn's own `exposes` contract has no reason to make on its own -- each one
+      # confirmed live against a real model, not just read off docs (a docs summary claimed
+      # `minLength` was ALSO unsupported by Anthropic; a flat schema with `minLength: 1` on a
+      # String field succeeded live regardless, so only what's actually confirmed failing is
+      # stripped, nothing broader):
+      #
+      # - `additionalProperties: false` on every fixed-shape object node. Anthropic's
+      #   `output_config.format.schema` REQUIRES this unconditionally -- confirmed live
+      #   ("output_config.format.schema: For 'object' type, 'additionalProperties' must be
+      #   explicitly set to false"), and it deletes any `strict:` key before validating
+      #   (protocols/anthropic/chat.rb#build_output_config), so there is no "non-strict" escape
+      #   hatch there. OpenAI's strict mode has the same requirement (its own docs / community
+      #   reports). Injected on every object node that declares `properties` and has no
+      #   `additionalProperties` of its own -- which excludes a map (`type: Hash, of: {...}`),
+      #   whose `additionalProperties` already names its value schema and must stay that way.
+      # - `minProperties`/`maxProperties` stripped from every object node. axn emits
+      #   `minProperties: 1` on any fixed-shape object (a nested `type: Hash, shape: {...}` field,
+      #   or a map) by default -- Anthropic's schema validator rejects it outright, confirmed live
+      #   ("output_config.format.schema: For 'object' type, property 'minProperties' is not
+      #   supported"). No prose-restatement fallback (the kind PRO-3172's now-deleted Gemini
+      #   workaround used) -- that machinery existed for a fixed-whitelist converter Gemini no
+      #   longer has (see tool_adapter.rb); reintroducing it here for a narrower, less common case
+      #   (a nested fixed-shape object's entry-count bound) isn't worth the complexity back.
+      # - `strict: false`, pinned rather than left to RubyLLM's own inference. Chat#with_schema's
+      #   strict_schema? (chat_completions/chat.rb) infers `strict: true` whenever every property
+      #   is required -- the common case for an Axn's output contract (see the README example) --
+      #   and OpenAI's *full* strict mode additionally requires every property to be listed in
+      #   `required` even when conceptually optional (via a nullable type), which axn's reflection
+      #   doesn't promise. Pinned rather than relying on the `additionalProperties: false` fix
+      #   above to make strict inference merely harmless.
+      def resolved_schema
+        return schema unless schema.is_a?(::Class) && schema.respond_to?(:output_schema)
+
+        { name: schema.name || "response", schema: sanitize_output_schema(schema.output_schema), strict: false }
+      end
+
+      # Builds a new Hash/Array throughout rather than mutating -- axn may hand back a memoized
+      # output_schema, and mutating it would corrupt every other reader.
+      #
+      # Every adjustment is gated on `object_node` -- the recursion walks every Hash in the
+      # schema, but not every Hash IS a schema node. `properties` is a name-to-schema map, so an
+      # Axn with a field literally named `minProperties`/`maxProperties`/`additionalProperties`
+      # puts a Hash at exactly the key this pass reads; an ungated delete/injection would corrupt
+      # that container instead of a schema node's own keywords (confirmed live: an Axn exposing
+      # `minProperties` had that field silently dropped from `properties` while `required` still
+      # named it -- an invalid schema). The `properties` container itself never carries `type`, so
+      # gating on `object_node` -- true only for an actual object-schema node -- keeps the pass off
+      # it entirely.
+      def sanitize_output_schema(node)
+        case node
+        when Hash
+          rebuilt = node.transform_values { |value| sanitize_output_schema(value) }
+          object_node = rebuilt[:type] == "object" || Array(rebuilt[:type]).include?("object")
+          if object_node
+            rebuilt.delete(:minProperties)
+            rebuilt.delete(:maxProperties)
+            rebuilt[:additionalProperties] = false if rebuilt.key?(:properties) && !rebuilt.key?(:additionalProperties)
+          end
+          rebuilt
+        when Array
+          node.map { |value| sanitize_output_schema(value) }
+        else
+          node
+        end
+      end
+
       # `tools:` accepts a mix of bare Axn classes (wrapped here, so callers can pass their own Axns
       # straight in) and already-wrapped `::RubyLLM::Tool`s -- a class or an instance, the latter being
       # how you pass a tool that closed over explicit context via `Axn::RubyLLM.wrap(axn, ambient_context:)`.
@@ -237,22 +263,14 @@ module Axn
       end
 
       def record_otel_attributes!(input_tokens:, output_tokens:, cost:, response_model:, stubbed:)
-        # Telemetry is a best-effort side effect: it must never break the LLM call. Route it through
-        # axn core's guard (PRO-2950) rather than a bare rescue — it swallows + warn-logs on failure
-        # (and fails loud in dev when best_effort_raises_in_dev is set) instead of silently vanishing.
-        Axn::Extensions.best_effort("recording OpenTelemetry attributes on the axn.call span") do
-          next unless defined?(::OpenTelemetry::Trace)
-
-          span = ::OpenTelemetry::Trace.current_span
-          next unless span&.context&.valid?
-
-          span.set_attribute("gen_ai.request.model", resolved_model) if resolved_model
-          span.set_attribute("gen_ai.response.model", response_model) if response_model
-          span.set_attribute("gen_ai.usage.input_tokens", input_tokens) if input_tokens
-          span.set_attribute("gen_ai.usage.output_tokens", output_tokens) if output_tokens
-          span.set_attribute("gen_ai.usage.cost", cost) if cost
-          span.set_attribute("axn.ruby_llm.stubbed", stubbed) unless stubbed.nil?
-        end
+        Axn::Extensions::Tracing.annotate_span(
+          "gen_ai.request.model" => resolved_model,
+          "gen_ai.response.model" => response_model,
+          "gen_ai.usage.input_tokens" => input_tokens,
+          "gen_ai.usage.output_tokens" => output_tokens,
+          "gen_ai.usage.cost" => cost,
+          "axn.ruby_llm.stubbed" => stubbed,
+        )
       end
     end
   end

data/lib/axn/ruby_llm/rspec.rb

@@ -14,8 +14,7 @@ module Axn
         #
         # Usage in a spec:
         #   stub_axn_ruby_llm("Here is a summary.")
-        #   stub_axn_ruby_llm({ "key" => "value" })          # auto-JSON-serialized for json: true calls
-        #   stub_axn_ruby_llm({ "k" => "v" }, schema: MySchema) # Hash passed through unparsed
+        #   stub_axn_ruby_llm({ "k" => "v" }, schema: MySchema)  # Hash passed through as `parsed`
         #   stub_axn_ruby_llm("...", input_tokens: 100, output_tokens: 50, cost: 0.0023)
         #   stub_axn_ruby_llm("...", cache_read_tokens: 500, cache_write_tokens: 200)
         #   stub_axn_ruby_llm(response: "...")               # keyword form still works
@@ -28,63 +27,44 @@ module Axn
           raise ArgumentError, "stub_axn_ruby_llm requires a response (positionally or as `response:`)" if response.equal?(UNSET)
 
           resolved_model_id = model || Axn::RubyLLM.config.default_model
-          llm_message = _stub_axn_ruby_llm_message(response, resolved_model_id, input_tokens, output_tokens,
-                                                   cache_read_tokens:, cache_write_tokens:, schema:)
-          chat_instance = _stub_axn_ruby_llm_chat(model, llm_message, schema:)
-          _stub_axn_ruby_llm_cost(llm_message, resolved_model_id, cost)
-          chat_instance
+          llm_message = _stub_axn_ruby_llm_message(response, resolved_model_id, schema:)
+          _stub_axn_ruby_llm_chat(model, llm_message, input_tokens:, output_tokens:,
+                                                      cache_read_tokens:, cache_write_tokens:, cost:)
         end
 
         private
 
-        def _stub_axn_ruby_llm_message(response, model_id, input_tokens, output_tokens, cache_read_tokens:,
-                                       cache_write_tokens:, schema:)
-          content = if schema
-                      response
-                    elsif response.is_a?(Hash)
-                      response.to_json
-                    else
-                      response.to_s
-                    end
-          instance_double(::RubyLLM::Message,
-                          content:, input_tokens:, output_tokens:,
-                          cache_read_tokens:, cache_write_tokens:, model_id:)
+        # `content` mirrors real ::RubyLLM::Message#content (a read-only String, JSON text when
+        # `schema:` is set); `parsed` mirrors #parsed (the Hash `schema:` callers actually want back
+        # via Ask's `parsed_response`, which reads `.parsed` -- not `.content` -- once schema is set).
+        def _stub_axn_ruby_llm_message(response, model_id, schema:)
+          content = schema ? response.to_json : response.to_s
+          parsed = schema ? response : nil
+          instance_double(::RubyLLM::Message, content:, parsed:, model: model_id)
         end
 
-        def _stub_axn_ruby_llm_chat(model, llm_message, schema:)
+        def _stub_axn_ruby_llm_chat(model, llm_message, input_tokens:, output_tokens:,
+                                    cache_read_tokens:, cache_write_tokens:, cost:)
           chat_instance = instance_double(::RubyLLM::Chat)
           if model
             allow(::RubyLLM).to receive(:chat).with(model:).and_return(chat_instance)
           else
             allow(::RubyLLM).to receive(:chat).and_return(chat_instance)
           end
-          %i[with_instructions with_params with_tools].each do |method|
+          %i[with_instructions with_schema with_temperature with_provider_options with_tools].each do |method|
             allow(chat_instance).to receive(method).and_return(chat_instance)
           end
-          # Always stub with_schema so specs don't blow up if production code passes schema:
-          # even when the helper is called without schema:. Use a tight matcher when schema
-          # is known so the stub still validates the correct class is passed.
-          if schema
-            allow(chat_instance).to receive(:with_schema).with(schema).and_return(chat_instance)
-          else
-            allow(chat_instance).to receive(:with_schema).and_return(chat_instance)
-          end
           allow(chat_instance).to receive(:ask).and_return(llm_message)
-          # Ask sums usage across the chat's assistant turns; a stubbed call is single-turn.
-          allow(chat_instance).to receive(:messages).and_return([llm_message])
-          chat_instance
-        end
-
-        def _stub_axn_ruby_llm_cost(llm_message, model_id, cost)
-          model_info = instance_double("RubyLLM::Model")
-          allow(::RubyLLM.models).to receive(:find).with(model_id).and_return(model_info)
-          # Default to zero cost so specs exercise the "model found, cost computed" path.
+          # Ask reads usage off the chat's own ledger (Chat#tokens / Chat#cost), not per-message --
+          # a stubbed call is single-turn, so the ledger is just these values directly.
+          allow(chat_instance).to receive(:tokens).and_return(
+            instance_double(::RubyLLM::Tokens, input: input_tokens, output: output_tokens,
+                                               cache_read: cache_read_tokens, cache_write: cache_write_tokens),
+          )
+          # Default to zero cost so specs exercise the "cost computed" path.
           # Pass cost: explicitly to assert a specific value.
-          cost_total = cost || 0.0
-          # tokens?: true — the stubbed message is a billable assistant turn, and Ask's cost_breakdown
-          # keeps only token-bearing (tokens?) costs before its one?-vs-aggregate decision.
-          cost_struct = instance_double(::RubyLLM::Cost, total: cost_total, tokens?: true)
-          allow(llm_message).to receive(:cost).with(model: model_info).and_return(cost_struct)
+          allow(chat_instance).to receive(:cost).and_return(instance_double(::RubyLLM::Cost, total: cost || 0.0))
+          chat_instance
         end
       end
     end