Fix some broken Gemini tests, make SSEStreamParsing more resilient

Author: thmsmlr

lib/instructor/adapters/gemini.ex

@@ -23,10 +23,11 @@ defmodule Instructor.Adapters.Gemini do
   """
 
   @behaviour Instructor.Adapter
+  alias Instructor.SSEStreamParser
   alias Instructor.Adapters
   alias Instructor.JSONSchema
 
-  @supported_modes [:tools, :json_schema]
+  @supported_modes [:json_schema]
 
   @doc """
   Run a completion against Google's Gemini API
@@ -102,7 +103,7 @@ defmodule Instructor.Adapters.Gemini do
           generation_config =
             generation_config
             |> Map.put("response_mime_type", "application/json")
-            |> Map.put("response_schema", map_schema(schema))
+            |> Map.put("response_schema", normalize_json_schema(schema))
 
           params
           |> Map.put(:generationConfig, generation_config)
@@ -116,7 +117,7 @@ defmodule Instructor.Adapters.Gemini do
                   %{
                     name: tool["name"],
                     description: tool["description"],
-                    parameters: map_schema(tool["parameters"])
+                    parameters: normalize_json_schema(tool["parameters"])
                   }
                 end)
             }
@@ -155,23 +156,7 @@ defmodule Instructor.Adapters.Gemini do
               json: params,
               rpc_function: :streamGenerateContent,
               into: fn {:data, data}, {req, resp} ->
-                chunks =
-                  data
-                  |> String.split("\n")
-                  |> Enum.filter(fn line ->
-                    String.starts_with?(line, "data: {")
-                  end)
-                  |> Enum.map(fn line ->
-                    line
-                    |> String.replace_prefix("data: ", "")
-                    |> Jason.decode!()
-                    |> then(&parse_stream_chunk_for_mode(mode, &1))
-                  end)
-
-                for chunk <- chunks do
-                  send(pid, chunk)
-                end
-
+                send(pid, data)
                 {:cont, {req, resp}}
               end
             )
@@ -196,6 +181,8 @@ defmodule Instructor.Adapters.Gemini do
       end,
       fn task -> Task.await(task) end
     )
+    |> SSEStreamParser.parse()
+    |> Stream.map(fn chunk -> parse_stream_chunk_for_mode(mode, chunk) end)
   end
 
   defp do_chat_completion(mode, params, config) do
@@ -266,18 +253,70 @@ defmodule Instructor.Adapters.Gemini do
     chunk
   end
 
-  defp map_schema(schema) do
-    JSONSchema.traverse_and_update(schema, fn
-      %{"type" => _} = x
-      when is_map_key(x, "format") or is_map_key(x, "pattern") or
-             is_map_key(x, "title") or is_map_key(x, "additionalProperties") ->
-        Map.drop(x, ["format", "pattern", "title", "additionalProperties"])
+  defp normalize_json_schema(schema) do
+    JSONSchema.traverse_and_update(
+      schema,
+      fn
+        {%{"type" => _} = x, path}
+        when is_map_key(x, "format") or is_map_key(x, "pattern") or
+               is_map_key(x, "title") or is_map_key(x, "additionalProperties") ->
+          x
+          |> Map.drop(["format", "pattern", "title", "additionalProperties"])
+          |> case do
+            %{"type" => "object", "properties" => properties} when map_size(properties) == 0 ->
+              raise """
+              Invalid JSON Schema: object with no properties at path: #{inspect(path)}
+
+              Gemini does not support empty objects. This is likely because have have a naked :map type
+              without any fields at #{inspect(path)}. Try switching to an embedded schema instead.
+              """
+
+            x ->
+              x
+          end
 
-      x ->
-        x
-    end)
+        {x, _path} ->
+          x
+      end,
+      include_path: true
+    )
+    |> inline_defs()
+  end
+
+  defp inline_defs(schema) do
+    # First extract the definitions map for reference
+    {defs, schema} = Map.pop(schema, "$defs", %{})
+
+    # Traverse and replace all $refs with their definitions
+    traverse_and_inline(schema, defs)
+  end
+
+  defp traverse_and_inline(schema, defs) when is_map(schema) do
+    cond do
+      # If we find a $ref, replace it with the inlined definition
+      Map.has_key?(schema, "$ref") ->
+        ref = schema["$ref"]
+        def_key = String.replace_prefix(ref, "#/$defs/", "")
+        definition = Map.get(defs, def_key, %{})
+        # Recursively inline any nested refs in the definition
+        traverse_and_inline(definition, defs)
+
+      # Otherwise traverse all values in the map
+      true ->
+        schema
+        |> Enum.map(fn {k, v} -> {k, traverse_and_inline(v, defs)} end)
+        |> Enum.into(%{})
+    end
   end
 
