rasti-ai 2.0.1 → 2.0.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 6fb23b7f32a6b10d293df78ff4b805156f9fb9145b371a4b7fa43c6332acdd5f
-  data.tar.gz: '08adffd14dd7b7b183f80bf7b36b3e068690cf1dbe805970e01fef6e4bd062f4'
+  metadata.gz: 47f94217b69770cc1a2238f28995a4b8cda405554134d0ba7a9dc8a6418080c7
+  data.tar.gz: fc5522ee0c3ba8bc03f71c50f2c0853be9c2f7c992db7c6a65856ea555b79107
 SHA512:
-  metadata.gz: eae7bda6a78593f178a72fa95b587137fd9399eac4b3ad4959cd6305ce5661f4a37ca165e4f6cd20ad139b01bd7843eb72cacce295d02b70db1dd6dfd5ffbf62
-  data.tar.gz: d9fb1f4c18d4bcd24972d53789d42a8d5ef35f2eaa8d3700743230425af1506261ce63fd76f98f612d4cb43aded7920da0a84eaa8b6c85b12bffa69d19bd05f4
+  metadata.gz: 5e6ecfcb985c408abc48c12e89ec3b106056d0e8b03958a02fe3a0ee18418a0260cc6a8d736f25e82d15f9c454aa3983b322bafec69122e24149b7d4d90c2f08
+  data.tar.gz: 966afa955522d3aa598a047abef9457e0dac1e9636c80a3c1df2c0509c8343436ed4d59b585a9587cda9bce96848d6e06adf57550259525cc37f0a92e4470baf

data/lib/rasti/ai/assistant.rb

@@ -4,7 +4,7 @@ module Rasti
 
       attr_reader :state
 
-      def initialize(client:nil, json_schema:nil, state:nil, model:nil, tools:[], mcp_servers:{}, logger:nil, usage_tracker:nil)
+      def initialize(client:nil, json_schema:nil, state:nil, model:nil, tools:[], mcp_servers:{}, logger:nil)
         @client = client || build_default_client
         @json_schema = json_schema
         @state = state || AssistantState.new
@@ -12,7 +12,6 @@ module Rasti
         @tools = {}
         @serialized_tools = []
         @logger = logger || Rasti::AI.logger
-        @usage_tracker = usage_tracker || Rasti::AI.usage_tracker
 
         register_tools(tools)
         register_mcp_servers(mcp_servers)
@@ -23,7 +22,6 @@ module Rasti
 
         loop do
           response = request_completion
-          track_usage response
 
           tool_calls = parse_tool_calls(response)
 
@@ -47,18 +45,12 @@ module Rasti
 
       private
 
-      attr_reader :client, :json_schema, :model, :tools, :serialized_tools, :logger, :usage_tracker
+      attr_reader :client, :json_schema, :model, :tools, :serialized_tools, :logger
 
       def messages
         state.messages
       end
 
-      def track_usage(response)
-        return unless usage_tracker
-        usage = parse_usage response
-        usage_tracker.call usage if usage
-      end
-
       # --- Shared behavior ---
 
       def register_tools(tools)
@@ -140,10 +132,6 @@ module Rasti
         raise NotImplementedError
       end
 
-      def parse_usage(response)
-        raise NotImplementedError
-      end
-
       def extract_tool_call_info(tool_call)
         raise NotImplementedError
       end

data/lib/rasti/ai/client.rb

@@ -4,17 +4,28 @@ module Rasti
 
       RETRYABLE_STATUS_CODES = [502, 503, 504].freeze
 
-      def initialize(api_key:nil, logger:nil, http_connect_timeout:nil, http_read_timeout:nil, http_max_retries:nil)
+      def initialize(api_key:nil, logger:nil, http_connect_timeout:nil, http_read_timeout:nil, http_max_retries:nil, usage_tracker:nil)
         @api_key = api_key || default_api_key
         @logger = logger || Rasti::AI.logger
         @http_connect_timeout = http_connect_timeout || Rasti::AI.http_connect_timeout
         @http_read_timeout = http_read_timeout || Rasti::AI.http_read_timeout
         @http_max_retries = http_max_retries || Rasti::AI.http_max_retries
+        @usage_tracker = usage_tracker || Rasti::AI.usage_tracker
       end
 
       private
 
