axn-ruby_llm 0.1.2 → 0.2.1

data/lib/axn/ruby_llm/ask.rb

@@ -11,6 +11,7 @@ module Axn
       expects :model, optional: true
       expects :system_prompt, optional: true
       expects :temperature, optional: true
+      expects :tools, optional: true
 
       exposes :response
       exposes :raw_message
@@ -25,8 +26,39 @@ module Axn
 
       StubMessage = Data.define(:content, :input_tokens, :output_tokens, :cache_read_tokens, :cache_write_tokens, :model_id)
 
-      error prefix: "LLM request failed: "
-      error "Failed to parse JSON from LLM response", if: JSON::ParserError
+      # 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.
+      KNOWN_ERROR_CLASSES = [
+        ::RubyLLM::Error,
+        ::Faraday::Error,
+        ::RubyLLM::ConfigurationError,
+        ::RubyLLM::ModelNotFoundError,
+        ::RubyLLM::PromptNotFoundError,
+        ::RubyLLM::InvalidRoleError,
+        ::RubyLLM::InvalidToolChoiceError,
+        ::RubyLLM::UnsupportedAttachmentError,
+      ].freeze
+      KNOWN_ERROR = ->(exception:) { KNOWN_ERROR_CLASSES.any? { |k| exception.is_a?(k) } }
+      RETRYABLE_ERROR = lambda { |exception:|
+        [::RubyLLM::OverloadedError, ::RubyLLM::ServiceUnavailableError, ::RubyLLM::ServerError].any? { |k| exception.is_a?(k) }
+      }
+
+      # Base headlines for a consistent result.error / result.success surface: failures read
+      # "<error_headline>: <reason>" (configurable via Axn::RubyLLM.configure); successes read
+      # "LLM request completed", with any detail attached parenthetically via join: (e.g. the
+      # stubbed-values note on the disabled path below).
+      # Reason entries are ordered most-specific-last (axn checks most-recently-declared first), so a
+      # narrower match (retryable, context length, JSON parse) wins over the generic KNOWN_ERROR catch-all.
+      error { Axn::RubyLLM.config.error_headline }
+      error(if: KNOWN_ERROR, &:message)
+      error(if: RETRYABLE_ERROR) { |e| "Provider temporarily unavailable, try again later: #{e.message}" }
+      error(if: ::RubyLLM::ContextLengthExceededError) { |e| "Prompt exceeds the model's context window: #{e.message}" }
+      error "Response was not valid JSON", if: JSON::ParserError
+      success "LLM request completed", join: ->(base, reason) { "#{base} (#{reason})" }
 
       before do
         if disabled?
@@ -38,7 +70,8 @@ module Axn
             response_model: nil,
             stubbed: true,
           )
-          done!("disabled - returning stubbed values", **exposures)
+          # Reason attaches to the "LLM request completed" base via the parenthetical join: above.
+          done!("using stubbed values - actual LLM request disabled", **exposures)
         end
       end
 
@@ -46,20 +79,20 @@ module Axn
         expose(
           response: parsed_response,
           raw_message: llm_response,
-          input_tokens: llm_response.input_tokens,
-          output_tokens: llm_response.output_tokens,
-          cache_read_tokens: llm_response.cache_read_tokens,
-          cache_write_tokens: llm_response.cache_write_tokens,
+          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),
           prompt_tokens: total_input_tokens,
           cost_breakdown:,
           cost: cost_breakdown&.total,
           stubbed: false,
         )
         record_otel_attributes!(
-          input_tokens: llm_response.input_tokens,
-          output_tokens: llm_response.output_tokens,
+          input_tokens: sum_across(:input_tokens),
+          output_tokens: sum_across(:output_tokens),
           cost: cost_breakdown&.total,
-          response_model: llm_response.model_id,
+          response_model: response_message&.model_id,
           stubbed: false,
         )
       rescue ::RubyLLM::RateLimitError => e
@@ -68,7 +101,7 @@ module Axn
 
       private
 
-      def disabled? = !Axn::RubyLLM.configuration.enabled?
+      def disabled? = !Axn::RubyLLM.enabled?
 
       def stubbed_exposures
         content = schema || json ? { "stubbed" => true } : "stubbed response value"
