Merge pull request #18 from SergoHUH/feat/handle-structured-output-responses

feat(llm): handle structured output responses

Author: DenisDenis9331

lib/aireview/review_pipeline.rb

@@ -120,7 +120,34 @@ def critique_candidates(merge_request:, changes_text:, jira_issue:, candidates:)
     end
 
     def parse_with_repair(raw:, kind:, expected:, repair_stage:, critique_candidate_ids: nil)
-      parse_expected_json(raw, expected, critique_candidate_ids: critique_candidate_ids)
+      if raw.is_a?(Hash)
+        return parse_structured_result(
+          raw,
+          kind: kind,
+          expected: expected,
+          critique_candidate_ids: critique_candidate_ids
+        )
+      end
+
+      raise ParseError, "LLM returned unsupported #{kind} type: #{raw.class}" unless raw.is_a?(String)
+
+      parse_string_with_repair(
+        raw: raw,
+        kind: kind,
+        expected: expected,
+        repair_stage: repair_stage,
+        critique_candidate_ids: critique_candidate_ids
+      )
+    end
+
+    def parse_structured_result(raw, kind:, expected:, critique_candidate_ids: nil)
+      parse_expected_result(raw, expected, critique_candidate_ids: critique_candidate_ids)
+    rescue SchemaError => e
+      raise ParseError, "LLM returned invalid #{kind}: #{e.message}"
+    end
+
+    def parse_string_with_repair(raw:, kind:, expected:, repair_stage:, critique_candidate_ids: nil)
+      parse_expected_result(raw, expected, critique_candidate_ids: critique_candidate_ids)
     rescue JSON::ParserError, SchemaError => e
       @logger.warn("Invalid #{kind} JSON, requesting one repair: #{e.message}")
       repaired = repair_json(
@@ -131,14 +158,14 @@ def parse_with_repair(raw:, kind:, expected:, repair_stage:, critique_candidate_
         critique_candidate_ids: critique_candidate_ids
       )
       begin
-        parse_expected_json(repaired, expected, critique_candidate_ids: critique_candidate_ids)
+        parse_expected_result(repaired, expected, critique_candidate_ids: critique_candidate_ids)
       rescue JSON::ParserError, SchemaError => second_error
         raise ParseError, "LLM returned invalid #{kind} JSON after repair: #{second_error.message}"
       end
     end
 
-    def parse_expected_json(raw, expected, critique_candidate_ids: nil)
-      parsed = JSON.parse(strip_code_fences(raw.to_s))
+    def parse_expected_result(raw, expected, critique_candidate_ids: nil)
+      parsed = raw.is_a?(Hash) ? raw : JSON.parse(strip_code_fences(raw.to_s))
 
       case expected
       when :generate
@@ -192,15 +219,15 @@ def normalize_critique_result(parsed, critique_candidate_ids:)
     end
 
     def validate_generate_result_shape!(parsed)
-      return if parsed.is_a?(Hash) && parsed['candidates'].is_a?(Array)
-
-      raise SchemaError, 'expected an object with summary and candidates array'
+      valid_shape = parsed.is_a?(Hash) && parsed['candidates'].is_a?(Array)
+      raise SchemaError, 'expected an object with summary and candidates array' unless valid_shape
+      raise SchemaError, 'each generate candidate must be an object' unless parsed['candidates'].all?(Hash)
     end
 
     def validate_critique_result_shape!(parsed)
-      return if parsed.is_a?(Hash) && parsed['verdicts'].is_a?(Array)
-
-      raise SchemaError, 'expected an object with verdicts array'
+      valid_shape = parsed.is_a?(Hash) && parsed['verdicts'].is_a?(Array)
+      raise SchemaError, 'expected an object with verdicts array' unless valid_shape
+      raise SchemaError, 'each critique verdict must be an object' unless parsed['verdicts'].all?(Hash)
     end
 
     def validate_identifiers!(identifiers, missing_message:, duplicate_prefix:)

spec/review_pipeline_spec.rb

@@ -68,6 +68,13 @@ def generate_result(candidates)
     )
   end
 
