rcrewai 0.2.1 → 0.4.0

data/lib/rcrewai/llm_clients/azure.rb

@@ -1,158 +1,53 @@
 # frozen_string_literal: true
 
-require_relative 'base'
+require_relative 'openai'
 
 module RCrewAI
   module LLMClients
-    class Azure < Base
+    # Azure OpenAI: same wire format as OpenAI but routes through a deployment
+    # path with an api-version query param, and authenticates with an api-key
+    # header instead of Authorization: Bearer.
+    class Azure < OpenAI
       def initialize(config = RCrewAI.configuration)
         super
-        @base_url = config.base_url || build_azure_url
         @api_version = config.api_version || '2024-02-01'
         @deployment_name = config.deployment_name || config.model
-      end
-
-      def chat(messages:, **options)
-        payload = {
-          messages: format_messages(messages),
-          temperature: options[:temperature] || config.temperature,
-          max_tokens: options[:max_tokens] || config.max_tokens
-        }
-
-        # Add additional OpenAI-compatible options
-        payload[:top_p] = options[:top_p] if options[:top_p]
-        payload[:frequency_penalty] = options[:frequency_penalty] if options[:frequency_penalty]
-        payload[:presence_penalty] = options[:presence_penalty] if options[:presence_penalty]
-        payload[:stop] = options[:stop] if options[:stop]
-
-        url = "#{@base_url}/openai/deployments/#{@deployment_name}/chat/completions?api-version=#{@api_version}"
-        log_request(:post, url, payload)
-
-        response = http_client.post(url, payload, build_headers.merge(authorization_header))
-        log_response(response)
-
-        result = handle_response(response)
-        format_response(result)
-      end
-
-      def complete(prompt:, **options)
-        # For older models that use completions endpoint
-        payload = {
-          prompt: prompt,
-          temperature: options[:temperature] || config.temperature,
-          max_tokens: options[:max_tokens] || config.max_tokens
-        }
-
-        url = "#{@base_url}/openai/deployments/#{@deployment_name}/completions?api-version=#{@api_version}"
-        log_request(:post, url, payload)
-
-        response = http_client.post(url, payload, build_headers.merge(authorization_header))
-        log_response(response)
-
-        result = handle_response(response)
-        format_completion_response(result)
+        @base_url = build_endpoint_url
       end
 
       def models
-        # Azure OpenAI uses deployments instead of models
         url = "#{@base_url}/openai/deployments?api-version=#{@api_version}"
-        response = http_client.get(url, {}, build_headers.merge(authorization_header))
+        response = http_client.get(url, {}, build_headers.merge(auth_header))
         result = handle_response(response)
-        
-        if result['data']
-          result['data'].map { |deployment| deployment['id'] }
-        else
-          [@deployment_name].compact
-        end
+        Array(result['data']).map { |d| d['id'] }
+      rescue StandardError
+        [@deployment_name].compact
       end
 
       private
 
-      def authorization_header
-        { 'api-key' => config.api_key }
+      def chat_url
+        "#{@base_url}/openai/deployments/#{@deployment_name}/chat/completions?api-version=#{@api_version}"
       end
 
-      def format_messages(messages)
-        messages.map do |msg|
-          if msg.is_a?(Hash)
-            msg
-          else
-            { role: 'user', content: msg.to_s }
-          end
-        end
+      def build_endpoint_url
+        endpoint = config.base_url || ENV['AZURE_OPENAI_ENDPOINT'] || ENV['AZURE_ENDPOINT']
+        endpoint&.chomp('/')
       end
 
-      def format_response(response)
-        choice = response.dig('choices', 0)
-        return nil unless choice
-
-        {
-          content: choice.dig('message', 'content'),
-          role: choice.dig('message', 'role'),
-          finish_reason: choice['finish_reason'],
-          usage: response['usage'],
-          model: @deployment_name,
-          provider: :azure
-        }
+      def provider_name
+        :azure
       end
 
-      def format_completion_response(response)
-        choice = response.dig('choices', 0)
-        return nil unless choice
-
-        {
-          content: choice['text'],
-          finish_reason: choice['finish_reason'],
-          usage: response['usage'],
-          model: @deployment_name,
-          provider: :azure
-        }
+      def auth_header
+        { 'api-key' => config.azure_api_key || config.api_key }
       end
 
       def validate_config!
