rcrewai 0.2.1 → 0.3.0

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

data/lib/rcrewai/llm_clients/ollama.rb

@@ -1,18 +1,32 @@
 # frozen_string_literal: true
 
+require 'faraday'
+require 'json'
 require_relative 'base'
+require_relative '../events'
+require_relative '../provider_schema'
+require_relative '../pricing'
 
 module RCrewAI
   module LLMClients
     class Ollama < Base
       DEFAULT_URL = 'http://localhost:11434'
 
+      NATIVE_TOOL_MODELS = %w[
+        llama3.1 llama3.1:8b llama3.1:70b llama3.1:405b
+        llama3.2 llama3.2:1b llama3.2:3b
+        qwen2.5 qwen2.5:7b qwen2.5:14b qwen2.5:32b qwen2.5:72b
+        mistral-nemo mistral-large
+        command-r command-r-plus
+        firefunction-v2
+      ].freeze
+
       def initialize(config = RCrewAI.configuration)
         super
         @base_url = config.base_url || ollama_url || DEFAULT_URL
       end
 
-      def chat(messages:, **options)
+      def chat(messages:, tools: nil, tool_choice: :auto, stream: nil, **options) # rubocop:disable Lint/UnusedMethodArgument
         payload = {
           model: config.model,
           messages: format_messages(messages),
@@ -24,147 +38,174 @@ module RCrewAI
             repeat_penalty: options[:repeat_penalty]
           }.compact
         }
-
-        # Add stop sequences if provided
         payload[:options][:stop] = options[:stop] if options[:stop]
 
-        url = "#{@base_url}/api/chat"
-        log_request(:post, url, payload)
-
-        response = http_client.post(url, payload, build_headers)
-        log_response(response)
+        if tools && !tools.empty?
+          payload[:tools] = ProviderSchema.for_many(:ollama, tools)
+        end
 
-        result = handle_response(response)
-        format_response(result)
+        url = "#{@base_url}/api/chat"
+        if stream
+          payload[:stream] = true
+          stream_chat(url, payload, stream)
+        else
+          payload[:stream] = false
+          plain_chat(url, payload)
+        end
       end
 
-      def complete(prompt:, **options)
-        payload = {
-          model: config.model,
-          prompt: prompt,
-          options: {
-            temperature: options[:temperature] || config.temperature,
-            num_predict: options[:max_tokens] || config.max_tokens,
-            top_p: options[:top_p],
-            top_k: options[:top_k],
-            repeat_penalty: options[:repeat_penalty]
-          }.compact
-        }
-
-        payload[:options][:stop] = options[:stop] if options[:stop]
-
-        url = "#{@base_url}/api/generate"
-        log_request(:post, url, payload)
-
-        response = http_client.post(url, payload, build_headers)
-        log_response(response)
+      def supports_native_tools?(model: config.model)
+        override = RCrewAI.configuration.respond_to?(:ollama_native_tools) ? RCrewAI.configuration.ollama_native_tools : nil
+        return override unless override.nil?
 
-        result = handle_response(response)
-        format_completion_response(result)
+        base = model.to_s.split(':').first
+        NATIVE_TOOL_MODELS.any? { |m| m == model || m.split(':').first == base }
       end
 
       def models
         url = "#{@base_url}/api/tags"
         response = http_client.get(url, {}, build_headers)
         result = handle_response(response)
-        
-        if result['models']
-          result['models'].map { |model| model['name'] }
-        else
-          []
-        end
-      rescue => e
+        Array(result['models']).map { |m| m['name'] }
+      rescue StandardError => e
         logger.warn "Failed to fetch Ollama models: #{e.message}"
         []
       end
 
       def pull_model(model_name)
-        payload = { name: model_name }
         url = "#{@base_url}/api/pull"
-        
-        response = http_client.post(url, payload, build_headers)
+        response = http_client.post(url, { name: model_name }, build_headers)
         handle_response(response)
       end
 
-      def model_info(model_name = nil)
-        model_name ||= config.model
-        payload = { name: model_name }
-        url = "#{@base_url}/api/show"
-        
+      private
+
+      def plain_chat(url, payload)
+        log_request(:post, url, payload)
         response = http_client.post(url, payload, build_headers)
-        handle_response(response)
-      rescue => e
-        logger.warn "Failed to get model info for #{model_name}: #{e.message}"
-        nil
+        log_response(response)
+        body = handle_response(response)
+        normalize_non_streaming(body)
       end
 
-      private
+      def stream_chat(url, payload, sink)
+        log_request(:post, url, payload)
 
-      def format_messages(messages)
-        messages.map do |msg|
-          if msg.is_a?(Hash)
-            {
-              role: msg[:role],
-              content: msg[:content]
-            }
-          else
-            { role: 'user', content: msg.to_s }
+        assembled_text = +''
+        tool_calls = []
+        finish_reason = nil
+        prompt_tokens = nil
+        completion_tokens = nil
+        buffer = String.new(encoding: Encoding::UTF_8)
+
+        process_line = lambda do |line|
+          line = line.strip
+          return if line.empty?
+
+          data = JSON.parse(line)
+          if (msg = data['message'])
+            if msg['content']
+              assembled_text << msg['content']
+              sink.call(Events::TextDelta.new(type: :text_delta, timestamp: Time.now,
+                                              agent: nil, iteration: nil,
+                                              text: msg['content']))
+            end
+            Array(msg['tool_calls']).each do |tc|
+              fn = tc['function'] || {}
+              tool_calls << {
+                id: tc['id'],
+                name: fn['name'],
+                arguments: fn['arguments'].is_a?(String) ? JSON.parse(fn['arguments']) : (fn['arguments'] || {})
+              }
+            end
+          end
+          if data['done']
+            finish_reason = tool_calls.any? ? :tool_calls : :stop
+            prompt_tokens = data['prompt_eval_count']
+            completion_tokens = data['eval_count']
           end
         end
