diff --git a/lib/openlayer/integrations/google_conversational_search_tracer.rb b/lib/openlayer/integrations/google_conversational_search_tracer.rb index 23ed398..5103c59 100644 --- a/lib/openlayer/integrations/google_conversational_search_tracer.rb +++ b/lib/openlayer/integrations/google_conversational_search_tracer.rb @@ -345,6 +345,9 @@ def self.extract_citation_sources(citation) # Extract references from answer # + # Each reference's `content` is a oneof: exactly one of chunk_info, + # structured_document_info, or unstructured_document_info is present. + # # @param answer [Object] Answer object from response # @return [Array, nil] Array of reference hashes or nil def self.extract_references(answer) @@ -354,25 +357,85 @@ def self.extract_references(answer) return nil if references.nil? || !references.respond_to?(:each_with_index) references.each_with_index.map do |reference, index| - chunk_info = safe_extract(reference, :chunk_info) - next nil if chunk_info.nil? - - doc_metadata = safe_extract(chunk_info, :document_metadata) - - { - reference_id: index.to_s, - content: safe_extract(chunk_info, :content), - relevance_score: safe_extract(chunk_info, :relevance_score)&.to_f, - document_id: doc_metadata ? safe_extract(doc_metadata, :document) : nil, - uri: doc_metadata ? safe_extract(doc_metadata, :uri) : nil, - title: doc_metadata ? safe_extract(doc_metadata, :title) : nil - }.compact + extract_reference(reference, index) end.compact rescue StandardError => e warn_if_debug("[Openlayer] Failed to extract references: #{e.message}") nil end + # Extract a single reference hash from whichever `content` oneof + # variant is present on it + # + # @param reference [Object] Reference object + # @param index [Integer] Reference's index in the references array + # @return [Hash, nil] Reference hash, or nil if no known variant is set + def self.extract_reference(reference, index) + if (chunk_info = safe_extract(reference, :chunk_info)) + extract_chunk_info_reference(chunk_info, index) + elsif (structured_document_info = safe_extract(reference, :structured_document_info)) + extract_structured_document_info_reference(structured_document_info, index) + elsif (unstructured_document_info = safe_extract(reference, :unstructured_document_info)) + extract_unstructured_document_info_reference(unstructured_document_info, index) + end + rescue StandardError => e + warn_if_debug("[Openlayer] Failed to extract reference: #{e.message}") + nil + end + + # Build a reference hash from a ChunkInfo variant + # + # @param chunk_info [Object] ChunkInfo object + # @param index [Integer] Reference's index in the references array + # @return [Hash] Reference hash + def self.extract_chunk_info_reference(chunk_info, index) + doc_metadata = safe_extract(chunk_info, :document_metadata) + + { + reference_id: index.to_s, + content: safe_extract(chunk_info, :content), + relevance_score: safe_extract(chunk_info, :relevance_score)&.to_f, + document_id: doc_metadata ? safe_extract(doc_metadata, :document) : nil, + uri: doc_metadata ? safe_extract(doc_metadata, :uri) : nil, + title: doc_metadata ? safe_extract(doc_metadata, :title) : nil + }.compact + end + + # Build a reference hash from a StructuredDocumentInfo variant + # + # @param structured_document_info [Object] StructuredDocumentInfo object + # @param index [Integer] Reference's index in the references array + # @return [Hash] Reference hash + def self.extract_structured_document_info_reference(structured_document_info, index) + { + reference_id: index.to_s, + document_id: safe_extract(structured_document_info, :document), + uri: safe_extract(structured_document_info, :uri), + title: safe_extract(structured_document_info, :title), + struct_data: safe_extract(structured_document_info, :struct_data) + }.compact + end + + # Build a reference hash from an UnstructuredDocumentInfo variant + # + # @param unstructured_document_info [Object] UnstructuredDocumentInfo object + # @param index [Integer] Reference's index in the references array + # @return [Hash] Reference hash + def self.extract_unstructured_document_info_reference(unstructured_document_info, index) + chunk_contents = safe_extract(unstructured_document_info, :chunk_contents) + content = if chunk_contents.respond_to?(:map) + chunk_contents.map { |chunk| safe_extract(chunk, :content) }.compact.join("\n") + end + + { + reference_id: index.to_s, + content: (content unless content.nil? || content.empty?), + document_id: safe_extract(unstructured_document_info, :document), + uri: safe_extract(unstructured_document_info, :uri), + title: safe_extract(unstructured_document_info, :title) + }.compact + end + # Extract related questions from answer # # @param answer [Object] Answer object from response @@ -717,6 +780,10 @@ def self.warn_if_debug(message) :extract_citations, :extract_citation_sources, :extract_references, + :extract_reference, + :extract_chunk_info_reference, + :extract_structured_document_info_reference, + :extract_unstructured_document_info_reference, :extract_related_questions, :extract_steps, :extract_step_data, diff --git a/test/openlayer/integrations/google_conversational_search_tracer_test.rb b/test/openlayer/integrations/google_conversational_search_tracer_test.rb index 05620c0..037dc5b 100644 --- a/test/openlayer/integrations/google_conversational_search_tracer_test.rb +++ b/test/openlayer/integrations/google_conversational_search_tracer_test.rb @@ -13,19 +13,83 @@ module Integrations class Openlayer::Test::Integrations::GoogleConversationalSearchTracerTest < Minitest::Test Tracer = Openlayer::Integrations::GoogleConversationalSearchTracer + class FakeDocumentMetadata + attr_reader :document, :uri, :title + + def initialize(document: nil, uri: nil, title: nil) + @document = document + @uri = uri + @title = title + end + end + + class FakeChunkInfo + attr_reader :content, :relevance_score, :document_metadata + + def initialize(content: nil, relevance_score: nil, document_metadata: nil) + @content = content + @relevance_score = relevance_score + @document_metadata = document_metadata + end + end + + class FakeStructuredDocumentInfo + attr_reader :document, :struct_data, :title, :uri + + def initialize(document: nil, struct_data: nil, title: nil, uri: nil) + @document = document + @struct_data = struct_data + @title = title + @uri = uri + end + end + + class FakeChunkContent + attr_reader :content + + def initialize(content: nil) + @content = content + end + end + + class FakeUnstructuredDocumentInfo + attr_reader :document, :uri, :title, :chunk_contents + + def initialize(document: nil, uri: nil, title: nil, chunk_contents: []) + @document = document + @uri = uri + @title = title + @chunk_contents = chunk_contents + end + end + + # Duck-types Discovery Engine's Answer::Reference, whose `content` field is + # a oneof: exactly one of chunk_info / structured_document_info / + # unstructured_document_info is non-nil, mirroring real protobuf behavior. + class FakeReference + attr_reader :chunk_info, :structured_document_info, :unstructured_document_info + + def initialize(chunk_info: nil, structured_document_info: nil, unstructured_document_info: nil) + @chunk_info = chunk_info + @structured_document_info = structured_document_info + @unstructured_document_info = unstructured_document_info + end + end + class FakeAnswer - attr_reader :answer_text + attr_reader :answer_text, :references - def initialize(answer_text) + def initialize(answer_text, references: []) @answer_text = answer_text + @references = references end end class FakeResponse attr_reader :answer - def initialize(answer_text) - @answer = FakeAnswer.new(answer_text) + def initialize(answer_text, references: []) + @answer = FakeAnswer.new(answer_text, references: references) end end @@ -141,4 +205,93 @@ def test_trace_client_strips_additional_columns_before_forwarding_to_google_clie assert_equal("hi", response.answer.answer_text) assert_equal("abc-123", @openlayer_client.last_row[:trace_id]) end + + def test_chunk_info_reference_is_preserved + reference = FakeReference.new( + chunk_info: FakeChunkInfo.new( + content: "chunk text", + relevance_score: 0.9, + document_metadata: FakeDocumentMetadata.new(document: "doc-1", uri: "https://a.test", title: "A") + ) + ) + row = trace_row(response: FakeResponse.new("hi", references: [reference])) + + references = row[:steps][0][:references] + assert_equal(1, references.length) + assert_equal("chunk text", references[0][:content]) + assert_equal(["chunk text"], row[:context]) + end + + def test_structured_document_info_reference_is_preserved + document_id = "projects/example/dataStores/example/documents/example-doc" + reference = FakeReference.new( + structured_document_info: FakeStructuredDocumentInfo.new( + document: document_id, + uri: "https://example.test/doc", + title: "Example structured document", + struct_data: {kind: "example"} + ) + ) + row = trace_row(response: FakeResponse.new("hi", references: [reference])) + + references = row[:steps][0][:references] + assert_equal(1, references.length) + assert_equal(document_id, references[0][:document_id]) + assert_equal("https://example.test/doc", references[0][:uri]) + assert_equal("Example structured document", references[0][:title]) + assert_equal({kind: "example"}, references[0][:struct_data]) + end + + def test_unstructured_document_info_reference_is_preserved + reference = FakeReference.new( + unstructured_document_info: FakeUnstructuredDocumentInfo.new( + document: "doc-2", + uri: "https://b.test", + title: "B", + chunk_contents: [ + FakeChunkContent.new(content: "chunk one"), + FakeChunkContent.new(content: "chunk two") + ] + ) + ) + row = trace_row(response: FakeResponse.new("hi", references: [reference])) + + references = row[:steps][0][:references] + assert_equal(1, references.length) + assert_equal("doc-2", references[0][:document_id]) + assert_equal("https://b.test", references[0][:uri]) + assert_equal("B", references[0][:title]) + assert_equal("chunk one\nchunk two", references[0][:content]) + assert_equal(["chunk one\nchunk two"], row[:context]) + end + + def test_reference_with_no_known_variant_is_skipped_without_raising + row = trace_row(response: FakeResponse.new("hi", references: [FakeReference.new])) + + assert_nil(row[:steps][0][:references]) + assert_nil(row[:context]) + end + + def test_mixed_reference_variants_are_all_preserved_in_order + chunk_ref = FakeReference.new(chunk_info: FakeChunkInfo.new(content: "chunk")) + structured_ref = FakeReference.new( + structured_document_info: FakeStructuredDocumentInfo.new(title: "Structured") + ) + row = trace_row(response: FakeResponse.new("hi", references: [chunk_ref, structured_ref])) + + references = row[:steps][0][:references] + assert_equal(2, references.length) + assert_equal("chunk", references[0][:content]) + assert_equal("Structured", references[1][:title]) + end + + def test_references_count_metadata_matches_traced_references_for_structured_document_info + reference = FakeReference.new( + structured_document_info: FakeStructuredDocumentInfo.new(title: "Structured") + ) + row = trace_row(response: FakeResponse.new("hi", references: [reference])) + + assert_equal(1, row[:steps][0][:metadata][:references_count]) + assert_equal(1, row[:steps][0][:references].length) + end end