-        super
-        raise ConfigurationError, "Azure API key is required" unless config.azure_api_key || config.api_key
-        raise ConfigurationError, "Azure base URL or endpoint is required" unless config.base_url || azure_endpoint
-        raise ConfigurationError, "Azure deployment name is required" unless config.deployment_name || config.model
-      end
-
-      def build_azure_url
-        endpoint = azure_endpoint
-        return nil unless endpoint
-        
-        # Remove trailing slash and add proper path
-        endpoint = endpoint.chomp('/')
-        "#{endpoint}"
-      end
-
-      def azure_endpoint
-        # Try multiple environment variable names
-        ENV['AZURE_OPENAI_ENDPOINT'] || 
-        ENV['AZURE_ENDPOINT'] || 
-        config.instance_variable_get(:@azure_endpoint)
-      end
-
-      def handle_response(response)
-        case response.status
-        when 200..299
-          response.body
-        when 400
-          error_details = response.body.dig('error', 'message') || response.body
-          raise APIError, "Bad request: #{error_details}"
-        when 401
-          raise AuthenticationError, "Invalid API key or authentication failed"
-        when 403
-          raise AuthenticationError, "Access denied - check your API key and permissions"
-        when 404
-          raise ModelNotFoundError, "Deployment '#{@deployment_name}' not found"
-        when 429
-          raise RateLimitError, "Rate limit exceeded or quota exhausted"
-        when 500..599
-          raise APIError, "Azure OpenAI service error: #{response.status}"
-        else
-          raise APIError, "Unexpected response: #{response.status}"
-        end
+        raise ConfigurationError, 'Azure API key is required' unless config.azure_api_key || config.api_key
+        raise ConfigurationError, 'Azure endpoint is required' unless config.base_url || ENV['AZURE_OPENAI_ENDPOINT'] || ENV['AZURE_ENDPOINT']
+        raise ConfigurationError, 'Azure deployment name is required' unless config.deployment_name || config.model
       end
     end
   end
-end
+end

data/lib/rcrewai/llm_clients/base.rb

@@ -16,8 +16,12 @@ module RCrewAI
         validate_config!
       end
 
-      def chat(messages:, **options)
-        raise NotImplementedError, "Subclasses must implement #chat method"
+      def chat(messages:, tools: nil, tool_choice: :auto, stream: nil, **options)
+        raise NotImplementedError, 'Subclasses must implement #chat method'
+      end
+
+      def supports_native_tools?(model: config.model) # rubocop:disable Lint/UnusedMethodArgument
+        true
       end
 
       def complete(prompt:, **options)
@@ -27,8 +31,8 @@ module RCrewAI
       protected
 
       def validate_config!
-        raise ConfigurationError, "API key is required" unless config.api_key
-        raise ConfigurationError, "Model is required" unless config.model
+        raise ConfigurationError, 'API key is required' unless config.api_key
+        raise ConfigurationError, 'Model is required' unless config.model
       end
 
       def build_headers
@@ -54,9 +58,9 @@ module RCrewAI
         when 400
           raise APIError, "Bad request: #{response.body}"
         when 401
-          raise AuthenticationError, "Invalid API key"
+          raise AuthenticationError, 'Invalid API key'
         when 429
-          raise RateLimitError, "Rate limit exceeded"
+          raise RateLimitError, 'Rate limit exceeded'
         when 500..599
           raise APIError, "Server error: #{response.status}"
         else
@@ -79,4 +83,4 @@ module RCrewAI
     class RateLimitError < APIError; end
     class ModelNotFoundError < APIError; end
   end
-end
+end

data/lib/rcrewai/llm_clients/google.rb

@@ -1,137 +1,201 @@
 # frozen_string_literal: true
 
+require 'faraday'
+require 'json'
 require_relative 'base'
+require_relative '../events'
+require_relative '../sse_parser'
+require_relative '../provider_schema'
+require_relative '../pricing'
 
 module RCrewAI
   module LLMClients
     class Google < Base
       BASE_URL = 'https://generativelanguage.googleapis.com/v1beta'
 
+      FINISH_REASON_MAP = {
+        'STOP' => :stop,
+        'MAX_TOKENS' => :length,
+        'SAFETY' => :stop,
+        'RECITATION' => :stop
+      }.freeze
+
       def initialize(config = RCrewAI.configuration)
         super
         @base_url = BASE_URL
       end
 
-      def chat(messages:, **options)
-        # Convert messages to Gemini format
-        formatted_contents = format_messages(messages)
-        
+      def chat(messages:, tools: nil, tool_choice: :auto, stream: nil, **options) # rubocop:disable Lint/UnusedMethodArgument
+        contents = format_messages(messages)
         payload = {
-          contents: formatted_contents,
+          contents: contents,
           generationConfig: {
             temperature: options[:temperature] || config.temperature,
-            maxOutputTokens: options[:max_tokens] || config.max_tokens || 2048,
-            topP: options[:top_p] || 0.8,
-            topK: options[:top_k] || 10
-          }
+            maxOutputTokens: options[:max_tokens] || config.max_tokens || 2048
+          }.compact
         }
