ask-agent 0.2.2 → 0.3.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 38ad7fd8e88fbb94137ceff0dacd127b598e3487c69b1b58feaab1299c3e0d35
-  data.tar.gz: 6b86830a8ef5547305efdb5d8aa36b38de9d4ea0b1ce55c6c207668575ff554c
+  metadata.gz: 26c124febbf2652616754dd539cc8df5d24afe5fbe8c03381f760598578a051f
+  data.tar.gz: 50d87794d60901a85da426f50a1a619807faf96bbe1fedb58145dd2d5d7fbeb1
 SHA512:
-  metadata.gz: c9e4653c5d70595143daf009735193f984e7f972bc075658508b8baf26d3fcdb044fa251c50d4de95593b268ec2715334c67bc2e7a7578fee419f2da4eab8cd0
-  data.tar.gz: e375fe521571910a21c00dfe169b506fc617e0c46bc235e108369c197ee866c05e42813905dcbc9445c898d8143d1b826666dd6593816feefa5fde7b8f539a8c
+  metadata.gz: 5927098dcbd1fe84faa0cf2a7cddbc63e544a7113e47c61f972044e598053758417edce0e63d9a55d7922fbbc65ad9f552ef5c3fb1c07ac71ba1b718d1cbb815
+  data.tar.gz: 0436ac54f1dcd57ac06706d7e6a8ad5030574341170e0c7e99ca174d3dd94341651142c78bf5d138be1f99a978d71eb21f03a57df1ef4b6cb038ebb22778a637

data/CHANGELOG.md

@@ -1,3 +1,27 @@
+## [0.3.1] — 2026-07-17
+
+### Added
+
+- **Rate-limit aware retry in Chat** — `Chat#ask` retries up to 3 times on `RateLimitError`, using `retry_after` from the error when available, otherwise exponential backoff with jitter.
+
+### Fixed
+
+- **`retryable_error_name?` in ToolExecutor** — fixed duplicate `Ask::RateLimitError` and non-existent `Ask::ServiceUnavailableError`. Now uses class hierarchy matching so subclasses are also retried. (Backport from LiteLLM error classification.)
+
+## [0.3.0] — 2026-07-17
+
+### Added
+
+- **Token and cost tracking** — `ResponseMessage` and `ChatChunk` now carry `input_tokens`, `output_tokens`, and `cost` fields. Token counts are extracted from provider responses and streaming chunks.
+- **Instrumentation events** — `Chat#ask` emits `chat.ask` and `chat.stream.ask` events via `Ask::Instrumentation`, unlocking the full monitoring pipeline (ask-agent → ask-instrumentation → ask-monitoring).
+- **Cost in agent events** — `SessionEnd` and `TurnEnd` events now include `input_tokens`, `output_tokens`, and `cost` fields, accumulated across all turns in the session.
+- **Cumulative session costs** — `Session` tracks `total_input_tokens`, `total_output_tokens`, and `total_cost` across all turns and reflection rounds.
+
+### Changed
+
+- **Dependency added** — `ask-instrumentation >= 0.1` added to gemspec. Instrumentation is optional (emission is wrapped in `defined?` check).
+- **Gemfile** — now uses local path resolution for sibling ask-* gems during development.
+
 ## [0.2.1] - 2026-06-25
 
 ### Changed

data/lib/ask/agent/chat.rb

@@ -3,42 +3,27 @@
 module Ask
   module Agent
     # Response message returned by {Chat#ask}.
-    # Presents a message-like response interface for ask-agent internal use.
-    ResponseMessage = Data.define(:content, :tool_calls, :thinking) do
+    # Includes token counts and cost when available from the provider.
+    ResponseMessage = Data.define(:content, :tool_calls, :thinking, :input_tokens, :output_tokens, :cost) do
       def tool_call? = !tool_calls.empty?
       def to_s = content.to_s
     end
 
-    # Tool call data used in {ResponseMessage} and {ChatChunk}.
     ToolCallInfo = Data.define(:id, :name, :arguments)
 