-      attr_reader :api_key, :logger, :http_connect_timeout, :http_read_timeout, :http_max_retries
+      attr_reader :api_key, :logger, :http_connect_timeout, :http_read_timeout, :http_max_retries, :usage_tracker
+
+      def track_usage(response)
+        return unless usage_tracker
+        usage = parse_usage response
+        usage_tracker.call usage if usage
+      end
+
+      def parse_usage(response)
+        raise NotImplementedError
+      end
 
       def default_api_key
         raise NotImplementedError
@@ -63,7 +74,9 @@ module Rasti
             raise Errors::RequestFail.new(url, body, response)
           end
 
-          JSON.parse response.body
+          parsed_response = JSON.parse response.body
+          track_usage parsed_response
+          parsed_response
 
         rescue SocketError, Net::OpenTimeout, Net::ReadTimeout, Errors::RequestFail => e
           if retry_count < max_retries

data/lib/rasti/ai/gemini/assistant.rb

@@ -60,19 +60,6 @@ module Rasti
           !response.dig('candidates', 0, 'finishReason').nil?
         end
 
-        def parse_usage(response)
-          usage = response['usageMetadata']
-          return unless usage
-          Usage.new(
-            provider: :gemini,
-            model: response['modelVersion'],
-            input_tokens: usage['promptTokenCount'],
-            output_tokens: usage['candidatesTokenCount'],
-            cached_tokens: usage['cachedContentTokenCount'] || 0,
-            reasoning_tokens: usage['thoughtsTokenCount'] || 0
-          )
-        end
-
         def extract_tool_call_info(tool_call)
           fc = tool_call['functionCall']
           [fc['name'], fc['args'] || {}]

data/lib/rasti/ai/gemini/client.rb

@@ -17,6 +17,20 @@ module Rasti
 
         private
 
+        def parse_usage(response)
+          usage = response['usageMetadata']
+          return unless usage
+          Usage.new(
+            provider: 'gemini',
+            model: response['modelVersion'],
+            input_tokens: usage['promptTokenCount'],
+            output_tokens: usage['candidatesTokenCount'],
+            cached_tokens: usage['cachedContentTokenCount'] || 0,
+            reasoning_tokens: usage['thoughtsTokenCount'] || 0,
+            raw: usage
+          )
+        end
+
         def default_api_key
           Rasti::AI.gemini_api_key
         end

data/lib/rasti/ai/open_ai/assistant.rb

@@ -51,19 +51,6 @@ module Rasti
           !response.dig('choices', 0, 'finish_reason').nil?
         end
 
-        def parse_usage(response)
-          usage = response['usage']
-          return unless usage
-          Usage.new(
-            provider: :open_ai,
-            model: response['model'],
-            input_tokens: usage['prompt_tokens'],
-            output_tokens: usage['completion_tokens'],
-            cached_tokens: usage.dig('prompt_tokens_details', 'cached_tokens') || 0,
-            reasoning_tokens: usage.dig('completion_tokens_details', 'reasoning_tokens') || 0
-          )
-        end
-
         def extract_tool_call_info(tool_call)
           name = tool_call['function']['name']
           args = JSON.parse(tool_call['function']['arguments'])

data/lib/rasti/ai/open_ai/client.rb

@@ -18,6 +18,20 @@ module Rasti
 
         private
 
+        def parse_usage(response)
+          usage = response['usage']
+          return unless usage
+          Usage.new(
+            provider: 'open_ai',
+            model: response['model'],
+            input_tokens: usage['prompt_tokens'],
+            output_tokens: usage['completion_tokens'],
+            cached_tokens: usage.dig('prompt_tokens_details', 'cached_tokens') || 0,
+            reasoning_tokens: usage.dig('completion_tokens_details', 'reasoning_tokens') || 0,
+            raw: usage
+          )
+        end
+
         def default_api_key
           Rasti::AI.openai_api_key
         end

data/lib/rasti/ai/usage.rb

@@ -2,12 +2,13 @@ module Rasti
   module AI
     class Usage < Rasti::Model
 
-      attribute :provider,         Rasti::Types::Symbol
+      attribute :provider,         Rasti::Types::String
       attribute :model,            Rasti::Types::String
       attribute :input_tokens,     Rasti::Types::Integer
       attribute :output_tokens,    Rasti::Types::Integer
       attribute :cached_tokens,    Rasti::Types::Integer
       attribute :reasoning_tokens, Rasti::Types::Integer
+      attribute :raw
 
     end
   end

data/lib/rasti/ai/version.rb

@@ -1,5 +1,5 @@
 module Rasti
   module AI
-    VERSION = '2.0.1'
+    VERSION = '2.0.2'
   end
 end

data/spec/gemini/assistant_spec.rb