-
-        # Add safety settings if provided
-        if options[:safety_settings]
-          payload[:safetySettings] = options[:safety_settings]
+        payload[:generationConfig][:topP] = options[:top_p] if options[:top_p]
+        payload[:generationConfig][:topK] = options[:top_k] if options[:top_k]
+        payload[:generationConfig][:stopSequences] = options[:stop_sequences] if options[:stop_sequences]
+        payload[:safetySettings] = options[:safety_settings] if options[:safety_settings]
+
+        if tools && !tools.empty?
+          payload[:tools] = [ProviderSchema.for_many(:google, tools)]
+          # tool_choice: Google has toolConfig; left as default unless explicitly set
         end
 
-        # Add stop sequences if provided
-        if options[:stop_sequences]
-          payload[:generationConfig][:stopSequences] = options[:stop_sequences]
+        api_key = config.google_api_key || config.api_key
+        if stream
+          url = "#{@base_url}/models/#{config.model}:streamGenerateContent?alt=sse&key=#{api_key}"
+          stream_chat(url, payload, stream)
+        else
+          url = "#{@base_url}/models/#{config.model}:generateContent?key=#{api_key}"
+          plain_chat(url, payload)
         end
+      end
 
-        url = "#{@base_url}/models/#{config.model}:generateContent?key=#{config.api_key}"
-        log_request(:post, url, payload)
-
-        response = http_client.post(url, payload, build_headers)
-        log_response(response)
-
-        result = handle_response(response)
-        format_response(result)
+      def supports_native_tools?(model: config.model) # rubocop:disable Lint/UnusedMethodArgument
+        true
       end
 
       def models
-        # Google AI Studio doesn't provide a models list endpoint with API key auth
-        # Return known Gemini models
-        [
-          'gemini-pro',
-          'gemini-pro-vision',
-          'gemini-1.5-pro',
-          'gemini-1.5-flash',
-          'text-bison-001',
-          'chat-bison-001'
-        ]
+        %w[gemini-pro gemini-1.5-pro gemini-1.5-flash gemini-pro-vision]
       end
 
       private
 
-      def format_messages(messages)
-        contents = []
-        
-        messages.each do |msg|
-          role = case msg[:role]
-                when 'user'
-                  'user'
-                when 'assistant'
-                  'model'
-                when 'system'
-                  # Gemini doesn't have system role, prepend to first user message
-                  next
-                else
-                  'user'
-                end
-
-          content = if msg.is_a?(Hash)
-                     msg[:content]
-                   else
-                     msg.to_s
-                   end
-
-          contents << {
-            role: role,
-            parts: [{ text: content }]
-          }
-        end
+      def plain_chat(url, payload)
+        log_request(:post, url, payload)
+        response = http_client.post(url, payload, build_headers)
+        log_response(response)
+        body = handle_response(response)
+        normalize_non_streaming(body)
+      end
+
+      def stream_chat(url, payload, sink)
+        log_request(:post, url, payload)
+
+        assembled_text = +''
+        tool_calls = []
+        finish_reason = nil
+        usage = nil
+
+        parser = SSEParser.new do |sse|
+          data = JSON.parse(sse[:data])
+          candidate = data.dig('candidates', 0) || {}
+          parts = candidate.dig('content', 'parts') || []
+
+          parts.each do |part|
+            if part['text']
+              text = part['text']
+              assembled_text << text
+              sink.call(Events::TextDelta.new(type: :text_delta, timestamp: Time.now,
+                                              agent: nil, iteration: nil, text: text))
+            elsif part['functionCall']
+              fc = part['functionCall']
+              tool_calls << { id: nil, name: fc['name'], arguments: fc['args'] || {} }
+            end
+          end
 
-        # Handle system message by prepending to first user message
-        system_msg = messages.find { |m| m[:role] == 'system' }
-        if system_msg && contents.any?
-          first_user_content = contents.find { |c| c[:role] == 'user' }
-          if first_user_content
-            original_text = first_user_content[:parts].first[:text]
-            first_user_content[:parts].first[:text] = "#{system_msg[:content]}\n\n#{original_text}"
+          finish_reason ||= FINISH_REASON_MAP[candidate['finishReason']]
+          if data['usageMetadata']
+            usage = {
+              prompt_tokens: data.dig('usageMetadata', 'promptTokenCount'),
+              completion_tokens: data.dig('usageMetadata', 'candidatesTokenCount'),
+              total_tokens: data.dig('usageMetadata', 'totalTokenCount')
+            }
           end
         end
 