-    # Chunk yielded during streaming from {Chat#ask}.
-    ChatChunk = Data.define(:content, :tool_calls, :thinking) do
+    ChatChunk = Data.define(:content, :tool_calls, :thinking, :input_tokens, :output_tokens) do
       def tool_call? = !tool_calls.empty?
     end
 
-    # Thin wrapper around {Ask::Provider} + an internal message array that
-    # presents a Chat-like API for ask-agent internal use.
-    #
-    # Manages conversation history, resolves the correct provider/model,
-    # handles streaming chunk accumulation, and normalises tool call
-    # formats between Ask::Provider (Array of Hashes) and ask-agent
-    # internal usage (Hash of { id => ToolCallInfo }).
     class Chat
-      # @return [String] model ID (e.g. "gpt-4o")
-        attr_reader :model_id
-      # @return [String] model ID (e.g. "gpt-4o")
+      attr_reader :model_id
+
       def model
         @model_id
       end
-      # @return [Array<Ask::Message>] all messages in the conversation
+
       attr_reader :messages
 
-      # @param model [String, #ask] model ID or chat-like object
-      # @param tools [Array<Ask::Tool>] tool instances available to the chat
-      # @param temperature [Float, nil] sampling temperature
-      # @param schema [Ask::Schema, Hash, nil] structured output schema
-      # @param provider [String, Symbol, nil] provider slug (overrides catalog lookup)
       def initialize(model:, tools: [], temperature: nil, schema: nil, provider: nil, **)
         @model_id = model.respond_to?(:id) ? model.id : model.to_s
         @model_info = Ask::ModelCatalog.find(@model_id)
@@ -50,40 +35,21 @@ module Ask
         @provider = nil
       end
 
-      # Send a user message and get a completion response.
-      #
-      # @param message [String, nil] user message text
-      # @yield [ChatChunk] streaming chunks (only when a block is given)
-      # @return [ResponseMessage] the assistant's response
       def ask(message = nil, &block)
         @messages << Ask::Message.new(role: :user, content: message.to_s) if message
 
         stream = block_given?
         tool_defs = @tools.map { |t| Ask::ToolDef.from_tool(t) }
 
-        # Accumulator for tool calls during streaming (keyed by index)
         calls_acc = {}
 
-        result = provider.chat(          @messages.map(&:to_h),
-          model: @model_id,
-          tools: tool_defs,
-          temperature: @temperature,
-          stream: stream,
-          schema: @schema&.respond_to?(:to_json_schema) ? @schema.to_json_schema : @schema,
-          **(@extra_params || {})
-        ) do |raw_chunk|
-          next unless block_given?
-
-          # Accumulate tool calls by index during streaming
-          accumulate_tool_calls(raw_chunk, calls_acc)
-
-          # Yield adapted chunk with current tool call state
-          yield ChatChunk.new(
-            content: raw_chunk.content,
-            tool_calls: build_current_tool_calls(calls_acc),
-            thinking: raw_chunk.respond_to?(:thinking) ? raw_chunk.thinking : nil
-          )
-        end
+        provider_model = @model_id
+        provider_tools = tool_defs
+        provider_temp = @temperature
+        provider_schema = @schema&.respond_to?(:to_json_schema) ? @schema.to_json_schema : @schema
+        provider_params = @extra_params || {}
+
+        result = chat_with_retry(stream, calls_acc, &block)
 
         response_msg = if stream
           build_stream_response(result, calls_acc)
@@ -91,24 +57,24 @@ module Ask
           build_response(result)
         end
 
-        # Store assistant response in conversation history
         @messages << Ask::Message.new(
           role: :assistant,
           content: response_msg.content,
           tool_calls: response_msg.tool_calls&.values&.map { |tc|
             { id: tc.id, type: "function", name: tc.name, arguments: tc.arguments }
-          }
+          },
+          metadata: {
+            input_tokens: response_msg.input_tokens,
+            output_tokens: response_msg.output_tokens,
+            cost: response_msg.cost
+          }.compact
         )
 
+        emit_instrumentation(stream, response_msg)
+
         response_msg
       end
 