@@ -132,73 +132,7 @@ describe Rasti::AI::Gemini::Assistant do
 
   end
 
-  describe 'Usage tracker' do
 
-    it 'Track usage' do
-      stub_gemini_generate_content question: question, answer: answer
-
-      tracked = []
-      tracker = ->(usage) { tracked << usage }
-
-      assistant = Rasti::AI::Gemini::Assistant.new usage_tracker: tracker
-
-      assistant.call question
-
-      assert_equal 1, tracked.count
-
-      usage = tracked[0]
-      assert_instance_of Rasti::AI::Usage, usage
-      assert_equal :gemini, usage.provider
-      assert_equal 'gemini-test', usage.model
-      assert_equal 10, usage.input_tokens
-      assert_equal 50, usage.output_tokens
-      assert_equal 0, usage.cached_tokens
-      assert_equal 0, usage.reasoning_tokens
-    end
-
-    it 'Track usage with tool calls' do
-      client = Minitest::Mock.new
-
-      tool_response = read_json_resource(
-        'gemini/tool_response.json',
-        name: 'goals_by_player',
-        arguments: {player: 'Lionel Messi', team: 'Barcelona'}
-      )
-
-      basic_resp = read_json_resource('gemini/basic_response.json', content: answer)
-
-      client.expect :generate_content, tool_response do |params| true end
-      client.expect :generate_content, basic_resp do |params| true end
-
-      tool = GoalsByPlayer.new
-
-      tracked = []
-      tracker = ->(usage) { tracked << usage }
-
-      assistant = Rasti::AI::Gemini::Assistant.new client: client, tools: [tool], usage_tracker: tracker
-
-      assistant.call question
-
-      assert_equal 2, tracked.count
-      assert_equal 20, tracked[0].input_tokens
-      assert_equal 10, tracked[0].output_tokens
-      assert_equal 10, tracked[1].input_tokens
-      assert_equal 50, tracked[1].output_tokens
-
-      client.verify
-    end
-
-    it 'Without tracker' do
-      stub_gemini_generate_content question: question, answer: answer
-
-      assistant = Rasti::AI::Gemini::Assistant.new
-
-      response = assistant.call question
-
-      assert_equal answer, response
-    end
-
-  end
 
   describe 'Tools' do
 

data/spec/gemini/client_spec.rb

@@ -83,6 +83,56 @@ describe Rasti::AI::Gemini::Client do
       refute_empty log_output.string
     end
 
+    describe 'Usage tracker' do
+
+      it 'Track usage' do
+        stub_gemini_generate_content
+
+        tracked = []
+        tracker = ->(usage) { tracked << usage }
+
+        client = Rasti::AI::Gemini::Client.new usage_tracker: tracker
+
+        client.generate_content contents: [user_content(question)]
+
+        assert_equal 1, tracked.count
+
+        expected_raw = {
+          'promptTokenCount' => 4,
+          'candidatesTokenCount' => 18,
+          'totalTokenCount' => 275,
+          'promptTokensDetails' => [
+            {
+              'modality' => 'TEXT',
+              'tokenCount' => 4
+            }
+          ],
+          'thoughtsTokenCount' => 253
+        }
+
+        usage = tracked[0]
+        assert_instance_of Rasti::AI::Usage, usage
+        assert_equal 'gemini', usage.provider
+        assert_equal 'gemini-test', usage.model
+        assert_equal 4, usage.input_tokens
+        assert_equal 18, usage.output_tokens
+        assert_equal 0, usage.cached_tokens
+        assert_equal 253, usage.reasoning_tokens
+        assert_equal expected_raw, usage.raw
+      end
+
+      it 'Without tracker' do
+        stub_gemini_generate_content
+
+        client = Rasti::AI::Gemini::Client.new
+
+        response = client.generate_content contents: [user_content(question)]
+
+        assert_response_content response, answer
+      end
+
+    end
+
   end
 
   it 'Request error' do

data/spec/open_ai/assistant_spec.rb

@@ -126,73 +126,7 @@ describe Rasti::AI::OpenAI::Assistant do
 
   end
 
-  describe 'Usage tracker' do
 
