feat: Add structured output support with generate_program_structured! for E2E tests (#65) (#67)

* feat: Add structured output support with generate_program_structured! for E2E tests

- Implement PtcRunner.Schema.to_llm_schema/0 with flattened schema optimized for LLMs
- Use anyOf pattern to list all operations at top level, avoiding  recursion
- Add PtcRunner.TestSupport.LLMClient.generate_program_structured!/1 for structured output mode
- Uses ReqLLM.generate_object! for guaranteed valid JSON from LLMs
- Add E2E tests for both text mode (text mode) and structured output mode
- E2E tests verify filter, sum, and chained operations work correctly
- Addresses issue #65 review comments on schema validation and error handling

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>

* test: Add comprehensive unit tests for to_llm_schema/0

Adds 13 new tests for the to_llm_schema/0 function following the same
patterns as the existing to_json_schema/0 tests. Tests verify:
- Basic schema structure with title, type, properties, anyOf
- Program property with correct structure
- Correct count of 33 operation schemas
- Each operation has required structure
- Literal operation schema correctness
- Operations with no fields have only 'op' as required
- Operations with optional fields exclude them from required
- Expr types use inline object schema (not $ref)
- List of expr types have correct inline schema
- List of string types have correct schema
- Non-negative integer types have minimum constraint
- Recursive operations like pipe validate correctly
- Generated schema is valid JSON

Fixes issue from PR #67 review.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: 3 people

lib/ptc_runner/schema.ex

@@ -352,6 +352,121 @@ defmodule PtcRunner.Schema do
     base_required ++ field_required
   end
 
+  @doc """
+  Generate a flattened JSON Schema optimized for LLM structured output.
+
+  This schema uses `anyOf` to list all operations at the top level, avoiding
+  the recursive `$ref` patterns that LLMs struggle with. The schema is designed
+  to work with ReqLLM.generate_object! for structured output mode.
+
+  ## Returns
+    A map representing the flattened JSON Schema for the PTC DSL.
+  """
+  @spec to_llm_schema() :: map()
+  def to_llm_schema do
+    operation_schemas =
+      @operations
+      |> Enum.map(fn {op_name, op_def} ->
+        operation_to_llm_schema(op_name, op_def)
+      end)
+
+    %{
+      "title" => "PTC Program",
+      "type" => "object",
+      "properties" => %{
+        "program" => %{
+          "description" => "The PTC program operation",
+          "anyOf" => operation_schemas
+        }
+      },
+      "required" => ["program"],
+      "additionalProperties" => false
+    }
+  end
+
+  # Convert a single operation to its flattened schema representation for LLM use
+  defp operation_to_llm_schema(op_name, op_def) do
+    fields = op_def["fields"]
+    properties = build_llm_properties(op_name, fields)
+    required_fields = build_required(fields)
+
+    %{
+      "type" => "object",
+      "description" => op_def["description"],
+      "properties" => properties,
+      "required" => required_fields,
+      "additionalProperties" => false
+    }
+  end
+
+  # Build the properties map for an LLM schema operation (flattened, no $ref)
+  defp build_llm_properties(op_name, fields) do
+    base_properties = %{
+      "op" => %{"const" => op_name}
+    }
+
+    field_properties =
+      fields
+      |> Enum.map(fn {field_name, field_spec} ->
+        {field_name, type_to_llm_json_schema(field_spec["type"])}
+      end)
+      |> Enum.into(%{})
+
+    Map.merge(base_properties, field_properties)
+  end
+
+  # Convert Elixir type to flattened JSON Schema type for LLM use (no $ref)
+  defp type_to_llm_json_schema(:any), do: %{}
+
+  defp type_to_llm_json_schema(:string) do
+    %{"type" => "string"}
+  end
+
+  # For nested expressions, use a placeholder schema instead of $ref
+  # This allows the LLM to generate valid operations
+  defp type_to_llm_json_schema(:expr) do
+    %{
+      "description" => "A PTC operation (any valid operation type)",
+      "type" => "object",
+      "properties" => %{
+        "op" => %{"type" => "string"}
+      },
+      "required" => ["op"]
+    }
+  end
+
+  defp type_to_llm_json_schema({:list, :expr}) do
+    %{
+      "type" => "array",
+      "items" => %{
+        "description" => "A PTC operation (any valid operation type)",
+        "type" => "object",
+        "properties" => %{
+          "op" => %{"type" => "string"}
+        },
+        "required" => ["op"]
+      }
+    }
+  end
+
+  defp type_to_llm_json_schema({:list, :string}) do
+    %{
+      "type" => "array",
+      "items" => %{"type" => "string"}
+    }
+  end
+
+  defp type_to_llm_json_schema(:map) do
+    %{"type" => "object"}
+  end
+
+  defp type_to_llm_json_schema(:non_neg_integer) do
+    %{
+      "type" => "integer",
+      "minimum" => 0
+    }
+  end
+
   # Convert Elixir type to JSON Schema type specification
   defp type_to_json_schema(:any), do: %{}
 

test/ptc_runner/e2e_test.exs

@@ -5,7 +5,7 @@ defmodule PtcRunner.E2ETest do
 
   alias PtcRunner.TestSupport.LLMClient
 
-  describe "LLM program generation" do
+  describe "LLM program generation - text mode" do
     test "generates valid filter program" do
       task = "Filter products where price is greater than 10"
       json_schema = PtcRunner.Schema.to_json_schema()
@@ -64,4 +64,61 @@ defmodule PtcRunner.E2ETest do
       assert result == 2
     end
   end
+
+  describe "LLM program generation - structured output mode" do
+    test "generates valid filter program with structured output" do
+      task = "Filter products where price is greater than 10"
+
+      program_json = LLMClient.generate_program_structured!(task)
+
+      context = %{
+        "input" => [
+          %{"name" => "Apple", "price" => 5},
+          %{"name" => "Book", "price" => 15},
+          %{"name" => "Laptop", "price" => 999}
+        ]
+      }
+
+      assert {:ok, result, _metrics} = PtcRunner.run(program_json, context: context)
+
+      # Constraint assertions - not exact values
+      assert is_list(result)
+      assert length(result) == 2
+      assert Enum.all?(result, fn item -> item["price"] > 10 end)
+    end
+
+    test "generates valid sum aggregation with structured output" do
+      task = "Calculate the sum of all prices"
+
+      program_json = LLMClient.generate_program_structured!(task)
+
+      context = %{
+        "input" => [
+          %{"price" => 10},
+          %{"price" => 20},
+          %{"price" => 30}
+        ]
+      }
+
+      assert {:ok, result, _metrics} = PtcRunner.run(program_json, context: context)
+      assert result == 60
+    end
+
+    test "generates valid chained operations with structured output" do
+      task = "Filter products where price > 10, then count them"
+
+      program_json = LLMClient.generate_program_structured!(task)
+
+      context = %{
+        "input" => [
+          %{"name" => "A", "price" => 5},
+          %{"name" => "B", "price" => 15},
+          %{"name" => "C", "price" => 25}
+        ]
+      }
+
+      assert {:ok, result, _metrics} = PtcRunner.run(program_json, context: context)
+      assert result == 2
+    end
+  end
 end

test/ptc_runner/schema_test.exs

@@ -486,4 +486,160 @@ defmodule PtcRunner.SchemaTest do
              "Generated schema does not match priv/ptc_schema.json"
     end
   end
+
+  describe "to_llm_schema/0" do
+    test "returns a valid JSON Schema structure" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      assert is_map(schema)
+      assert schema["title"] == "PTC Program"
+      assert schema["type"] == "object"
+      assert is_map(schema["properties"])
+      assert is_list(schema["properties"]["program"]["anyOf"])
+    end
+
+    test "schema has program property with correct structure" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      assert is_map(schema["properties"])
+      assert is_map(schema["properties"]["program"])
+      assert is_list(schema["properties"]["program"]["anyOf"])
+      assert schema["required"] == ["program"]
+      assert schema["additionalProperties"] == false
+    end
+
+    test "generates 33 operation schemas" do
+      schema = PtcRunner.Schema.to_llm_schema()
+      assert length(schema["properties"]["program"]["anyOf"]) == 33
+    end
+
+    test "each operation schema has required structure" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      Enum.each(schema["properties"]["program"]["anyOf"], fn op_schema ->
+        assert is_map(op_schema)
+        assert op_schema["type"] == "object"
+        assert is_map(op_schema["properties"])
+        assert is_list(op_schema["required"])
+        assert op_schema["additionalProperties"] == false
+      end)
+    end
+
+    test "literal operation schema is correct" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      literal_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "literal"
+        end)
+
+      assert literal_schema != nil
+      assert "op" in literal_schema["required"]
+      assert "value" in literal_schema["required"]
+      assert literal_schema["properties"]["value"] == %{}
+    end
+
+    test "operations with no fields have only 'op' as required" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      count_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "count"
+        end)
+
+      assert count_schema["required"] == ["op"]
+    end
+
+    test "operations with optional fields do not include them in required" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      get_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "get"
+        end)
+
+      assert "path" in get_schema["required"]
+      assert "default" not in get_schema["required"]
+      assert Map.has_key?(get_schema["properties"], "default")
+    end
+
+    test "expr types use inline object schema (not $ref)" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      let_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "let"
+        end)
+
+      # Verify value and in fields use inline object schema, not $ref
+      assert let_schema["properties"]["value"]["type"] == "object"
+      assert let_schema["properties"]["value"]["properties"]["op"]["type"] == "string"
+      assert let_schema["properties"]["in"]["type"] == "object"
+      assert let_schema["properties"]["in"]["properties"]["op"]["type"] == "string"
+    end
+
+    test "list of expr types have correct inline schema" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      and_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "and"
+        end)
+
+      assert and_schema["properties"]["conditions"]["type"] == "array"
+      assert is_map(and_schema["properties"]["conditions"]["items"])
+      assert and_schema["properties"]["conditions"]["items"]["type"] == "object"
+
+      assert and_schema["properties"]["conditions"]["items"]["properties"]["op"]["type"] ==
+               "string"
+    end
+
+    test "list of string types have correct schema" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      select_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "select"
+        end)
+
+      assert select_schema["properties"]["fields"] == %{
+               "type" => "array",
+               "items" => %{"type" => "string"}
+             }
+    end
+
+    test "non_neg_integer types have minimum constraint" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      nth_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "nth"
+        end)
+
+      assert nth_schema["properties"]["index"]["type"] == "integer"
+      assert nth_schema["properties"]["index"]["minimum"] == 0
+    end
+
+    test "recursive operations like pipe validate correctly" do
+      schema = PtcRunner.Schema.to_llm_schema()
+
+      pipe_schema =
+        Enum.find(schema["properties"]["program"]["anyOf"], fn op ->
+          op["properties"]["op"]["const"] == "pipe"
+        end)
+
+      assert pipe_schema != nil
+      assert pipe_schema["properties"]["steps"]["type"] == "array"
+      assert pipe_schema["properties"]["steps"]["items"]["type"] == "object"
+      assert pipe_schema["properties"]["steps"]["items"]["properties"]["op"]["type"] == "string"
+    end
+
+    test "generated schema is valid JSON" do
+      schema = PtcRunner.Schema.to_llm_schema()
+      json = Jason.encode!(schema)
+      decoded = Jason.decode!(json)
+
+      assert decoded == schema
+    end
+  end
 end