-      # Add a message to the conversation history.
-      #
-      # @param role [Symbol] :system, :user, :assistant, :tool
-      # @param content [String, nil] message content
-      # @param tool_call_id [String, nil] tool call ID (for tool results)
-      # @param tool_calls [Array<Hash>, nil] tool call invocations
       def add_message(role:, content: nil, tool_call_id: nil, tool_calls: nil)
         @messages << Ask::Message.new(
           role: role,
@@ -118,46 +84,28 @@ module Ask
         )
       end
 
-      # Set or replace the system prompt.
-      #
-      # @param prompt [String] system instructions
-      # @return [self]
       def with_instructions(prompt)
         @messages.reject! { |m| m.role == :system }
         @messages.unshift(Ask::Message.new(role: :system, content: prompt))
         self
       end
 
-      # Set the structured output schema and return self.
-      #
-      # @param schema [Ask::Schema, Hash] structured output schema
-      # @return [self]
-  # Set additional parameters for the provider call and return self.
-  #
-  # @param params [Hash] extra parameters passed to the provider
-  # @return [self]
-  def with_params(**params)
-    @extra_params = (@extra_params || {}).merge(params)
-    self
-  end
+      def with_params(**params)
+        @extra_params = (@extra_params || {}).merge(params)
+        self
+      end
 
-        # Set additional parameters forwarded to the provider call.
-  # @param params [Hash] extra keyword arguments for the provider
-  # @return [self]
-def with_schema(schema)
+      def with_schema(schema)
         @schema = schema.respond_to?(:to_json_schema) ? schema.to_json_schema : schema
         self
       end
 
-      # Clear all messages from the conversation.
       def reset_messages!
         @messages.clear
       end
 
       private
 
-      # Resolve model info from the catalog.
-      # Lazily resolve and instantiate the LLM provider.
       def provider
         @provider ||= build_provider
       end
@@ -177,7 +125,6 @@ def with_schema(schema)
         Ask::LLM::Config.new(config)
       end
 
-      # Accumulate partial tool calls from streaming chunks.
       def accumulate_tool_calls(raw_chunk, calls_acc)
         return unless raw_chunk.tool_call?
 
@@ -190,7 +137,6 @@ def with_schema(schema)
         end
       end
 
-      # Build current snapshot of tool calls from accumulator.
       def build_current_tool_calls(calls_acc)
         hash = {}
         calls_acc.each_value do |tc_data|
@@ -204,7 +150,6 @@ def with_schema(schema)
         hash
       end
 
-      # Convert Ask::Provider tool_calls (Array of Hashes) to Hash.
       def build_tool_call_hash(raw_calls)
         hash = {}
         raw_calls.each do |tc|
@@ -219,21 +164,111 @@ def with_schema(schema)
         hash
       end
 
-      # Build response from streaming result.
       def build_stream_response(stream, calls_acc)
-        thinking = stream.chunks.filter_map(&:thinking).last
+        tokens = accumulated_tokens(stream)
+        cost = calculate_cost(tokens[:input], tokens[:output])
         ResponseMessage.new(
           content: stream.accumulated_text,
           tool_calls: build_current_tool_calls(calls_acc),
-          thinking: thinking
+          thinking: stream.chunks.filter_map(&:thinking).last,
+          input_tokens: tokens[:input],
+          output_tokens: tokens[:output],
+          cost: cost
         )
       end
 
-      # Build response from non-streaming result.
       def build_response(msg)
         tool_calls = msg.tool_calls ? build_tool_call_hash(msg.tool_calls) : {}
         thinking = msg.respond_to?(:thinking) ? msg.thinking : nil