+  # Handle arrays by traversing each element
+  defp traverse_and_inline(schema, defs) when is_list(schema) do
+    Enum.map(schema, &traverse_and_inline(&1, defs))
+  end
+
+  # Base case - return non-map/list values as is
+  defp traverse_and_inline(schema, _defs), do: schema
+
   defp snake_to_camel(snake_case_string) do
     snake_case_string
     |> String.split("_")

lib/instructor/json_schema.ex

@@ -83,7 +83,7 @@ defmodule Instructor.JSONSchema do
   end
 
   defp uses_use_instructor(ecto_schema) when is_ecto_schema(ecto_schema) do
-    function_exported?(ecto_schema, :__llm_doc__, 0)
+    {:__llm_doc__, 0} in ecto_schema.__info__(:functions)
   end
 
   defp uses_use_instructor(_), do: false
@@ -414,28 +414,46 @@ defmodule Instructor.JSONSchema do
 
   ## Parameters
     - tree: The tree structure to traverse (can be a map, list, or any other type)
-    - fun: A function that takes an element and returns either:
+    - fun: A function that takes either:
+      - Just the element if include_path: false (default)
+      - A tuple of {element, path} if include_path: true, where path is a list of keys to reach this element
+    The function should return either:
       - An updated element
       - nil to remove the element
       - The original element if no changes are needed
+    - opts: Optional keyword list of options
+      - include_path: boolean, when true includes the path to each element in the callback (default: false)
 
   ## Returns
     The updated tree structure
   """
-  def traverse_and_update(tree, fun) when is_map(tree) do
+  def traverse_and_update(tree, fun, opts \\ []) do
+    do_traverse_and_update(tree, fun, [], opts)
+  end
+
+  defp do_traverse_and_update(tree, fun, path, opts) when is_map(tree) do
     tree
-    |> Enum.map(fn {k, v} -> {k, traverse_and_update(v, fun)} end)
+    |> Enum.map(fn {k, v} -> {k, do_traverse_and_update(v, fun, path ++ [k], opts)} end)
     |> Enum.filter(fn {_, v} -> v != nil end)
     |> Enum.into(%{})
-    |> fun.()
+    |> maybe_call_with_path(fun, path, opts)
   end
 
-  def traverse_and_update(tree, fun) when is_list(tree) do
+  defp do_traverse_and_update(tree, fun, path, opts) when is_list(tree) do
     tree
-    |> Enum.map(fn elem -> traverse_and_update(elem, fun) end)
+    |> Enum.with_index()
+    |> Enum.map(fn {elem, idx} -> do_traverse_and_update(elem, fun, path ++ [idx], opts) end)
     |> Enum.filter(&(&1 != nil))
-    |> fun.()
+    |> maybe_call_with_path(fun, path, opts)
   end
 
-  def traverse_and_update(tree, fun), do: fun.(tree)
+  defp do_traverse_and_update(tree, fun, path, opts), do: maybe_call_with_path(tree, fun, path, opts)
+
+  defp maybe_call_with_path(value, fun, path, opts) do
+    if Keyword.get(opts, :include_path, false) do
+      fun.({value, path})
+    else
+      fun.(value)
+    end
+  end
 end

lib/instructor/sse_stream_parser.ex

@@ -18,15 +18,15 @@ defmodule Instructor.SSEStreamParser do
       fn acc -> {[acc], nil} end,
       fn _acc -> nil end
     )
-    |> Stream.filter(fn line -> line != "" end)
+    |> Stream.filter(fn line -> String.trim(line) != "" end)
     |> Stream.transform(
       fn -> {:root, ""} end,
       fn
         "data: [DONE]" <> _, {:root, ""} ->
           {:halt, {:root, ""}}
 
         "data: " <> data, {:root, ""} ->
-          {[{:ok, Jason.decode!(data)}], {:root, ""}}
+          {[{:ok, decode_json!(data)}], {:root, ""}}
 
         "event: " <> _, {_, _} ->
           {[], {:root, ""}}
@@ -36,7 +36,7 @@ defmodule Instructor.SSEStreamParser do
       end,
       fn
         {:json, acc} ->
-          {[{:error, Jason.decode!(acc)}], {:root, ""}}
+          {[{:error, decode_json!(acc)}], {:root, ""}}
 
         {:root, ""} ->
           {[], {:root, ""}}
@@ -51,4 +51,11 @@ defmodule Instructor.SSEStreamParser do
         raise "Error from LLM: #{inspect(error)}"
     end)
   end
+
+  defp decode_json!(data) do
+    case Jason.decode(data) do
+      {:ok, decoded} -> decoded
+      {:error, err} -> raise "Error decoding: #{inspect(err)} \n\n #{inspect(data)}"
+    end
+  end
 end

lib/instructor/validator.ex

@@ -60,19 +60,6 @@ defmodule Instructor.Validator do
     end
   end
 
-  defmodule Validation do
-    use Ecto.Schema
-
-    @doc """
-    Validate if an attribute is correct and if not, return an error message
-    """
-    @primary_key false
-    embedded_schema do
-      field(:valid?, :boolean)
-      field(:reason, :string)
-    end
-  end
-
   @doc """
   Validate a changeset field using a language model
 
