rcrewai 0.2.1 → 0.4.0

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

data/lib/rcrewai/llm_clients/openai.rb

@@ -1,6 +1,12 @@
 # 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
@@ -12,113 +18,184 @@ module RCrewAI
         @base_url = BASE_URL
       end
 
-      def chat(messages:, **options)
+      def chat(messages:, tools: nil, tool_choice: :auto, stream: nil, **options)
         payload = {
           model: config.model,
-          messages: format_messages(messages),
+          messages: messages,
           temperature: options[:temperature] || config.temperature,
           max_tokens: options[:max_tokens] || config.max_tokens
-        }
+        }.compact
 
-        # Add additional OpenAI-specific 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}/chat/completions"
-        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
+        if tools && !tools.empty?
+          payload[:tools] = ProviderSchema.for_many(:openai, tools)
+          payload[:tool_choice] = tool_choice if tool_choice != :auto
+        end
 
-      def complete(prompt:, **options)
-        # For older models that use completions endpoint
-        if config.model.include?('davinci') || config.model.include?('curie') || 
-           config.model.include?('babbage') || config.model.include?('ada')
-          completion_request(prompt, **options)
+        if stream
+          payload[:stream] = true
+          payload[:stream_options] = { include_usage: true }
+          stream_chat(payload, stream)
         else
-          # Use chat endpoint for newer models
-          super
+          plain_chat(payload)
         end
       end
 
+      def supports_native_tools?(model: config.model) # rubocop:disable Lint/UnusedMethodArgument
+        true
+      end
+
       def models
         url = "#{@base_url}/models"
-        response = http_client.get(url, {}, build_headers.merge(authorization_header))
+        response = http_client.get(url, {}, build_headers.merge(auth_header))
         result = handle_response(response)
         result['data'].map { |model| model['id'] }
       end
 
       private
 
-      def authorization_header
-        { 'Authorization' => "Bearer #{config.api_key}" }
+      def chat_url
+        "#{@base_url}/chat/completions"
       end
 
-      def completion_request(prompt, **options)
-        payload = {
-          model: config.model,
-          prompt: prompt,
-          temperature: options[:temperature] || config.temperature,
-          max_tokens: options[:max_tokens] || config.max_tokens
-        }
+      def plain_chat(payload)
+        url = chat_url
+        log_request(:post, url, payload)
+        response = http_client.post(url, payload, build_headers.merge(auth_header))
+        log_response(response)
+        body = handle_response(response)
+        normalize_non_streaming(body)
+      end
 
-        url = "#{@base_url}/completions"
+      def stream_chat(payload, sink) # rubocop:disable Metrics/AbcSize
+        url = chat_url
         log_request(:post, url, payload)
 
-        response = http_client.post(url, payload, build_headers.merge(authorization_header))
-        log_response(response)
+        assembled_text = +''
+        tool_calls_by_index = {}
+        final_usage = nil
+        finish_reason = nil
+
+        parser = SSEParser.new do |sse|
+          data_str = sse[:data]
+          next if data_str == '[DONE]'
+
+          data = JSON.parse(data_str)
+          choice = data.dig('choices', 0) || {}
+          delta = choice['delta'] || {}
+
+          if delta['content']
+            assembled_text << delta['content']
+            sink.call(Events::TextDelta.new(
+                        type: :text_delta, timestamp: Time.now, agent: nil, iteration: nil,
+                        text: delta['content']
+                      ))
+          end
 
-        result = handle_response(response)
-        format_completion_response(result)
-      end
+          Array(delta['tool_calls']).each do |tc|
+            idx = tc['index']
+            tool_calls_by_index[idx] ||= { id: nil, name: nil, arguments: +'' }
+            tool_calls_by_index[idx][:id]   ||= tc['id']
+            tool_calls_by_index[idx][:name] ||= tc.dig('function', 'name')
+            tool_calls_by_index[idx][:arguments] << (tc.dig('function', 'arguments') || '')
+          end
+
+          finish_reason ||= choice['finish_reason']&.to_sym
 
