n2b 0.5.0 → 0.7.1

data/lib/n2b/llm/open_ai.rb

@@ -22,20 +22,23 @@ module N2M
         model_name
       end
 
-      def make_request(content)
+      def make_request(content, expect_json: true)
         request = Net::HTTP::Post.new(API_URI)
         request.content_type = 'application/json'
         request['Authorization'] = "Bearer #{@config['access_key']}"
 
-        request.body = JSON.dump({
+        body_hash = {
           "model" => get_model_name,
-          response_format: { type: 'json_object' },
           "messages" => [
             {
               "role" => "user",
               "content" => content
-            }]
-        })
+            }
+          ]
+        }
+        body_hash["response_format"] = { "type" => "json_object" } if expect_json
+
+        request.body = JSON.dump(body_hash)
 
         response = Net::HTTP.start(API_URI.hostname, API_URI.port, use_ssl: true) do |http|
           http.request(request)
@@ -46,34 +49,38 @@ module N2M
           raise N2B::LlmApiError.new("LLM API Error: #{response.code} #{response.message} - #{response.body}")
         end
         answer = JSON.parse(response.body)['choices'].first['message']['content']
-        begin
-          # remove everything before the first { and after the last }
-          answer = answer.sub(/.*\{(.*)\}.*/m, '{\1}') unless answer.start_with?('{')
-          answer = JSON.parse(answer)
-        rescue JSON::ParserError
-          answer = { 'commands' => answer.split("\n"), explanation: answer }
+        if expect_json
+          begin
+            # remove everything before the first { and after the last }
+            answer = answer.sub(/.*\{(.*)\}.*/m, '{\1}') unless answer.start_with?('{')
+            answer = JSON.parse(answer)
+          rescue JSON::ParserError
+            answer = { 'commands' => answer.split("\n"), explanation: answer }
+          end
         end
         answer
       end
 
-      def analyze_code_diff(prompt_content)
+      def analyze_code_diff(prompt_content, expect_json: true)
         # This method assumes prompt_content is the full, ready-to-send prompt
         # including all instructions for the LLM (system message, diff, user additions, JSON format).
         request = Net::HTTP::Post.new(API_URI)
         request.content_type = 'application/json'
         request['Authorization'] = "Bearer #{@config['access_key']}"
 
-        request.body = JSON.dump({
+        body_hash = {
           "model" => get_model_name,
-          "response_format" => { "type" => "json_object" }, # Crucial for OpenAI to return JSON
           "messages" => [
             {
-              "role" => "user", # The entire prompt is passed as a single user message
+              "role" => "user",
               "content" => prompt_content
             }
           ],
-          "max_tokens" => @config['max_tokens'] || 1500 # Allow overriding, ensure it's enough for JSON
-        })
+          "max_tokens" => @config['max_tokens'] || 1500
+        }
+        body_hash["response_format"] = { "type" => "json_object" } if expect_json
+
+        request.body = JSON.dump(body_hash)
 
         response = Net::HTTP.start(API_URI.hostname, API_URI.port, use_ssl: true) do |http|
           http.request(request)
@@ -83,9 +90,8 @@ module N2M
           raise N2B::LlmApiError.new("LLM API Error: #{response.code} #{response.message} - #{response.body}")
         end
 
-        # Return the raw JSON string. CLI's call_llm_for_diff_analysis will handle parsing.
-        # OpenAI with json_object mode should return the JSON directly in 'choices'.first.message.content
-        JSON.parse(response.body)['choices'].first['message']['content']
+        answer = JSON.parse(response.body)['choices'].first['message']['content']
+        answer
       end
     end
   end