rasti-ai 3.1.0 → 3.2.0

data/README.md

@@ -53,6 +53,9 @@ Rasti::AI.configure do |config|
   config.huawei_maas_api_key = 'sk-maas-12345' # Default ENV['HUAWEI_MAAS_API_KEY']
   config.huawei_maas_default_model = 'DeepSeek-V3' # Default ENV['HUAWEI_MAAS_DEFAULT_MODEL']
 
+  # Default provider used by Rasti::AI.create_assistant and Rasti::AI.generate_text
+  config.default_provider = :open_ai # Default ENV['AI_DEFAULT_PROVIDER']
+
   # Usage tracking
   config.usage_tracker = ->(usage) { puts "#{usage.provider}: #{usage.input_tokens} in / #{usage.output_tokens} out" }
 end
@@ -70,10 +73,132 @@ All providers share the same interface. The examples below use OpenAI, but apply
 
 OpenRouter and Huawei MaaS are [OpenAI-compatible](https://platform.openai.com/docs/api-reference/chat) APIs — they use the same chat completions format, tool calling, and Bearer token auth. Their clients and assistants inherit from the OpenAI provider, so everything in the examples below works identically.
 
+### Choosing provider and model at runtime
+
+Instead of referencing a provider-specific class, the provider can be a parameter or come from the configuration. This makes swapping models a config change instead of a code change.
+
+```ruby
+# One-shot text generation
+result = Rasti::AI.generate_text prompt: 'who is the best player', provider: :open_ai
+
+result.content # => 'The best player is Lionel Messi'
+result.usage   # => Rasti::AI::Usage
+result.raw     # => raw provider response
+
+# Provider and model
+Rasti::AI.generate_text prompt: 'who is the best player', provider: :anthropic, model: 'claude-sonnet-4-5'
+
+# Both in a single string — handy for ENV vars
+Rasti::AI.generate_text prompt: 'who is the best player', model: ENV['AI_MODEL'] # 'gemini:gemini-2.0-flash'
+
+# Provider from configuration (Rasti::AI.default_provider)
+Rasti::AI.generate_text prompt: 'who is the best player'
+```
+
+The response is normalized across providers: `content` is the text the model returned (or the JSON string when `json_schema:` is used), `usage` is a `Rasti::AI::Usage`, and `raw` is the untouched provider payload as an escape hatch.
+
+Instead of a single `prompt:`, a conversation can be passed with `messages:` using generic roles — each provider translates them to its own format (`assistant` becomes `model` on Gemini, `system` becomes `system_instruction`, etc.):
+
+```ruby
+Rasti::AI.generate_text provider: :gemini, messages: [
+  {role: Rasti::AI::Roles::SYSTEM, content: 'Act as sports journalist'},
+  {role: Rasti::AI::Roles::USER, content: 'who is the best player'},
+  {role: Rasti::AI::Roles::ASSISTANT, content: 'Lionel Messi'},
+  {role: Rasti::AI::Roles::USER, content: 'how many goals did he score?'}
+]
+```
+
+`generate_text` is a single request without tools. For tool calling, multi-turn conversations and stateful interactions use an assistant.
+
+### Assistants — tool calling and multi-turn conversations
+
+`Rasti::AI.create_assistant` is the primary entry point for assistants. It resolves the provider at runtime — just like `generate_text` — and returns the provider-specific `Assistant` instance, so swapping models is a config change instead of a code change. Everything documented below (tools, state, structured responses, thinking, MCP) works the same across all providers:
+
+```ruby
+assistant = Rasti::AI.create_assistant provider: :open_ai,
+                                       model: 'gpt-4o',
+                                       system: 'Act as sports journalist',
+                                       tools: [GetCurrentWeather.new],
+                                       mcp_servers: {weather: mcp_client},
+                                       thinking: 'medium'
+
+assistant.call 'who is the best player' # => 'The best player is Lionel Messi'
+assistant # => #<Rasti::AI::OpenAI::Assistant>
+```
+
+Omit `provider:` to use `Rasti::AI.default_provider`, or pass `'provider:model'` as `model:` to set both in one string — handy for ENV-driven configuration:
+
+```ruby
+assistant = Rasti::AI.create_assistant model: ENV['AI_MODEL'], # 'anthropic:claude-sonnet-4-5'
+                                       system: 'Act as sports journalist',
+                                       tools: [GetCurrentWeather.new]
+```
+
+Both methods accept the same options:
+
+| Option | Purpose | `generate_text` | `create_assistant` |
+|---|---|---|---|
+| `provider:` | `:open_ai`, `:gemini`, `:anthropic`, `:open_router`, `:huawei_maas` (aliases: `:openai`, `:openrouter`, `:huawei`). Defaults to `Rasti::AI.default_provider` | ✓ | ✓ |
+| `model:` | Model name, or `'provider:model'` to set both. Defaults to the provider's configured default model | ✓ | ✓ |
+| `api_key:` | Overrides the configured API key — useful for multi-tenant apps | ✓ | ✓ |
+| `usage_tracker:` | Per-call usage tracker lambda | ✓ | ✓ |
+| `logger:` | Overrides the configured logger | ✓ | ✓ |
+| `http_connect_timeout:`, `http_read_timeout:`, `http_max_retries:` | Per-call HTTP settings | ✓ | ✓ |
+| `client:` | A pre-built client (skips building one from the options above) | ✓ | ✓ |
+| `json_schema:` | Schema for structured responses | ✓ | ✓ |
+| `thinking:` | `'low'`, `'medium'` or `'high'` | ✓ | ✓ |
+| `prompt:` | A single user message | ✓ | |
+| `messages:` | Conversation with generic roles — mutually exclusive with `prompt:` | ✓ | |
+| `system:` | System prompt. On `create_assistant` it's a shortcut for `state: AssistantState.new(context: ...)` | ✓ | ✓ |
+| `state:` | An existing `AssistantState` to continue a conversation | | ✓ |
+| `tools:` | Array of tool instances | | ✓ |
+| `mcp_servers:` | Hash of `{name => MCP::Client}` | | ✓ |
+
+### Provider
+
+A provider bundles everything needed to reach a model: credentials, default model, usage tracking and HTTP settings. `Rasti::AI.provider` returns a reusable instance — useful when the same credentials are used for many calls, e.g. per tenant:
+
+```ruby
+provider = Rasti::AI.provider :anthropic,
+                              model: 'claude-sonnet-4-5',
+                              api_key: tenant.anthropic_key,
+                              usage_tracker: ->(usage) { tenant.track usage },
+                              http_read_timeout: 300
+
+provider.name  # => :anthropic
+provider.model # => 'claude-sonnet-4-5'
+
+provider.generate_text prompt: 'who is the best player'          # => Rasti::AI::Result
+provider.generate_text prompt: '...', model: 'claude-opus-4-5'   # per call model override
+provider.create_assistant tools: [GetCurrentWeather.new]         # => Rasti::AI::Anthropic::Assistant
+provider.build_client                                            # => Rasti::AI::Anthropic::Client
+```
+
+The available providers are fixed: `Rasti::AI::Provider.names # => [:open_ai, :gemini, :anthropic, :open_router, :huawei_maas]`. An unknown name raises `Rasti::AI::Errors::UnknownProvider`.
+
+A provider instance can also be passed directly as `provider:` to the module-level methods:
+
+```ruby
+Rasti::AI.generate_text provider: provider, prompt: 'who is the best player'
+```
+
 ### Assistant
 
+`Rasti::AI.create_assistant` is the recommended way to build an assistant — it keeps the provider as a runtime value and shares the same options as `generate_text`:
+
 ```ruby
-assistant = Rasti::AI::OpenAI::Assistant.new
+assistant = Rasti::AI.create_assistant provider: :open_ai,
+                                       system: 'Act as sports journalist',
+                                       tools: [GetCurrentWeather.new]
+
+assistant.call 'who is the best player' # => 'The best player is Lionel Messi'
+```
+
+Provider-specific classes are also available and are fully equivalent — useful when the provider is fixed by the application:
+
+```ruby
+assistant = Rasti::AI::OpenAI::Assistant.new system: 'Act as sports journalist',
+                                             tools: [GetCurrentWeather.new]
 assistant.call 'who is the best player' # => 'The best player is Lionel Messi'
 ```
 
@@ -138,7 +263,7 @@ tools = [
   GetCurrentWeather.new
 ]
 
-assistant = Rasti::AI::OpenAI::Assistant.new tools: tools
+assistant = Rasti::AI.create_assistant provider: :open_ai, tools: tools
 
 assistant.call 'what time is it' # => 'The current time is 3:03 PM on April 28, 2025.'
 
@@ -149,7 +274,7 @@ assistant.call 'what is the weather in Buenos Aires' # => 'In Buenos Aires it is
 ```ruby
 state = Rasti::AI::AssistantState.new context: 'Act as sports journalist'
 
-assistant = Rasti::AI::OpenAI::Assistant.new state: state
+assistant = Rasti::AI.create_assistant provider: :open_ai, state: state
 
 assistant.call 'who is the best player'
 
@@ -157,14 +282,15 @@ state.context  # => 'Act as sports journalist'
 state.messages # Array of provider-specific message hashes
 ```
 
-The state keeps the conversation history, enabling multi-turn interactions. It also caches tool call results to avoid duplicate executions.
+The state keeps the conversation history, enabling multi-turn interactions. It also caches tool call results to avoid duplicate executions. The `system:` option is a shortcut for `state: AssistantState.new(context: ...)`.
 
 ### Structured responses (JSON Schema)
 ```ruby
-assistant = Rasti::AI::OpenAI::Assistant.new json_schema: {
-  player: 'string',
-  sport: 'string'
-}
+assistant = Rasti::AI.create_assistant provider: :open_ai,
+                                       json_schema: {
+                                         player: 'string',
+                                         sport: 'string'
+                                       }
 
 response = assistant.call 'who is the best player'
 JSON.parse response # => {"player" => "Lionel Messi", "sport" => "Football"}
@@ -173,7 +299,7 @@ JSON.parse response # => {"player" => "Lionel Messi", "sport" => "Football"}
 ### Custom model and client
 ```ruby
 # Override model
-assistant = Rasti::AI::OpenAI::Assistant.new model: 'gpt-4o'
+assistant = Rasti::AI.create_assistant provider: :open_ai, model: 'gpt-4o'
 
 # Custom client with per-client HTTP settings
 client = Rasti::AI::OpenAI::Client.new(
@@ -182,15 +308,15 @@ client = Rasti::AI::OpenAI::Client.new(
   http_max_retries: 5
 )
 
-assistant = Rasti::AI::OpenAI::Assistant.new client: client
+assistant = Rasti::AI.create_assistant provider: :open_ai, client: client
 
-# Anthropic client
+# Anthropic client — Claude can be slow on long responses
 client = Rasti::AI::Anthropic::Client.new(
   http_connect_timeout: 120,
-  http_read_timeout: 300  # Claude can be slow on long responses
+  http_read_timeout: 300
 )
 
-assistant = Rasti::AI::Anthropic::Assistant.new client: client
+assistant = Rasti::AI.create_assistant provider: :anthropic, client: client
 ```
 
 ### Thinking / extended reasoning
@@ -198,7 +324,7 @@ assistant = Rasti::AI::Anthropic::Assistant.new client: client
 Some providers support extended reasoning ("thinking") to improve accuracy on complex tasks. Pass `thinking:` with a level of `'low'`, `'medium'`, or `'high'` when creating an assistant:
 
 ```ruby
-assistant = Rasti::AI::Anthropic::Assistant.new thinking: 'high'
+assistant = Rasti::AI.create_assistant provider: :anthropic, thinking: 'high'
 assistant.call 'Solve this step by step: ...'
 ```
 
@@ -213,7 +339,7 @@ tracked_usage = []
 tracker = ->(usage) { tracked_usage << usage }
 
 client = Rasti::AI::OpenAI::Client.new usage_tracker: tracker
-assistant = Rasti::AI::OpenAI::Assistant.new client: client
+assistant = Rasti::AI.create_assistant provider: :open_ai, client: client
 assistant.call 'who is the best player'
 
 usage = tracked_usage.first
@@ -426,10 +552,8 @@ You can use MCP clients as tools for any assistant:
 ```ruby
 mcp_client = Rasti::AI::MCP::Client.new url: 'https://mcp.server.ai/mcp'
 
-
-assistant = Rasti::AI::OpenAI::Assistant.new(
-  mcp_servers: {my_mcp: mcp_client}
-)
+assistant = Rasti::AI.create_assistant provider: :open_ai,
+                                       mcp_servers: {my_mcp: mcp_client}
 
 # The assistant can now call tools from the MCP server
 assistant.call 'What is 5 plus 3?'

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

@@ -5,18 +5,8 @@ module Rasti
 
         ALLOWED_SCHEMA_FIELDS = %w[type description properties required enum items format nullable anyOf].freeze
 
-        THINKING_LEVELS = {
-          'low'    => {type: 'enabled', budget_tokens: 1_024}.freeze,
-          'medium' => {type: 'enabled', budget_tokens: 8_000}.freeze,
-          'high'   => {type: 'enabled', budget_tokens: 16_000}.freeze
-        }.freeze
-
         private
 
-        def build_default_client
-          Client.new
-        end
-
         def build_user_message(prompt)
           {role: Roles::USER, content: prompt}
         end
@@ -61,7 +51,7 @@ module Rasti
         end
 
         def thinking_config
-          THINKING_LEVELS[thinking]
+          Provider::THINKING_LEVELS[thinking]
         end
 
         def parse_tool_calls(response)

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

@@ -8,14 +8,14 @@ module Rasti
 
         def messages(messages:, model:nil, system:nil, tools:[], tool_choice:nil, max_tokens:nil, thinking:nil)
           body = {
-            model:      model || Rasti::AI.anthropic_default_model,
+            model: model || default_model,
             max_tokens: max_tokens || DEFAULT_MAX_TOKENS,
-            messages:   messages
+            messages: messages
           }
 
-          body[:thinking]    = thinking    if thinking
-          body[:system]      = system      if system
-          body[:tools]       = tools       unless tools.empty?
+          body[:thinking] = thinking if thinking
+          body[:system] = system if system
+          body[:tools] = tools unless tools.empty?
           body[:tool_choice] = tool_choice if tool_choice
 
           post '/messages', body
@@ -23,31 +23,9 @@ module Rasti
 
         private
 
-        def parse_usage(response)
-          usage = response['usage']
-          return unless usage
-          Usage.new(
-            provider:          'anthropic',
-            model:             response['model'],
-            input_tokens:      usage['input_tokens'],
-            output_tokens:     usage['output_tokens'],
-            cached_tokens:     usage['cache_read_input_tokens'] || 0,
-            reasoning_tokens:  0,
-            raw:               usage
-          )
-        end
-
-        def default_api_key
-          Rasti::AI.anthropic_api_key
-        end
-
-        def base_url
-          'https://api.anthropic.com/v1'
-        end
-
         def build_request(uri)
           request = super
-          request['x-api-key']        = api_key
+          request['x-api-key'] = api_key
           request['anthropic-version'] = ANTHROPIC_VERSION
           request
         end

data/lib/rasti/ai/anthropic/provider.rb

@@ -0,0 +1,87 @@
+module Rasti
+  module AI
+    module Anthropic
+      class Provider < Rasti::AI::Provider
+
+        STRUCTURED_OUTPUT_TOOL = 'structured_output'.freeze
+
+        THINKING_LEVELS = {
+          'low' => {type: 'enabled', budget_tokens: 1_024}.freeze,
+          'medium' => {type: 'enabled', budget_tokens: 8_000}.freeze,
+          'high' => {type: 'enabled', budget_tokens: 16_000}.freeze
+        }.freeze
+
+        def name
+          :anthropic
+        end
+
+        def default_model
+          Rasti::AI.anthropic_default_model
+        end
+
+        def default_api_key
+          Rasti::AI.anthropic_api_key
+        end
+
+        def base_url
+          'https://api.anthropic.com/v1'
+        end
+
+        def parse_usage(response)
+          usage = response['usage']
+          return nil unless usage
+
+          Usage.new(
+            provider: name.to_s,
+            model: response['model'],
+            input_tokens: usage['input_tokens'],
+            output_tokens: usage['output_tokens'],
+            cached_tokens: usage['cache_read_input_tokens'] || 0,
+            reasoning_tokens: 0,
+            raw: usage
+          )
+        end
+
+        private
+
+        def request(client:, messages:, system:, model:, json_schema:, thinking:)
+          tools = json_schema ? [structured_output_tool(json_schema)] : []
+          tool_choice = json_schema ? {type: 'tool', name: STRUCTURED_OUTPUT_TOOL} : nil
+
+          client.messages messages: messages,
+                          model: model,
+                          system: system,
+                          tools: tools,
+                          tool_choice: tool_choice,
+                          thinking: thinking
+        end
+
+        def encode_message(message)
+          {role: message[:role], content: message[:content]}
+        end
+
+        def parse_content(response)
+          content = response['content'] || []
+
+          structured = content.find { |block| block['type'] == 'tool_use' && block['name'] == STRUCTURED_OUTPUT_TOOL }
+          return JSON.dump(structured['input']) if structured
+
+          text_block = content.find { |block| block['type'] == 'text' }
+          text_block['text'] if text_block
+        end
+
+        def structured_output_tool(json_schema)
+          {
+            name: STRUCTURED_OUTPUT_TOOL,
+            description: 'Return the structured response',
+            input_schema: {
+              type: 'object',
+              properties: json_schema
+            }
+          }
+        end
+
+      end
+    end
+  end
+end

data/lib/rasti/ai/assistant.rb

@@ -2,14 +2,17 @@ module Rasti
   module AI
     class Assistant
 
+      include ProviderAware
+
       attr_reader :state, :model, :thinking
 
       VALID_THINKING_LEVELS = %w[low medium high].freeze
 
-      def initialize(client:nil, json_schema:nil, state:nil, model:nil, thinking:nil, tools:[], mcp_servers:{}, logger:nil)
+      def initialize(client:nil, provider:nil, json_schema:nil, state:nil, model:nil, thinking:nil, tools:[], mcp_servers:{}, logger:nil)
         raise ArgumentError, "Invalid thinking level '#{thinking}'. Valid: #{VALID_THINKING_LEVELS.join(', ')}" if thinking && !VALID_THINKING_LEVELS.include?(thinking)
 
-        @client = client || build_default_client
+        @provider = provider
+        @client = client
         @json_schema = json_schema
         @state = state || AssistantState.new
         @model = model
@@ -50,14 +53,16 @@ module Rasti
 
       private
 
-      attr_reader :client, :json_schema, :tools, :serialized_tools, :logger
+      attr_reader :json_schema, :tools, :serialized_tools, :logger
+
+      def client
+        @client ||= provider.build_client
+      end
 
       def messages
         state.messages
       end
 
-      # --- Shared behavior ---
-
       def register_tools(tools)
         tools.each do |tool|
           serialization = wrap_tool_serialization(ToolSerializer.serialize(tool.class))
@@ -99,12 +104,6 @@ module Rasti
         "Error: #{ex.message}"
       end
 
-      # --- Template methods ---
-
-      def build_default_client
-        raise NotImplementedError
-      end
-
       def build_user_message(prompt)
         raise NotImplementedError
       end

data/lib/rasti/ai/client.rb

@@ -2,9 +2,12 @@ module Rasti
   module AI
     class Client
 
+      include ProviderAware
+
       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, usage_tracker:nil)
+      def initialize(provider:nil, api_key:nil, logger:nil, http_connect_timeout:nil, http_read_timeout:nil, http_max_retries:nil, usage_tracker:nil)
+        @provider = provider
         @api_key = api_key || default_api_key
         @logger = logger || Rasti::AI.logger
         @http_connect_timeout = http_connect_timeout || Rasti::AI.http_connect_timeout
@@ -17,22 +20,22 @@ module Rasti
 
       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
+      def default_api_key
+        provider.default_api_key
       end
 
-      def parse_usage(response)
-        raise NotImplementedError
+      def default_model
+        provider.default_model
       end
 
-      def default_api_key
-        raise NotImplementedError
+      def track_usage(response)
+        return unless usage_tracker
+        usage = provider.parse_usage response
+        usage_tracker.call usage if usage
       end
 
       def base_url
-        raise NotImplementedError
+        provider.base_url
       end
 
       def build_url(relative_url)

data/lib/rasti/ai/errors.rb

@@ -34,6 +34,14 @@ module Rasti
 
       end
 
+      class UnknownProvider < StandardError
+
+        def initialize(provider, valid_providers)
+          super "Unknown provider #{provider.inspect}. Valid providers: #{valid_providers.join(', ')}"
+        end
+
+      end
+
     end
   end
 end

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

@@ -5,18 +5,8 @@ module Rasti
 
         ALLOWED_SCHEMA_FIELDS = %w[type description properties required enum items format nullable anyOf].freeze
 
-        THINKING_LEVELS = {
-          'low'    => {thinking_budget: 1_024}.freeze,
-          'medium' => {thinking_budget: 8_192}.freeze,
-          'high'   => {thinking_budget: 24_576}.freeze
-        }.freeze
-
         private
 
-        def build_default_client
-          Client.new
-        end
-
         def build_user_message(prompt)
           {role: Roles::USER, parts: [{text: prompt}]}
         end
@@ -112,7 +102,7 @@ module Rasti
         end
 
         def thinking_config
-          THINKING_LEVELS[thinking]
+          Provider::THINKING_LEVELS[thinking]
         end
 
         def generation_config

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

@@ -4,7 +4,7 @@ module Rasti
       class Client < Rasti::AI::Client
 
         def generate_content(contents:, model:nil, tools:[], system_instruction:nil, generation_config:nil)
-          model_name = model || Rasti::AI.gemini_default_model
+          model_name = model || default_model
 
           body = {contents: contents}
 
@@ -17,30 +17,8 @@ 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
-
-        def base_url
-          'https://generativelanguage.googleapis.com/v1beta'
-        end
-
         def build_url(relative_url)
-          "#{base_url}#{relative_url}?key=#{api_key}"
+          "#{super}?key=#{api_key}"
         end
 
       end

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

@@ -0,0 +1,90 @@
+module Rasti
+  module AI
+    module Gemini
+      class Provider < Rasti::AI::Provider
+
+        THINKING_LEVELS = {
+          'low' => {thinking_budget: 1_024}.freeze,
+          'medium' => {thinking_budget: 8_192}.freeze,
+          'high' => {thinking_budget: 24_576}.freeze
+        }.freeze
+
+        def name
+          :gemini
+        end
+
+        def default_model
+          Rasti::AI.gemini_default_model
+        end
+
+        def default_api_key
+          Rasti::AI.gemini_api_key
+        end
+
+        def base_url
+          'https://generativelanguage.googleapis.com/v1beta'
+        end
+
+        def parse_usage(response)
+          usage = response['usageMetadata']
+          return nil unless usage
+
+          Usage.new(
+            provider: name.to_s,
+            model: response['modelVersion'],
+            input_tokens: usage['promptTokenCount'],
+            output_tokens: usage['candidatesTokenCount'],
+            cached_tokens: usage['cachedContentTokenCount'] || 0,
+            reasoning_tokens: usage['thoughtsTokenCount'] || 0,
+            raw: usage
+          )
+        end
+
+        private
+
+        def request(client:, messages:, system:, model:, json_schema:, thinking:)
+          client.generate_content contents: messages,
+                                  model: model,
+                                  system_instruction: system_instruction(system),
+                                  generation_config: generation_config(json_schema, thinking)
+        end
+
+        def encode_message(message)
+          {
+            role: encode_role(message[:role]),
+            parts: [{text: message[:content]}]
+          }
+        end
+
+        def encode_role(role)
+          role == Rasti::AI::Roles::ASSISTANT ? Roles::MODEL : Roles::USER
+        end
+
+        def parse_content(response)
+          parts = response.dig('candidates', 0, 'content', 'parts') || []
+          text_part = parts.find { |part| part.key? 'text' }
+          text_part['text'] if text_part
+        end
+
+        def system_instruction(system)
+          return nil if system.nil?
+          {parts: [{text: system}]}
+        end
+
+        def generation_config(json_schema, thinking)
+          config = {}
+
+          config[:thinking_config] = thinking if thinking
+
+          if json_schema
+            config[:response_mime_type] = 'application/json'
+            config[:response_schema] = json_schema
+          end
+
+          config.empty? ? nil : config
+        end
+
+      end
+    end
+  end
+end