rasti-ai 2.0.1 → 3.0.0

data/spec/anthropic/assistant_spec.rb

@@ -0,0 +1,349 @@
+require 'minitest_helper'
+
+describe Rasti::AI::Anthropic::Assistant do
+
+  let(:api_url) { 'https://api.anthropic.com/v1/messages' }
+
+  let(:question) { 'How many goals has Messi scored for Barca?' }
+
+  let(:answer) { 'Lionel Messi scored 672 goals in 778 official matches for FC Barcelona.' }
+
+  def stub_anthropic_messages(question:, answer:, model:nil, json_schema:nil)
+    model ||= Rasti::AI.anthropic_default_model
+
+    body = read_json_resource('anthropic/basic_request.json', model: model, prompt: question)
+
+    if json_schema
+      body['tools'] = [{
+        'name'         => 'structured_output',
+        'description'  => 'Return the structured response',
+        'input_schema' => {'type' => 'object', 'properties' => json_schema}
+      }]
+      body['tool_choice'] = {'type' => 'tool', 'name' => 'structured_output'}
+
+      stub_request(:post, api_url)
+        .with(body: JSON.dump(body))
+        .to_return(body: read_resource('anthropic/tool_response.json', name: 'structured_output', arguments: json_schema))
+    else
+      stub_request(:post, api_url)
+        .with(body: JSON.dump(body))
+        .to_return(body: read_resource('anthropic/basic_response.json', content: answer))
+    end
+  end
+
+  it 'Default' do
+    stub_anthropic_messages question: question, answer: answer
+
+    assistant = Rasti::AI::Anthropic::Assistant.new
+
+    response = assistant.call question
+
+    assert_equal answer, response
+  end
+
+  describe 'Customized' do
+
+    it 'Client' do
+      client_arguments = [
+        {
+          model:       nil,
+          system:      nil,
+          tools:       [],
+          tool_choice: nil,
+          thinking:    nil,
+          messages: [
+            {
+              role:    Rasti::AI::Anthropic::Roles::USER,
+              content: question
+            }
+          ]
+        }
+      ]
+
+      client_response = read_json_resource 'anthropic/basic_response.json', content: answer
+
+      client = Minitest::Mock.new
+      client.expect :messages, client_response, client_arguments
+
+      assistant = Rasti::AI::Anthropic::Assistant.new client: client
+
+      response = assistant.call question
+
+      assert_equal answer, response
+
+      client.verify
+    end
+
+    it 'State' do
+      context = 'Act as sports journalist'
+      state   = Rasti::AI::AssistantState.new context: context
+
+      request_body = {
+        model:     Rasti::AI.anthropic_default_model,
+        max_tokens: 4096,
+        messages: [
+          {
+            role:    Rasti::AI::Anthropic::Roles::USER,
+            content: question
+          }
+        ],
+        system: context
+      }
+
+      stub_request(:post, api_url)
+        .with(body: JSON.dump(request_body))
+        .to_return(body: read_resource('anthropic/basic_response.json', content: answer))
+
+      assistant = Rasti::AI::Anthropic::Assistant.new state: state
+
+      response = assistant.call question
+
+      expected_assistant_message = {
+        role:    Rasti::AI::Anthropic::Roles::ASSISTANT,
+        content: answer
+      }
+
+      assert_equal answer, response
+      assert_equal 2, state.messages.count
+      assert_equal expected_assistant_message, state.messages.last
+    end
+
+    it 'Model' do
+      model = SecureRandom.uuid
+
+      stub_anthropic_messages question: question, answer: answer, model: model
+
+      assistant = Rasti::AI::Anthropic::Assistant.new model: model
+
+      response = assistant.call question
+
+      assert_equal answer, response
+    end
+
+    it 'Thinking' do
+      body = read_json_resource('anthropic/basic_request.json', model: Rasti::AI.anthropic_default_model, prompt: question)
+      body['thinking'] = {'type' => 'enabled', 'budget_tokens' => 8_000}
+
+      stub_request(:post, api_url)
+        .with(body: JSON.dump(body))
+        .to_return(body: read_resource('anthropic/basic_response.json', content: answer))
+
+      assistant = Rasti::AI::Anthropic::Assistant.new thinking: 'medium'
+
+      response = assistant.call question
+
+      assert_equal answer, response
+    end
+
+    it 'JSON Schema' do
+      json_schema = {'answer' => 'Response answer'}
+
+      stub_anthropic_messages question: question, answer: answer, json_schema: json_schema
+
+      assistant = Rasti::AI::Anthropic::Assistant.new json_schema: json_schema
+
+      response = assistant.call question
+
+      assert_equal 'Response answer', JSON.parse(response)['answer']
+    end
+
+  end
+
+
+
+  describe 'Tools' do
+
+    let(:client) { Minitest::Mock.new }
+
+    let(:tool_response) do
+      read_json_resource(
+        'anthropic/tool_response.json',
+        name:      'goals_by_player',
+        arguments: {
+          player: 'Lionel Messi',
+          team:   'Barcelona'
+        }
+      )
+    end
+
+    let(:tool_result) { '672' }
+
+    let(:error_message) { 'There was an error using a tool' }
+
+    def basic_response(content)
+      read_json_resource 'anthropic/basic_response.json', content: content
+    end
+
+    def stub_client_request(role:, content:, response:, tools:[])
+      serialized_tools = tools.map do |tool|
+        raw = Rasti::AI::ToolSerializer.serialize(tool.class)
+        result = raw.dup
+        if result.key?(:inputSchema)
+          result[:input_schema] = result.delete(:inputSchema)
+        end
+        result
+      end
+
+      client.expect :messages, response do |params|
+        last_message = params[:messages].last
+        last_message[:role]    == role &&
+          last_message[:content] == content &&
+          params[:tools]         == serialized_tools
+      end
+    end
+
+    it 'Call function' do
+      tool = GoalsByPlayer.new
+
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  question,
+                          tools:    [tool],
+                          response: tool_response
+
+      expected_tool_result_content = [{
+        type:        'tool_result',
+        tool_use_id: 'toolu_01A09q90qw90lq917835lq9',
+        content:     tool_result
+      }]
+
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  expected_tool_result_content,
+                          tools:    [tool],
+                          response: basic_response(answer)
+
+      assistant = Rasti::AI::Anthropic::Assistant.new client: client, tools: [tool]
+
+      response = assistant.call question
+
+      assert_equal answer, response
+
+      client.verify
+    end
+
+    it 'Tool failure' do
+      tool = GoalsByPlayer.new
+      tool.define_singleton_method :call do |*args|
+        raise 'Broken tool'
+      end
+
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  question,
+                          tools:    [tool],
+                          response: tool_response
+
+      expected_tool_result_content = [{
+        type:        'tool_result',
+        tool_use_id: 'toolu_01A09q90qw90lq917835lq9',
+        content:     'Error: Broken tool'
+      }]
+
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  expected_tool_result_content,
+                          tools:    [tool],
+                          response: basic_response(error_message)
+
+      assistant = Rasti::AI::Anthropic::Assistant.new client: client, tools: [tool]
+
+      response = assistant.call question
+
+      assert_equal error_message, response
+
+      client.verify
+    end
+
+    it 'Undefined tool' do
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  question,
+                          response: tool_response
+
+      expected_tool_result_content = [{
+        type:        'tool_result',
+        tool_use_id: 'toolu_01A09q90qw90lq917835lq9',
+        content:     'Error: Undefined tool goals_by_player'
+      }]
+
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  expected_tool_result_content,
+                          response: basic_response(error_message)
+
+      assistant = Rasti::AI::Anthropic::Assistant.new client: client, tools: []
+
+      response = assistant.call question
+
+      assert_equal error_message, response
+
+      client.verify
+    end
+
+    it 'Cached result' do
+      mock = Minitest::Mock.new
+      mock.expect :call, tool_result, [{'player' => 'Lionel Messi', 'team' => 'Barcelona'}]
+
+      tool = GoalsByPlayer.new
+      tool.define_singleton_method :call do |*args|
+        mock.call(*args)
+      end
+
+      expected_tool_result_content = [{
+        type:        'tool_result',
+        tool_use_id: 'toolu_01A09q90qw90lq917835lq9',
+        content:     tool_result
+      }]
+
+      assistant = Rasti::AI::Anthropic::Assistant.new client: client, tools: [tool]
+
+      5.times do
+        stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                            content:  question,
+                            tools:    [tool],
+                            response: tool_response
+
+        stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                            content:  expected_tool_result_content,
+                            tools:    [tool],
+                            response: basic_response(answer)
+
+        response = assistant.call question
+
+        assert_equal answer, response
+      end
+
+      client.verify
+    end
+
+    it 'Custom logger' do
+      log_output = StringIO.new
+      logger     = Logger.new log_output
+
+      tool = GoalsByPlayer.new
+
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  question,
+                          tools:    [tool],
+                          response: tool_response
+
+      expected_tool_result_content = [{
+        type:        'tool_result',
+        tool_use_id: 'toolu_01A09q90qw90lq917835lq9',
+        content:     tool_result
+      }]
+
+      stub_client_request role:     Rasti::AI::Anthropic::Roles::USER,
+                          content:  expected_tool_result_content,
+                          tools:    [tool],
+                          response: basic_response(answer)
+
+      assistant = Rasti::AI::Anthropic::Assistant.new client: client, tools: [tool], logger: logger
+
+      response = assistant.call question
+
+      assert_equal answer, response
+
+      refute_empty log_output.string
+
+      client.verify
+    end
+
+  end
+
+end