@@ -87,6 +120,8 @@ 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)
@@ -96,52 +131,120 @@ module Axn
         json ? JSON.parse(llm_response.content) : llm_response.content
       end
 
+      # 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
+      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
+
       def total_input_tokens
-        vals = [llm_response.input_tokens, llm_response.cache_read_tokens, llm_response.cache_write_tokens]
+        vals = usage_messages.flat_map { |m| [m.input_tokens, m.cache_read_tokens, m.cache_write_tokens] }
         vals.all?(&:nil?) ? nil : vals.sum(&:to_i)
       end
 
-      def cost_breakdown
+      memo def cost_breakdown
         return nil unless model_info
 
-        llm_response.cost(model: 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
-        ::RubyLLM.models.find(llm_response.model_id)
+        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_tools(*resolved_tools) if resolved_tools.any?
         end
       end
 
       def resolved_model
-        model || Axn::RubyLLM.configuration.default_model
+        model || Axn::RubyLLM.config.default_model
+      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:)`.
+      # RubyLLM's `with_tools` accepts either a class or an instance, so wrapped classes register as-is.
+      def resolved_tools
+        Array(tools).map { |tool| _as_ruby_llm_tool(tool) }
+      end
+
+      def _as_ruby_llm_tool(tool)
+        return tool if tool.is_a?(::RubyLLM::Tool)
+        return tool if tool.is_a?(::Class) && tool < ::RubyLLM::Tool
+
+        Axn::RubyLLM.wrap(tool)
       end
 
       def record_otel_attributes!(input_tokens:, output_tokens:, cost:, response_model:, stubbed:)
-        return unless defined?(::OpenTelemetry::Trace)
-
-        span = ::OpenTelemetry::Trace.current_span
-        return 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?
-      rescue StandardError
-        # never let telemetry break the action
+        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

@@ -6,19 +6,28 @@ module Axn
   module RubyLLM
     module RSpec
       module Helpers
-        # Stubs RubyLLM so that Ask returns a canned response.
+        UNSET = Object.new
+        private_constant :UNSET
+
+        # Stubs RubyLLM so that Ask returns a canned response. The response can be given
+        # positionally (the common case) or as `response:` — both are equivalent.
         #
         # Usage in a spec:
-        #   stub_axn_ruby_llm(response: "Here is a summary.")
-        #   stub_axn_ruby_llm(response: { "key" => "value" })  # auto-JSON-serialized for json: true calls
-        #   stub_axn_ruby_llm(response: { "k" => "v" }, schema: MySchema) # Hash passed through unparsed
-        #   stub_axn_ruby_llm(response: "...", input_tokens: 100, output_tokens: 50, cost: 0.0023)
-        #   stub_axn_ruby_llm(response: "...", cache_read_tokens: 500, cache_write_tokens: 200)
+        #   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("...", 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
         #
         # Returns the chat instance double for further assertions if needed.
-        def stub_axn_ruby_llm(response:, model: nil, schema: nil, input_tokens: nil, output_tokens: nil,
-                              cache_read_tokens: nil, cache_write_tokens: nil, cost: nil)
-          resolved_model_id = model || Axn::RubyLLM.configuration.default_model
+        def stub_axn_ruby_llm(positional_response = UNSET, response: UNSET, model: nil, schema: nil,
+                              input_tokens: nil, output_tokens: nil, cache_read_tokens: nil,
+                              cache_write_tokens: nil, cost: nil)
+          response = positional_response unless positional_response.equal?(UNSET)
+          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:)
@@ -49,8 +58,9 @@ module Axn
           else
             allow(::RubyLLM).to receive(:chat).and_return(chat_instance)
           end
-          allow(chat_instance).to receive(:with_instructions).and_return(chat_instance)
-          allow(chat_instance).to receive(:with_params).and_return(chat_instance)
+          %i[with_instructions with_params 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.
@@ -60,6 +70,8 @@ module Axn
             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
 
@@ -69,7 +81,9 @@ module Axn
           # Default to zero cost so specs exercise the "model found, cost computed" path.
           # Pass cost: explicitly to assert a specific value.
           cost_total = cost || 0.0
-          cost_struct = instance_double(::RubyLLM::Cost, total: cost_total)
+          # 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)
         end
       end