-      def format_messages(messages)
-        messages.map do |msg|
-          if msg.is_a?(Hash)
-            msg
-          else
-            { role: 'user', content: msg.to_s }
+          if data['usage']
+            final_usage = {
+              prompt_tokens: data['usage']['prompt_tokens'],
+              completion_tokens: data['usage']['completion_tokens'],
+              total_tokens: data['usage']['total_tokens']
+            }
           end
         end
-      end
 
-      def format_response(response)
-        choice = response.dig('choices', 0)
-        return nil unless choice
+        streaming_post(url, payload) { |chunk| parser.feed(chunk) }
+
+        tool_calls = tool_calls_by_index.values.map do |tc|
+          {
+            id: tc[:id],
+            name: tc[:name],
+            arguments: tc[:arguments].empty? ? {} : JSON.parse(tc[:arguments])
+          }
+        end
+
+        if final_usage
+          sink.call(Events::Usage.new(
+                      type: :usage, timestamp: Time.now, agent: nil, iteration: nil,
+                      prompt_tokens: final_usage[:prompt_tokens],
+                      completion_tokens: final_usage[:completion_tokens],
+                      total_tokens: final_usage[:total_tokens],
+                      cost_usd: Pricing.cost_for(config.model,
+                                                 prompt_tokens: final_usage[:prompt_tokens],
+                                                 completion_tokens: final_usage[:completion_tokens])
+                    ))
+        end
 
         {
-          content: choice.dig('message', 'content'),
-          role: choice.dig('message', 'role'),
-          finish_reason: choice['finish_reason'],
-          usage: response['usage'],
-          model: response['model'],
-          provider: :openai
+          content: assembled_text.empty? ? nil : assembled_text,
+          tool_calls: tool_calls,
+          usage: final_usage || {},
+          finish_reason: finish_reason || :stop,
+          model: config.model,
+          provider: provider_name
         }
       end
 
-      def format_completion_response(response)
-        choice = response.dig('choices', 0)
-        return nil unless choice
+      def provider_name
+        :openai
+      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.merge(auth_header)
+          req.body = payload.to_json
+          req.options.on_data = proc { |chunk, _| on_chunk.call(chunk) }
+        end
+      end
+
+      def normalize_non_streaming(body)
+        choice = body.dig('choices', 0) || {}
+        msg = choice['message'] || {}
+        tool_calls = Array(msg['tool_calls']).map do |tc|
+          {
+            id: tc['id'],
+            name: tc.dig('function', 'name'),
+            arguments: JSON.parse(tc.dig('function', 'arguments') || '{}')
+          }
+        end
         {
-          content: choice['text'],
-          finish_reason: choice['finish_reason'],
-          usage: response['usage'],
-          model: response['model'],
-          provider: :openai
+          content: msg['content'],
+          tool_calls: tool_calls,
+          usage: {
+            prompt_tokens: body.dig('usage', 'prompt_tokens'),
+            completion_tokens: body.dig('usage', 'completion_tokens'),
+            total_tokens: body.dig('usage', 'total_tokens')
+          },
+          finish_reason: (choice['finish_reason'] || 'stop').to_sym,
+          model: body['model'] || config.model,
+          provider: provider_name
         }
       end
 
+      def auth_header
+        { 'Authorization' => "Bearer #{config.openai_api_key || config.api_key}" }
+      end
+
       def validate_config!
-        raise ConfigurationError, "OpenAI API key is required" unless config.openai_api_key || config.api_key
-        raise ConfigurationError, "Model is required" unless config.model
+        raise ConfigurationError, 'OpenAI API key is required' unless config.openai_api_key || config.api_key
+        raise ConfigurationError, 'Model is required' unless config.model
       end
     end
   end
-end
+end