n2b 0.5.1 → 0.7.1

data/lib/n2b/llm/claude.rb

@@ -32,7 +32,7 @@ module N2M
           "messages" => [
             {
               "role" => "user",
-              "content" => content 
+              "content" => content
             }
           ]
         })
@@ -44,29 +44,16 @@ module N2M
         if response.code != '200'
           raise N2B::LlmApiError.new("LLM API Error: #{response.code} #{response.message} - #{response.body}")
         end
-        answer = JSON.parse(response.body)['content'].first['text'] 
-        begin 
-          # The llm_response.json file is likely for debugging and can be kept or removed.
-          # For this refactoring, I'll keep it as it doesn't affect the error handling logic.
-          File.open('llm_response.json', 'w') do |f|
-            f.write(answer)
-          end
-          # remove everything before the first { and after the last }
-          
+        answer = JSON.parse(response.body)['content'].first['text']
+        begin
           answer = answer.sub(/.*?\{(.*)\}.*/m, '{\1}') unless answer.start_with?('{')
-          # gsub all \n with \\n that are inside "
-          # 
           answer.gsub!(/"([^"]*)"/) { |match| match.gsub(/\n/, "\\n") }
-          # The llm_response.json file is likely for debugging and can be kept or removed.
-          File.open('llm_response.json', 'w') do |f|
-            f.write(answer)
-          end
           answer = JSON.parse(answer)
         rescue JSON::ParserError
           # This specific JSON parsing error is about the LLM's *response content*, not an API error.
           # It should probably be handled differently, but the subtask is about LlmApiError.
           # For now, keeping existing behavior for this part.
-          puts "Error parsing JSON from LLM response: #{answer}" # Clarified error message
+          puts "Error parsing JSON from LLM response: #{answer}"
           answer = { 'explanation' => answer} # Default fallback
         end
         answer
@@ -83,10 +70,10 @@ module N2M
 
         request.body = JSON.dump({
           "model" => get_model_name,
-          "max_tokens" => @config['max_tokens'] || 1024, # Allow overriding max_tokens from config
+          "max_tokens" => @config['max_tokens'] || 1024,
           "messages" => [
             {
-              "role" => "user", # The entire prompt is passed as a single user message
+              "role" => "user",
               "content" => prompt_content
             }
           ]

data/lib/n2b/llm/gemini.rb

@@ -34,7 +34,10 @@ module N2M
             "parts" => [{
               "text" => content
             }]
-          }]
+          }],
+          "generationConfig" => {
+            "responseMimeType" => "application/json"
+          }
         })
 
         response = Net::HTTP.start(uri.hostname, uri.port, use_ssl: true) do |http|
@@ -82,14 +85,12 @@ module N2M
         request.body = JSON.dump({
           "contents" => [{
             "parts" => [{
-              "text" => prompt_content # The entire prompt is passed as text
+              "text" => prompt_content
             }]
           }],
-          # Gemini specific: Ensure JSON output if possible via generationConfig
-          # However, the primary method is instructing it within the prompt itself.
-          # "generationConfig": {
-          #   "responseMimeType": "application/json", # This might be too restrictive or not always work as expected
-          # }
+          "generationConfig" => {
+            "responseMimeType" => "application/json"
+          }
         })
 
         response = Net::HTTP.start(uri.hostname, uri.port, use_ssl: true) do |http|

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