fix: preserve structuredDocumentInfo/unstructuredDocumentInfo references in GoogleConversationalSearchTracer

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<Hash>, 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,

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