-        contents
-      end
+        streaming_post(url, payload) { |chunk| parser.feed(chunk) }
+
+        if usage
+          sink.call(Events::Usage.new(
+                      type: :usage, timestamp: Time.now, agent: nil, iteration: nil,
+                      prompt_tokens: usage[:prompt_tokens],
+                      completion_tokens: usage[:completion_tokens],
+                      total_tokens: usage[:total_tokens],
+                      cost_usd: Pricing.cost_for(config.model,
+                                                 prompt_tokens: usage[:prompt_tokens] || 0,
+                                                 completion_tokens: usage[:completion_tokens] || 0)
+                    ))
+        end
 
-      def format_response(response)
-        candidate = response.dig('candidates', 0)
-        return nil unless candidate
+        finish_reason = :tool_calls if tool_calls.any?
 
-        content = candidate.dig('content', 'parts', 0, 'text')
-        finish_reason = candidate['finishReason']
+        {
+          content: assembled_text.empty? ? nil : assembled_text,
+          tool_calls: tool_calls,
+          usage: usage || {},
+          finish_reason: finish_reason || :stop,
+          model: config.model,
+          provider: :google
+        }
+      end
+
+      def streaming_post(url, payload, &on_chunk)
+        conn = Faraday.new do |f|
+          f.request :json
+          f.options.timeout = config.timeout
+          f.adapter Faraday.default_adapter
+        end
+        conn.post(url) do |req|
+          req.headers = build_headers
+          req.body = payload.to_json
+          req.options.on_data = proc { |chunk, _| on_chunk.call(chunk) }
+        end
+      end
 
-        # Extract usage information if available
-        usage_metadata = response['usageMetadata']
-        usage = if usage_metadata
-                  {
-                    'prompt_tokens' => usage_metadata['promptTokenCount'],
-                    'completion_tokens' => usage_metadata['candidatesTokenCount'],
-                    'total_tokens' => usage_metadata['totalTokenCount']
-                  }
-                end
+      def normalize_non_streaming(body)
+        candidate = body.dig('candidates', 0) || {}
+        parts = candidate.dig('content', 'parts') || []
+        text = parts.select { |p| p['text'] }.map { |p| p['text'] }.join
+        tool_calls = parts.select { |p| p['functionCall'] }.map do |p|
+          { id: nil, name: p.dig('functionCall', 'name'), arguments: p.dig('functionCall', 'args') || {} }
+        end
+        usage_md = body['usageMetadata'] || {}
+        finish_reason = tool_calls.any? ? :tool_calls : (FINISH_REASON_MAP[candidate['finishReason']] || :stop)
 
         {
-          content: content,
-          role: 'assistant',
+          content: text.empty? ? nil : text,
+          tool_calls: tool_calls,
+          usage: {
+            prompt_tokens: usage_md['promptTokenCount'],
+            completion_tokens: usage_md['candidatesTokenCount'],
+            total_tokens: usage_md['totalTokenCount']
+          },
           finish_reason: finish_reason,
-          usage: usage,
           model: config.model,
           provider: :google
         }
       end
 
+      def format_messages(messages)
+        contents = []
+        system_msg = nil
+
+        messages.each do |msg|
+          case msg[:role]
+          when 'system'
+            system_msg = msg[:content]
+            next
+          when 'assistant'
+            contents << { role: 'model', parts: [{ text: msg[:content].to_s }] }
+          else
+            contents << { role: 'user', parts: [{ text: msg[:content].to_s }] }
+          end
+        end
+
+        if system_msg && (first_user = contents.find { |c| c[:role] == 'user' })
+          first_user[:parts].first[:text] = "#{system_msg}\n\n#{first_user[:parts].first[:text]}"
+        end
+
+        contents
+      end
+
       def validate_config!
-        super
-        raise ConfigurationError, "Google API key is required" unless config.google_api_key || config.api_key
+        raise ConfigurationError, 'Google API key is required' unless config.google_api_key || config.api_key
+        raise ConfigurationError, 'Model is required' unless config.model
       end
 
       def handle_response(response)
@@ -139,14 +203,14 @@ module RCrewAI
         when 200..299
           response.body
         when 400
-          error_details = response.body.dig('error', 'message') || response.body
-          raise APIError, "Bad request: #{error_details}"
+          details = response.body.is_a?(Hash) ? response.body.dig('error', 'message') : response.body
+          raise APIError, "Bad request: #{details}"
         when 401
-          raise AuthenticationError, "Invalid API key"
+          raise AuthenticationError, 'Invalid API key'
         when 403
-          raise AuthenticationError, "API key does not have permission"
+          raise AuthenticationError, 'API key does not have permission'
         when 429
-          raise RateLimitError, "Rate limit exceeded or quota exhausted"
+          raise RateLimitError, 'Rate limit exceeded'
         when 500..599
           raise APIError, "Server error: #{response.status}"
         else
@@ -155,4 +219,4 @@ module RCrewAI
       end
     end
   end
-end
+end