-        ResponseMessage.new(content: msg.content.to_s, tool_calls: tool_calls, thinking: thinking)
+        metadata = msg.metadata || {}
+        input_tokens = metadata[:input_tokens] || metadata["input_tokens"]
+        output_tokens = metadata[:output_tokens] || metadata["output_tokens"]
+        cost = calculate_cost(input_tokens, output_tokens)
+        ResponseMessage.new(
+          content: msg.content.to_s,
+          tool_calls: tool_calls,
+          thinking: thinking,
+          input_tokens: input_tokens,
+          output_tokens: output_tokens,
+          cost: cost
+        )
+      end
+
+      def accumulated_tokens(stream)
+        input = 0
+        output = 0
+        stream.chunks.each do |chunk|
+          if chunk.usage
+            input = chunk.usage[:input_tokens] || chunk.usage["input_tokens"] || input
+            output = chunk.usage[:output_tokens] || chunk.usage["output_tokens"] || output
+          end
+          output += 1 if chunk.content.to_s.length > 0
+        end
+        { input: input, output: output }
+      end
+
+      def calculate_cost(input_tokens, output_tokens)
+        return nil unless input_tokens || output_tokens
+        Ask::LLM::CostCalculator.calculate(@model_info, input_tokens: input_tokens || 0, output_tokens: output_tokens || 0)
+      rescue StandardError
+        nil
+      end
+
+      MAX_CHAT_RETRIES = 3
+
+      def chat_with_retry(stream, calls_acc, &block)
+        MAX_CHAT_RETRIES.times do |attempt|
+          begin
+            return provider.chat(
+              @messages.map(&:to_h),
+              model: @model_id,
+              tools: @tools.map { |t| Ask::ToolDef.from_tool(t) },
+              temperature: @temperature,
+              stream: stream,
+              schema: @schema&.respond_to?(:to_json_schema) ? @schema.to_json_schema : @schema,
+              **(@extra_params || {})
+            ) do |raw_chunk|
+              next unless block
+
+              accumulate_tool_calls(raw_chunk, calls_acc)
+
+              block.call(ChatChunk.new(
+                content: raw_chunk.content,
+                tool_calls: build_current_tool_calls(calls_acc),
+                thinking: raw_chunk.respond_to?(:thinking) ? raw_chunk.thinking : nil,
+                input_tokens: nil,
+                output_tokens: nil
+              ))
+            end
+          rescue Ask::RateLimitError => e
+            raise if attempt >= MAX_CHAT_RETRIES - 1
+
+            delay = e.retry_after || ((2 ** attempt) + rand(0.0..1.0))
+            sleep(delay)
+          end
+        end
+      end
+
+      def emit_instrumentation(stream, response_msg)
+        return unless defined?(Ask::Instrumentation)
+
+        payload = {
+          model: @model_id,
+          provider: @model_info.provider,
+          input_tokens: response_msg.input_tokens,
+          output_tokens: response_msg.output_tokens,
+          cost: response_msg.cost,
+          tool_calls: response_msg.tool_call?,
+          stream: stream
+        }.compact
+
+        if stream
+          Ask::Instrumentation.instrument("chat.stream.ask", payload)
+        else
+          Ask::Instrumentation.instrument("chat.ask", payload)
+        end
+      rescue StandardError
+        nil
       end
     end
   end

data/lib/ask/agent/events.rb

@@ -4,10 +4,10 @@ module Ask
   module Agent
     module Events
       SessionStart = Data.define
-      SessionEnd = Data.define(:result, :turn_count, :tool_calls_made)
+      SessionEnd = Data.define(:result, :turn_count, :tool_calls_made, :input_tokens, :output_tokens, :cost)
 
       TurnStart = Data.define
-      TurnEnd = Data.define(:tool_results, :turn_number)
+      TurnEnd = Data.define(:tool_results, :turn_number, :input_tokens, :output_tokens, :cost)
 
       MessageStart = Data.define
       TextDelta = Data.define(:content)

data/lib/ask/agent/loop.rb

@@ -6,7 +6,7 @@ module Ask
       LOOP_DETECTION_WINDOW = 3
       @max_consecutive_tool_turns = 6
 
-      attr_reader :turn_count
+      attr_reader :turn_count, :last_input_tokens, :last_output_tokens, :last_cost
 
       def initialize(max_turns: 25, max_consecutive_tool_turns: 6)
         @max_turns = max_turns
@@ -36,6 +36,10 @@ module Ask
           end
         end
 
+        @last_input_tokens = response.input_tokens
+        @last_output_tokens = response.output_tokens
+        @last_cost = response.cost
+
         event_emitter.emit(Events::MessageEnd.new(tool_calls: response.tool_call?))
         @turn_count += 1
 
@@ -65,7 +69,13 @@ module Ask
           return "Based on my investigation: #{summary}"
         end
 