-      end
-
-      def format_response(response)
-        message = response['message']
-        return nil unless message
 
-        # Ollama doesn't provide detailed usage stats by default
-        usage = {
-          'prompt_tokens' => response['prompt_eval_count'],
-          'completion_tokens' => response['eval_count'],
-          'total_tokens' => (response['prompt_eval_count'] || 0) + (response['eval_count'] || 0)
-        }.compact
+        streaming_post(url, payload) do |chunk|
+          chunk = chunk.dup.force_encoding(Encoding::UTF_8) unless chunk.encoding == Encoding::UTF_8
+          buffer << chunk
+          while (idx = buffer.index("\n"))
+            line = buffer.slice!(0, idx + 1)
+            process_line.call(line)
+          end
+        end
+        process_line.call(buffer) unless buffer.empty?
+
+        if prompt_tokens || completion_tokens
+          sink.call(Events::Usage.new(
+                      type: :usage, timestamp: Time.now, agent: nil, iteration: nil,
+                      prompt_tokens: prompt_tokens, completion_tokens: completion_tokens,
+                      total_tokens: (prompt_tokens || 0) + (completion_tokens || 0),
+                      cost_usd: nil
+                    ))
+        end
 
         {
-          content: message['content'],
-          role: message['role'] || 'assistant',
-          finish_reason: response['done'] ? 'stop' : nil,
-          usage: usage,
-          model: response['model'] || config.model,
+          content: assembled_text.empty? ? nil : assembled_text,
+          tool_calls: tool_calls,
+          usage: {
+            prompt_tokens: prompt_tokens,
+            completion_tokens: completion_tokens,
+            total_tokens: (prompt_tokens || 0) + (completion_tokens || 0)
+          },
+          finish_reason: finish_reason || :stop,
+          model: config.model,
           provider: :ollama
         }
       end
 
-      def format_completion_response(response)
+      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
+
+      def normalize_non_streaming(body)
+        msg = body['message'] || {}
+        text = msg['content']
+        tool_calls = Array(msg['tool_calls']).map do |tc|
+          fn = tc['function'] || {}
+          args = fn['arguments']
+          args = JSON.parse(args) if args.is_a?(String)
+          { id: tc['id'], name: fn['name'], arguments: args || {} }
+        end
+        prompt_tokens = body['prompt_eval_count']
+        completion_tokens = body['eval_count']
+
         {
-          content: response['response'],
-          finish_reason: response['done'] ? 'stop' : nil,
+          content: text && !text.empty? ? text : nil,
+          tool_calls: tool_calls,
           usage: {
-            'prompt_tokens' => response['prompt_eval_count'],
-            'completion_tokens' => response['eval_count'],
-            'total_tokens' => (response['prompt_eval_count'] || 0) + (response['eval_count'] || 0)
-          }.compact,
-          model: response['model'] || config.model,
+            prompt_tokens: prompt_tokens,
+            completion_tokens: completion_tokens,
+            total_tokens: (prompt_tokens || 0) + (completion_tokens || 0)
+          },
+          finish_reason: tool_calls.any? ? :tool_calls : :stop,
+          model: body['model'] || config.model,
           provider: :ollama
         }
       end
 
-      def validate_config!
-        # Ollama doesn't require an API key
-        raise ConfigurationError, "Model is required" unless config.model
-        
-        # Test connection to Ollama server
-        test_connection
-      end
-
-      def test_connection
-        url = "#{@base_url}/api/tags"
-        response = http_client.get(url, {}, build_headers)
-        
-        unless (200..299).include?(response.status)
-          raise ConfigurationError, "Cannot connect to Ollama server at #{@base_url}"
+      def format_messages(messages)
+        messages.map do |msg|
+          if msg.is_a?(Hash)
+            { role: msg[:role], content: msg[:content] }
+          else
+            { role: 'user', content: msg.to_s }
+          end
         end
-      rescue Faraday::ConnectionFailed
-        raise ConfigurationError, "Cannot connect to Ollama server at #{@base_url}. Is Ollama running?"
       end
 
       def ollama_url
@@ -172,22 +213,25 @@ module RCrewAI
       end
 
       def build_headers
-        # Ollama doesn't require special headers
         {
           'Content-Type' => 'application/json',
           'User-Agent' => "rcrewai/#{RCrewAI::VERSION}"
         }
       end
 
+      def validate_config!
+        raise ConfigurationError, 'Model is required' unless config.model
+      end
+
       def handle_response(response)
         case response.status
         when 200..299
           response.body
         when 400
-          error_details = response.body['error'] || response.body
-          raise APIError, "Bad request: #{error_details}"
+          details = response.body.is_a?(Hash) ? response.body['error'] : response.body
+          raise APIError, "Bad request: #{details}"
         when 404
-          raise ModelNotFoundError, "Model '#{config.model}' not found. Try running: ollama pull #{config.model}"
+          raise ModelNotFoundError, "Model '#{config.model}' not found. Try: ollama pull #{config.model}"
         when 500..599
           raise APIError, "Ollama server error: #{response.status}"
         else
@@ -196,4 +240,4 @@ module RCrewAI
       end
     end
   end
-end
+end