data/spec/anthropic/client_spec.rb

@@ -0,0 +1,203 @@
+require 'minitest_helper'
+
+describe Rasti::AI::Anthropic::Client do
+
+  let(:api_url) { 'https://api.anthropic.com/v1/messages' }
+
+  def user_message(content)
+    {
+      role: Rasti::AI::Anthropic::Roles::USER,
+      content: content
+    }
+  end
+
+  describe 'Basic message' do
+
+    let(:question) { 'who is Messi?' }
+
+    let(:answer) { 'Lionel Messi is the best player ever' }
+
+    def stub_anthropic_messages(api_key:nil, model:nil)
+      api_key ||= Rasti::AI.anthropic_api_key
+      model   ||= Rasti::AI.anthropic_default_model
+
+      stub_request(:post, api_url)
+        .with(
+          headers: {
+            'x-api-key'          => api_key,
+            'anthropic-version'  => '2023-06-01'
+          },
+          body: read_resource('anthropic/basic_request.json', model: model, prompt: question)
+        )
+        .to_return(body: read_resource('anthropic/basic_response.json', content: answer))
+    end
+
+    def assert_response_content(response, expected_content)
+      text_block = response['content'].find { |b| b['type'] == 'text' }
+      assert_equal expected_content, text_block['text']
+    end
+
+    it 'Default API key, model and logger' do
+      stub_anthropic_messages
+
+      client = Rasti::AI::Anthropic::Client.new
+
+      response = client.messages messages: [user_message(question)]
+
+      assert_response_content response, answer
+    end
+
+    it 'Custom API key' do
+      custom_api_key = SecureRandom.uuid
+
+      stub_anthropic_messages api_key: custom_api_key
+
+      client = Rasti::AI::Anthropic::Client.new api_key: custom_api_key
+
+      response = client.messages messages: [user_message(question)]
+
+      assert_response_content response, answer
+    end
+
+    it 'Custom model' do
+      custom_model = SecureRandom.uuid
+
+      stub_anthropic_messages model: custom_model
+
+      client = Rasti::AI::Anthropic::Client.new
+
+      response = client.messages messages: [user_message(question)], model: custom_model
+
+      assert_response_content response, answer
+    end
+
+    it 'Custom logger' do
+      log_output = StringIO.new
+      logger = Logger.new log_output
+
+      stub_anthropic_messages
+
+      client = Rasti::AI::Anthropic::Client.new logger: logger
+
+      response = client.messages messages: [user_message(question)]
+
+      assert_response_content response, answer
+
+      refute_empty log_output.string
+    end
+
+    describe 'Usage tracker' do
+
+      it 'Track usage' do
+        stub_anthropic_messages
+
+        tracked = []
+        tracker = ->(usage) { tracked << usage }
+
+        client = Rasti::AI::Anthropic::Client.new usage_tracker: tracker
+
+        client.messages messages: [user_message(question)]
+
+        assert_equal 1, tracked.count
+
+        expected_raw = {
+          'input_tokens'                => 25,
+          'output_tokens'               => 11,
+          'cache_creation_input_tokens' => 0,
+          'cache_read_input_tokens'     => 0
+        }
+
+        usage = tracked[0]
+        assert_instance_of Rasti::AI::Usage, usage
+        assert_equal 'anthropic',   usage.provider
+        assert_equal 'claude-test', usage.model
+        assert_equal 25,            usage.input_tokens
+        assert_equal 11,            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_anthropic_messages
+
+        client = Rasti::AI::Anthropic::Client.new
+
+        response = client.messages messages: [user_message(question)]
+
+        assert_response_content response, answer
+      end
+
+    end
+
+  end
+
+  it 'Request error' do
+    stub_request(:post, api_url)
+      .to_return(status: 400, body: '{"type":"error","error":{"type":"invalid_request_error","message":"Test error"}}')
+
+    client = Rasti::AI::Anthropic::Client.new
+
+    error = assert_raises(Rasti::AI::Errors::RequestFail) do
+      client.messages messages: ['invalid message']
+    end
+
+    assert_includes error.message, 'Response: 400'
+  end
+
+  it 'Tool call' do
+    question  = 'how many goals did messi for barca'
+    tool_name = 'player_goals'
+
+    tool = {
+      name:         tool_name,
+      description:  'Gets the number of goals scored by a player for a specific team',
+      input_schema: {
+        type:       'object',
+        properties: {
+          name: {
+            type:        'string',
+            description: 'Full name of the player'
+          },
+          team: {
+            type:        'string',
+            description: 'Name of the team the player was part of'
+          }
+        },
+        required: ['name', 'team']
+      }
+    }
+
+    arguments = {
+      name: 'Lionel Messi',
+      team: 'FC Barcelona'
+    }
+
+    stub_request(:post, api_url)
+      .with(
+        headers: {
+          'x-api-key'         => Rasti::AI.anthropic_api_key,
+          'anthropic-version' => '2023-06-01'
+        },
+        body: read_resource(
+          'anthropic/tool_request.json',
+          model:   Rasti::AI.anthropic_default_model,
+          prompt:  question,
+          tools:   [tool]
+        )
+      )
+      .to_return(body: read_resource('anthropic/tool_response.json', name: tool_name, arguments: arguments))
+
+    client = Rasti::AI::Anthropic::Client.new
+
+    response = client.messages messages:    [{role: 'user', content: question}],
+                                tools:       [tool],
+                                tool_choice: {type: 'auto'}
+
+    tool_use = response['content'].find { |b| b['type'] == 'tool_use' }
+
+    assert_equal tool_name,                          tool_use['name']
+    assert_equal({'name' => 'Lionel Messi', 'team' => 'FC Barcelona'}, tool_use['input'])
+  end
+
+end

data/spec/gemini/assistant_spec.rb

@@ -110,6 +110,21 @@ describe Rasti::AI::Gemini::Assistant do
       assert_equal answer, response
     end
 
+    it 'Thinking' do
+      body = read_json_resource('gemini/basic_request.json', prompt: question)
+      body['generation_config'] = {'thinking_config' => {'thinking_budget' => 8_192}}
+
+      stub_request(:post, api_url)
+        .with(body: JSON.dump(body))
+        .to_return(body: read_resource('gemini/basic_response.json', content: answer))
+
+      assistant = Rasti::AI::Gemini::Assistant.new thinking: 'medium'
+
+      response = assistant.call question
+
+      assert_equal answer, response
+    end
+
     it 'JSON Schema' do
       json_schema = {answer: 'Response answer'}
       json_answer = "{\\\"answer\\\": \\\"#{answer}\\\"}"
@@ -132,73 +147,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