-        event_emitter.emit(Events::TurnEnd.new(tool_results: tool_results, turn_number: @turn_count))
+        event_emitter.emit(Events::TurnEnd.new(
+          tool_results: tool_results,
+          turn_number: @turn_count,
+          input_tokens: @last_input_tokens,
+          output_tokens: @last_output_tokens,
+          cost: @last_cost
+        ))
 
         if compactor && compactor.should_compact?
           compactor.run(event_emitter: event_emitter)

data/lib/ask/agent/session.rb

@@ -7,7 +7,7 @@ module Ask
   module Agent
     class Session
       attr_reader :id, :chat, :tools, :turn_count, :created_at, :messages
-      attr_reader :tool_calls_made
+      attr_reader :tool_calls_made, :total_input_tokens, :total_output_tokens, :total_cost
 
       def reflection_count
         @reflector&.reflection_count || 0
@@ -33,6 +33,10 @@ module Ask
         @created_at = Time.now
         @_no_tools_instructed = false
 
+        @total_input_tokens = 0
+        @total_output_tokens = 0
+        @total_cost = 0.0
+
         @telemetry = telemetry.is_a?(Telemetry) ? telemetry : Telemetry.new(enabled: !!telemetry)
 
         @chat = build_chat(model, system_prompt, tools, **chat_options)
@@ -99,6 +103,10 @@ module Ask
             event_emitter: self,
             session_id: @id
           )
+
+          @total_input_tokens += @loop.last_input_tokens.to_i
+          @total_output_tokens += @loop.last_output_tokens.to_i
+          @total_cost += @loop.last_cost.to_f
         rescue MaxTurnsExceeded => e
           emit(Events::MaxTurnsExceeded.new(max_turns: @max_turns))
           @telemetry.log(:max_turns_exceeded, session_id: @id, max_turns: @max_turns)
@@ -143,6 +151,10 @@ module Ask
               event_emitter: self,
               session_id: @id
             )
+
+            @total_input_tokens += @loop.last_input_tokens.to_i
+            @total_output_tokens += @loop.last_output_tokens.to_i
+            @total_cost += @loop.last_cost.to_f
           end
         end
 
@@ -151,7 +163,14 @@ module Ask
           try_auto_meta_agent
         end
 
-        emit(Events::SessionEnd.new(result: response, turn_count: @turn_count, tool_calls_made: @tool_calls_made))
+        emit(Events::SessionEnd.new(
+          result: response,
+          turn_count: @turn_count,
+          tool_calls_made: @tool_calls_made,
+          input_tokens: @total_input_tokens,
+          output_tokens: @total_output_tokens,
+          cost: @total_cost
+        ))
         @messages = @chat.messages.dup
 
         response

data/lib/ask/agent/tool_executor.rb

@@ -182,10 +182,15 @@ module Ask
       end
 
       def retryable_error_name?(error_name)
-        retryable = %w[Timeout::Error Errno::ETIMEDOUT
-                       Ask::RateLimitError Ask::ServerError
-                       Ask::RateLimitError Ask::ServiceUnavailableError]
-        retryable.include?(error_name)
+        return false unless error_name
+
+        klass = Object.const_get(error_name) rescue nil
+        return false unless klass
+
+        klass <= Ask::RateLimitError ||
+        klass <= Ask::ServerError ||
+        klass <= Ask::ServiceUnavailable ||
+        %w[Timeout::Error Errno::ETIMEDOUT].include?(error_name)
       end
 
       def critical_error?(error_class_name)

data/lib/ask/agent/version.rb

@@ -2,6 +2,6 @@
 
 module Ask
   module Agent
-    VERSION = "0.2.2"
+    VERSION = "0.3.1"
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ask-agent
 version: !ruby/object:Gem::Version
-  version: 0.2.2
+  version: 0.3.1
 platform: ruby
 authors:
 - Kaka Ruto
@@ -65,6 +65,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0.1'
+- !ruby/object:Gem::Dependency
+  name: ask-instrumentation
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.1'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.1'
 - !ruby/object:Gem::Dependency
   name: minitest
   requirement: !ruby/object:Gem::Requirement