Handle malformed gemini replies

Author: SamSaffron

lib/ai_bot/tools/create_artifact.rb

@@ -21,8 +21,8 @@ def self.signature
                 required: true,
               },
               {
-                name: "html_content",
-                description: "The HTML content for the artifact",
+                name: "html_body",
+                description: "The HTML content for the BODY tag (do not include the BODY tag)",
                 type: "string",
                 required: true,
               },
@@ -43,7 +43,7 @@ def invoke
           post = Post.find_by(id: context[:post_id])
           return error_response("No post context found") unless post
 
-          html = parameters[:html_content].to_s
+          html = parameters[:html_body].to_s
           css = parameters[:css].to_s
           js = parameters[:js].to_s
 

lib/completions/endpoints/gemini.rb

@@ -173,25 +173,24 @@ def decode_chunk(chunk)
             .decode(chunk)
             .map do |parsed|
               update_usage(parsed)
-              parsed
-                .dig(:candidates, 0, :content, :parts)
-                .map do |part|
-                  if part[:text]
-                    part = part[:text]
-                    if part != ""
-                      part
-                    else
-                      nil
-                    end
-                  elsif part[:functionCall]
-                    @tool_index += 1
-                    ToolCall.new(
-                      id: "tool_#{@tool_index}",
-                      name: part[:functionCall][:name],
-                      parameters: part[:functionCall][:args],
-                    )
+              parts = parsed.dig(:candidates, 0, :content, :parts)
+              parts&.map do |part|
+                if part[:text]
+                  part = part[:text]
+                  if part != ""
+                    part
+                  else
+                    nil
                   end
+                elsif part[:functionCall]
+                  @tool_index += 1
+                  ToolCall.new(
+                    id: "tool_#{@tool_index}",
+                    name: part[:functionCall][:name],
+                    parameters: part[:functionCall][:args],
+                  )
                 end
+              end
             end
             .flatten
             .compact

spec/lib/completions/endpoints/gemini_spec.rb

@@ -324,6 +324,37 @@ def tool_response
     expect(log.response_tokens).to eq(4)
   end
 
+  it "Can correctly handle malformed responses" do
+    response = <<~TEXT
+      data: {"candidates": [{"content": {"parts": [{"text": "Certainly"}],"role": "model"}}],"usageMetadata": {"promptTokenCount": 399,"totalTokenCount": 399},"modelVersion": "gemini-1.5-pro-002"}
+
+      data: {"candidates": [{"content": {"parts": [{"text": "! I'll create a simple \\"Hello, World!\\" page where each letter"}],"role": "model"},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"}]}],"usageMetadata": {"promptTokenCount": 399,"totalTokenCount": 399},"modelVersion": "gemini-1.5-pro-002"}
+
+      data: {"candidates": [{"content": {"parts": [{"text": " has a different color using inline styles for simplicity.  Each letter will be wrapped"}],"role": "model"},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"}]}],"usageMetadata": {"promptTokenCount": 399,"totalTokenCount": 399},"modelVersion": "gemini-1.5-pro-002"}
+
+      data: {"candidates": [{"content": {"parts": [{"text": ""}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
+
+      data: {"candidates": [{"finishReason": "MALFORMED_FUNCTION_CALL"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
+
+    TEXT
+
+    llm = DiscourseAi::Completions::Llm.proxy("custom:#{model.id}")
+    url = "#{model.url}:streamGenerateContent?alt=sse&key=123"
+
+    output = []
+
+    stub_request(:post, url).to_return(status: 200, body: response)
+    llm.generate("Hello", user: user) { |partial| output << partial }
+
+    expect(output).to eq(
+      [
+        "Certainly",
+        "! I'll create a simple \"Hello, World!\" page where each letter",
+        " has a different color using inline styles for simplicity.  Each letter will be wrapped",
+      ],
+    )
+  end
+
   it "Can correctly handle streamed responses even if they are chunked badly" do
     data = +""
     data << "da|ta: |"