@@ -95,6 +82,20 @@ defmodule Instructor.Validator do
     end
   """
   def validate_with_llm(changeset, field, statement, opts \\ []) do
+    defmodule Validation do
+      use Ecto.Schema
+      use Instructor
+
+      @llm_doc """
+      Validate if an attribute is correct and if not, return an error message
+      """
+      @primary_key false
+      embedded_schema do
+        field(:valid?, :boolean)
+        field(:reason, :string)
+      end
+    end
+
     Ecto.Changeset.validate_change(changeset, field, fn field, value ->
       {:ok, response} =
         Instructor.chat_completion(
@@ -116,10 +117,10 @@ defmodule Instructor.Validator do
         )
 
       case response do
-        %Validation{valid?: true} ->
+        %{valid?: true} ->
           []
 
-        %Validation{reason: reason} ->
+        %{reason: reason} ->
           [
             {field, "is invalid, #{reason}"}
           ]

test/instructor_test.exs

@@ -268,7 +268,7 @@ defmodule InstructorTest do
               stream: true,
               response_model: {:array, President},
               messages: [
-                %{role: "user", content: "What are the first 3 presidents of the United States?"}
+                %{role: "user", content: "Who were the first 3 presidents of the United States?"}
               ]
             )
           )

test/json_schema_test.exs

@@ -74,12 +74,12 @@ defmodule JSONSchemaTest do
 
   test "includes documentation" do
     json_schema =
-      JSONSchema.from_ecto_schema(InstructorTest.DemoWithDocumentation)
+      JSONSchema.from_ecto_schema(InstructorTest.DemoWithUseInstructorAndNewDoc)
       |> Jason.decode!()
 
     expected_json_schema =
       %{
-        "description" => "Hello world\n",
+        "description" => "Hello world",
         "properties" => %{
           "string" => %{
             "title" => "string",
@@ -88,7 +88,7 @@ defmodule JSONSchemaTest do
           }
         },
         "required" => ["string"],
-        "title" => "InstructorTest.DemoWithDocumentation",
+        "title" => "InstructorTest.DemoWithUseInstructorAndNewDoc",
         "type" => "object",
         "additionalProperties" => false
       }
@@ -140,7 +140,7 @@ defmodule JSONSchemaTest do
   test "basic types" do
     defmodule Demo do
       use Ecto.Schema
-
+      use Instructor
       # Be explicit about all fields in this test
       @primary_key false
       embedded_schema do

test/support/test_schemas.ex

@@ -41,11 +41,10 @@ defmodule InstructorTest.DemoWithUseInstructorButOldDoc do
 end
 
 defmodule InstructorTest.DemoWithUseInstructorAndNewDoc do
-  use Instructor
   use Ecto.Schema
+  use Instructor
 
   @llm_doc "Hello world"
-
   @primary_key false
   embedded_schema do
     field(:string, :string)