+  def structured_generate_result(candidates)
+    {
+      'summary' => 'MR recalculates order totals during checkout.',
+      'candidates' => candidates.map { |candidate| candidate.transform_keys(&:to_s) }
+    }
+  end
+
   it 'renders only candidates accepted by critique' do
     allow(reviewer).to receive(:generate).and_return(generate_result(candidates))
     allow(reviewer).to receive(:critique).and_return(
@@ -104,6 +111,95 @@ def generate_result(candidates)
     expect(result).to include('ok')
   end
 
+  it 'processes structured Hash responses without JSON parsing or repair' do
+    allow(reviewer).to receive(:generate).and_return(structured_generate_result([candidates.first]))
+    allow(reviewer).to receive(:critique).and_return(
+      {
+        'verdicts' => [
+          { 'id' => 'C1', 'decision' => 'keep', 'reason' => 'confirmed by diff' }
+        ]
+      }
+    )
+    allow(JSON).to receive(:parse).and_call_original
+
+    result = pipeline.run(merge_request: merge_request, changes_text: changes_text)
+
+    expect(JSON).not_to have_received(:parse)
+    expect(reviewer).to have_received(:generate).once
+    expect(reviewer).to have_received(:critique).once
+    expect(result).to include('Tax is no longer included')
+  end
+
+  it 'continues processing JSON strings without repair' do
+    allow(reviewer).to receive(:generate).and_return(generate_result([candidates.first]))
+
+    result = pipeline.run(merge_request: merge_request, changes_text: changes_text, critique: false)
+
+    expect(reviewer).to have_received(:generate).once
+    expect(result).to include('Tax is no longer included')
+  end
+
+  it 'rejects an invalid structured generate result without repair' do
+    invalid_result = structured_generate_result([{ file: 'app/models/order.rb' }])
+    allow(reviewer).to receive(:generate).and_return(invalid_result)
+
+    expect do
+      pipeline.run(merge_request: merge_request, changes_text: changes_text, critique: false)
+    end.to raise_error(
+      Aireview::ParseError,
+      /invalid generate result: each generate candidate must include a non-empty id/
+    )
+    expect(reviewer).to have_received(:generate).once
+  end
+
+  it 'rejects malformed structured generate candidates without repair' do
+    allow(reviewer).to receive(:generate).and_return(
+      { 'summary' => 'Malformed result', 'candidates' => [nil] }
+    )
+
+    expect do
+      pipeline.run(merge_request: merge_request, changes_text: changes_text, critique: false)
+    end.to raise_error(
+      Aireview::ParseError,
+      /invalid generate result: each generate candidate must be an object/
+    )
+    expect(reviewer).to have_received(:generate).once
+  end
+
+  it 'rejects an unsupported response type without repair' do
+    allow(reviewer).to receive(:generate).and_return(nil)
+
+    expect do
+      pipeline.run(merge_request: merge_request, changes_text: changes_text, critique: false)
+    end.to raise_error(Aireview::ParseError, /unsupported generate result type: NilClass/)
+    expect(reviewer).to have_received(:generate).once
+  end
+
+  it 'rejects an invalid structured critique result without repair' do
+    allow(reviewer).to receive(:generate).and_return(structured_generate_result([candidates.first]))
+    allow(reviewer).to receive(:critique).and_return(
+      { 'verdicts' => [{ 'id' => 'C9', 'decision' => 'keep', 'reason' => 'hallucinated id' }] }
+    )
+
+    expect do
+      pipeline.run(merge_request: merge_request, changes_text: changes_text)
+    end.to raise_error(Aireview::ParseError, /invalid critique result: unknown verdict ids: C9/)
+    expect(reviewer).to have_received(:critique).once
+  end
+
+  it 'rejects malformed structured critique verdicts without repair' do
+    allow(reviewer).to receive(:generate).and_return(structured_generate_result([candidates.first]))
+    allow(reviewer).to receive(:critique).and_return({ 'verdicts' => [42] })
+
+    expect do
+      pipeline.run(merge_request: merge_request, changes_text: changes_text)
+    end.to raise_error(
+      Aireview::ParseError,
+      /invalid critique result: each critique verdict must be an object/
+    )
+    expect(reviewer).to have_received(:critique).once
+  end
+
   it 'repairs invalid generate JSON once' do
     allow(reviewer).to receive(:generate).and_return('not json', generate_result([candidates.first]))
 
@@ -113,6 +209,15 @@ def generate_result(candidates)
     expect(result).to include('Tax is no longer included')
   end
 
+  it 'accepts a structured Hash returned by a repair request for a JSON string' do
+    allow(reviewer).to receive(:generate).and_return('not json', structured_generate_result([candidates.first]))
+
+    result = pipeline.run(merge_request: merge_request, changes_text: changes_text, critique: false)
+
+    expect(reviewer).to have_received(:generate).twice
+    expect(result).to include('Tax is no longer included')
+  end
+
   it 'repairs generate candidates with invalid ids once' do
     invalid_generate = JSON.generate(summary: 'bad', candidates: [{ file: 'app/models/order.rb' }])
     allow(reviewer).to receive(:generate).and_return(invalid_generate, generate_result([candidates.first]))