improve Grounding for Google Search and Vertex AI Search

mscraftsman · Apr 12, 2024 · a923ca0 · a923ca0
1 parent 99f5eab
commit a923ca0
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 8 deletions.
diff --git a/src/Mscc.GenerativeAI/CHANGELOG.md b/src/Mscc.GenerativeAI/CHANGELOG.md
@@ -15,6 +15,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - implement Server-Sent Events (SSE)
 
 ### Changed
+
+- improve Grounding for Google Search and Vertex AI Search
+
 ### Fixed
 
 ## 1.1.2

diff --git a/src/Mscc.GenerativeAI/GenerativeModel.cs b/src/Mscc.GenerativeAI/GenerativeModel.cs
@@ -628,21 +628,45 @@ public async Task<GenerateContentResponse> GenerateContent(GenerateContentReques
             if (_useVertexAi)
             {
                 var fullText = new StringBuilder();
-                var contentResponseVertex = await Deserialize<List<GenerateContentResponse>>(response);
-                foreach (var item in contentResponseVertex)
+                GroundingMetadata groundingMetadata = null;
+                var contents = await Deserialize<List<GenerateContentResponse>>(response);
+                foreach (var content in contents)
                 {
-                    switch (item.Candidates[0].FinishReason)
+                    if (!(content.Candidates?[0].GroundingMetadata is null))
+                    {
+                        groundingMetadata = content.Candidates[0].GroundingMetadata;
+                        continue;
+                    }
+
+                    switch (content.Candidates?[0].FinishReason)
                     {
                         case FinishReason.Safety:
-                            return item;
+                            return content;
                             break;
                         case FinishReason.Unspecified:
                         default:
-                            fullText.Append(item.Text);
+                            fullText.Append(content.Text);
                             break;
                     }
                 }
-                var result = contentResponseVertex.LastOrDefault();
+                var result = contents.LastOrDefault();
+                if (result.Candidates is null)
+                {
+                    result.Candidates = new List<Candidate>()
+                    {
+                        new()
+                        {
+                            Content = new ContentResponse()
+                            {
+                                Parts = new List<Part>()
+                                {
+                                    new()
+                                }
+                            }
+                        }
+                    };
+                }
+                result.Candidates[0].GroundingMetadata = groundingMetadata;
                 result.Candidates[0].Content.Parts[0].Text = fullText.ToString();
                 return result;
             }

diff --git a/tests/Mscc.GenerativeAI/VertexAi_GeminiPro_Should.cs b/tests/Mscc.GenerativeAI/VertexAi_GeminiPro_Should.cs
@@ -365,7 +365,13 @@ public async void Generate_Content_WithGrounding_GoogleSearch()
             response.Text.Should().NotBeEmpty();
             _output.WriteLine(response?.Text);
             response.Candidates[0].GroundingMetadata.Should().NotBeNull();
-            _output.WriteLine($"GroundingMetadata: {response.Candidates[0].GroundingMetadata}");
+            response.Candidates[0].GroundingMetadata.WebSearchQueries.Should().NotBeNull();
+            response.Candidates[0].GroundingMetadata.WebSearchQueries.Count.Should().BeGreaterOrEqualTo(1);
+            _output.WriteLine($"{new string('-', 20)}");
+            foreach (string query in response.Candidates[0].GroundingMetadata.WebSearchQueries)
+            {
+                _output.WriteLine($"SearchQuery: {query}");
+            }
         }
 
         [Fact]
@@ -404,7 +410,13 @@ public async void Generate_Content_WithGrounding_VertexAiSearch()
             response.Text.Should().NotBeEmpty();
             _output.WriteLine(response?.Text);
             response.Candidates[0].GroundingMetadata.Should().NotBeNull();
-            _output.WriteLine($"GroundingMetadata: {response.Candidates[0].GroundingMetadata}");
+            response.Candidates[0].GroundingMetadata.WebSearchQueries.Should().NotBeNull();
+            response.Candidates[0].GroundingMetadata.WebSearchQueries.Count.Should().BeGreaterOrEqualTo(1);
+            _output.WriteLine($"{new string('-', 20)}");
+            foreach (string query in response.Candidates[0].GroundingMetadata.WebSearchQueries)
+            {
+                _output.WriteLine($"SearchQuery: {query}");
+            }
         }
 
         [Theory]