-    it 'Track usage' do
-      stub_open_ai_chat_completions question: question, answer: answer
-
-      tracked = []
-      tracker = ->(usage) { tracked << usage }
-
-      assistant = Rasti::AI::OpenAI::Assistant.new usage_tracker: tracker
-
-      assistant.call question
-
-      assert_equal 1, tracked.count
-
-      usage = tracked[0]
-      assert_instance_of Rasti::AI::Usage, usage
-      assert_equal :open_ai, usage.provider
-      assert_equal 'gpt-4o-mini-2024-07-18', usage.model
-      assert_equal 27, usage.input_tokens
-      assert_equal 229, usage.output_tokens
-      assert_equal 0, usage.cached_tokens
-      assert_equal 0, usage.reasoning_tokens
-    end
-
-    it 'Track usage with tool calls' do
-      client = Minitest::Mock.new
-
-      tool_response = read_json_resource(
-        'open_ai/tool_response.json',
-        name: 'goals_by_player',
-        arguments: {player: 'Lionel Messi', team: 'Barcelona'}
-      )
-
-      basic_resp = read_json_resource('open_ai/basic_response.json', content: answer)
-
-      client.expect :chat_completions, tool_response do |params| true end
-      client.expect :chat_completions, basic_resp do |params| true end
-
-      tool = GoalsByPlayer.new
-
-      tracked = []
-      tracker = ->(usage) { tracked << usage }
-
-      assistant = Rasti::AI::OpenAI::Assistant.new client: client, tools: [tool], usage_tracker: tracker
-
-      assistant.call question
-
-      assert_equal 2, tracked.count
-      assert_equal 83, tracked[0].input_tokens
-      assert_equal 23, tracked[0].output_tokens
-      assert_equal 27, tracked[1].input_tokens
-      assert_equal 229, tracked[1].output_tokens
-
-      client.verify
-    end
-
-    it 'Without tracker' do
-      stub_open_ai_chat_completions question: question, answer: answer
-
-      assistant = Rasti::AI::OpenAI::Assistant.new
-
-      response = assistant.call question
-
-      assert_equal answer, response
-    end
-
-  end
 
   describe 'Tools' do
 

data/spec/open_ai/client_spec.rb

@@ -83,6 +83,59 @@ describe Rasti::AI::OpenAI::Client do
       refute_empty log_output.string
     end
 
+    describe 'Usage tracker' do
+
+      it 'Track usage' do
+        stub_open_ai_chat_completions
+
+        tracked = []
+        tracker = ->(usage) { tracked << usage }
+
+        client = Rasti::AI::OpenAI::Client.new usage_tracker: tracker
+
+        client.chat_completions messages: [user_message(question)]
+
+        assert_equal 1, tracked.count
+
+        expected_raw = {
+          'prompt_tokens' => 27,
+          'completion_tokens' => 229,
+          'total_tokens' => 256,
+          'prompt_tokens_details' => {
+            'cached_tokens' => 0,
+            'audio_tokens' => 0
+          },
+          'completion_tokens_details' => {
+            'reasoning_tokens' => 0,
+            'audio_tokens' => 0,
+            'accepted_prediction_tokens' => 0,
+            'rejected_prediction_tokens' => 0
+          }
+        }
+
+        usage = tracked[0]
+        assert_instance_of Rasti::AI::Usage, usage
+        assert_equal 'open_ai', usage.provider
+        assert_equal 'gpt-4o-mini-2024-07-18', usage.model
+        assert_equal 27, usage.input_tokens
+        assert_equal 229, usage.output_tokens
+        assert_equal 0, usage.cached_tokens
+        assert_equal 0, usage.reasoning_tokens
+        assert_equal expected_raw, usage.raw
+      end
+
+      it 'Without tracker' do
+        stub_open_ai_chat_completions
+
+        client = Rasti::AI::OpenAI::Client.new
+
+        response = client.chat_completions messages: [user_message(question)]
+
+        assert_response_content response, answer
+      end
+
+    end
+
   end
 
   it 'Request error' do

data/spec/resources/gemini/basic_response.json

@@ -14,9 +14,16 @@
     }
   ],
   "usageMetadata": {
-    "promptTokenCount": 10,
-    "candidatesTokenCount": 50,
-    "totalTokenCount": 60
+    "promptTokenCount": 4,
+    "candidatesTokenCount": 18,
+    "totalTokenCount": 275,
+    "promptTokensDetails": [
+      {
+        "modality": "TEXT",
+        "tokenCount": 4
+      }
+    ],
+    "thoughtsTokenCount": 253
   },
   "modelVersion": "gemini-test"
 }

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rasti-ai
 version: !ruby/object:Gem::Version
-  version: 2.0.1
+  version: 2.0.2
 platform: ruby
 authors:
 - Gabriel Naiman
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2026-03-11 00:00:00.000000000 Z
+date: 2026-03-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: multi_require