feat(compression): support grouped search output

Author: masterkain

lib/codex_pooler/gateway/request_compression/content_detector.ex

@@ -35,6 +35,12 @@ defmodule CodexPooler.Gateway.RequestCompression.ContentDetector do
   @search_heading_regex ~r/\b(?:search results|results for|matches for|found\s+\d+\s+(?:results|matches))\b/i
   @numbered_result_regex ~r/^\s*(?:\d+[\).]\s+|[-*]\s+).+\s-\s.+$/m
   @path_match_regex ~r/^\s*[\w.\/-]+:\d+(?::\d+)?:\s*\S/m
+  @nul_path_match_regex ~r/^[^\x00\r\n]+\x00\d+(?::\d+)?:\s*\S/m
+  @grouped_line_match_regex ~r/^\s*\d+(?::\d+)?:\s*\S/
+  @grouped_heading_path_regex ~r/^[\w.\/-]+$/u
+  @grouped_heading_extension_regex ~r/(?:^|\/)[\w.-]+\.[A-Za-z0-9][A-Za-z0-9_-]*$/u
+  @grouped_heading_sentence_punctuation_regex ~r/[!?;]|\.\s*$/
+  @max_grouped_heading_bytes 240
   @url_regex ~r/https?:\/\/\S+/i
   @snippet_separator_regex ~r/\s-\s/
 
@@ -167,14 +173,18 @@ defmodule CodexPooler.Gateway.RequestCompression.ContentDetector do
   defp search_points(content) do
     numbered_results = scan_count(@numbered_result_regex, content)
     path_matches = scan_count(@path_match_regex, content)
+    nul_matches = scan_count(@nul_path_match_regex, content)
+    grouped_matches = grouped_search_match_count(content)
+    structural_matches = path_matches + nul_matches + grouped_matches
     urls = scan_count(@url_regex, content)
     separators = scan_count(@snippet_separator_regex, content)
 
     [
       score(regex_match?(@search_heading_regex, content), 30),
       cond_score(numbered_results, [{2, 30}, {1, 15}]),
-      cond_score(path_matches, [{3, 45}, {2, 35}, {1, 15}]),
-      score(path_matches >= 2, 15),
+      cond_score(structural_matches, [{3, 45}, {2, 35}, {1, 15}]),
+      score(structural_matches >= 2, 15),
+      score(nul_matches >= 2 or grouped_matches >= 2, 10),
       cond_score(urls, [{2, 20}, {1, 10}]),
       score(separators >= 2, 10)
     ]
@@ -223,6 +233,40 @@ defmodule CodexPooler.Gateway.RequestCompression.ContentDetector do
     |> min(100)
   end
 
+  defp grouped_search_match_count(content) do
+    content
+    |> lines()
+    |> Enum.with_index()
+    |> Enum.map(fn {line, index} ->
+      path = String.trim(line)
+
+      if grouped_heading?(path) do
+        count_grouped_lines_after(content, index)
+      else
+        0
+      end
+    end)
+    |> Enum.sum()
+  end
+
+  defp count_grouped_lines_after(content, heading_index) do
+    matches =
+      content
+      |> lines()
+      |> Enum.drop(heading_index + 1)
+      |> Enum.take_while(&Regex.match?(@grouped_line_match_regex, &1))
+      |> length()
+
+    if matches >= 2, do: matches, else: 0
+  end
+
+  defp grouped_heading?(path) do
+    byte_size(path) in 1..@max_grouped_heading_bytes and
+      Regex.match?(@grouped_heading_path_regex, path) and
+      (String.contains?(path, "/") or Regex.match?(@grouped_heading_extension_regex, path)) and
+      not Regex.match?(@grouped_heading_sentence_punctuation_regex, path)
+  end
+
   defp repeated_line?(lines) do
     lines
     |> Enum.map(&String.trim/1)

lib/codex_pooler/gateway/request_compression/strategies/search_results.ex

@@ -9,8 +9,14 @@ defmodule CodexPooler.Gateway.RequestCompression.Strategies.SearchResults do
   @default_max_files 20
   @default_max_matches 60
   @default_max_matches_per_file 3
+  @default_model "gpt-4o"
+  @max_heading_bytes 240
 
   @match_line_regex ~r/^\s*(?<path>[\w.\/-][\w.\/-]*):(?<line>\d+)(?::(?<column>\d+))?:\s*(?<text>\S.*)$/u
+  @line_match_regex ~r/^\s*(?<line>\d+)(?::(?<column>\d+))?:\s*(?<text>\S.*)$/u
+  @heading_path_regex ~r/^[\w.\/-]+$/u
+  @heading_extension_regex ~r/(?:^|\/)[\w.-]+\.[A-Za-z0-9][A-Za-z0-9_-]*$/u
+  @heading_sentence_punctuation_regex ~r/[!?;]|\.\s*$/
 
   @spec compress(term(), Strategies.opts()) :: Strategies.result()
   def compress(content, opts \\ [])
@@ -30,7 +36,7 @@ defmodule CodexPooler.Gateway.RequestCompression.Strategies.SearchResults do
         compressed_lines = render_groups(groups, selected_groups, compressed_match_count)
         compressed = Strategies.join_lines(compressed_lines)
 
-        Strategies.finalize(
+        finalize(
           @strategy,
           content,
           compressed,
@@ -66,41 +72,115 @@ defmodule CodexPooler.Gateway.RequestCompression.Strategies.SearchResults do
   def compress(_content, _opts), do: :skip
 
   defp parse_matches(lines) do
-    lines
-    |> Enum.with_index()
-    |> Enum.reduce([], fn {line, index}, matches ->
-      case parse_match(line, index) do
-        {:ok, match} -> [match | matches]
-        :skip -> matches
-      end
-    end)
-    |> Enum.reverse()
+    direct_matches =
+      lines
+      |> Enum.with_index()
+      |> Enum.reduce([], fn {line, index}, matches ->
+        case parse_match(line, index) do
+          {:ok, match} -> [match | matches]
+          :skip -> matches
+        end
+      end)
+
+    grouped_matches = parse_grouped_matches(lines)
+
+    (direct_matches ++ grouped_matches)
+    |> Enum.sort_by(& &1.index)
   end
 
   defp parse_match(line, index) do
+    if String.contains?(line, <<0>>) do
+      parse_nul_match(line, index)
+    else
+      parse_classic_match(line, index)
+    end
+  end
+
+  defp parse_nul_match(line, index) do
+    case String.split(line, <<0>>) do
+      [path, fragment] ->
+        parse_line_fragment(path, fragment, index)
+
+      _malformed ->
+        :skip
+    end
+  end
+
+  defp parse_classic_match(line, index) do
     case Regex.named_captures(@match_line_regex, line) do
       %{"path" => path, "line" => line_number, "column" => column, "text" => text} ->
-        path = String.trim(path)
-        text = String.trim(text)
+        build_match(path, line_number, column, text, index)
 
-        if path == "" or text == "" do
-          :skip
-        else
-          {:ok,
-           %{
-             path: path,
-             line: line_number,
-             column: column,
-             text: text,
-             index: index
-           }}
+      _no_match ->
+        :skip
+    end
+  end
+
+  defp parse_grouped_matches(lines) do
+    lines
+    |> Enum.with_index()
+    |> Enum.flat_map(fn {line, index} ->
+      path = String.trim(line)
+
+      if path_like_heading?(path) do
+        grouped_matches_after(lines, path, index)
+      else
+        []
+      end
+    end)
+  end
+
+  defp grouped_matches_after(lines, path, heading_index) do
+    matches =
+      lines
+      |> Enum.drop(heading_index + 1)
+      |> Enum.with_index(heading_index + 1)
+      |> Enum.reduce_while([], fn {line, index}, matches ->
+        case parse_line_fragment(path, line, index) do
+          {:ok, match} -> {:cont, [match | matches]}
+          :skip -> {:halt, matches}
         end
+      end)
+      |> Enum.reverse()
+
+    if length(matches) >= 2, do: matches, else: []
+  end
+
+  defp parse_line_fragment(path, fragment, index) do
+    case Regex.named_captures(@line_match_regex, fragment) do
+      %{"line" => line_number, "column" => column, "text" => text} ->
+        build_match(path, line_number, column, text, index)
 
       _no_match ->
         :skip
     end
   end
 
+  defp build_match(path, line_number, column, text, index) do
+    path = String.trim(path)
+    text = String.trim(text)
+
+    if path == "" or text == "" or String.contains?(path, <<0>>) do
+      :skip
+    else
+      {:ok,
+       %{
+         path: path,
+         line: line_number,
+         column: column,
+         text: text,
+         index: index
+       }}
+    end
+  end
+
+  defp path_like_heading?(path) do
+    byte_size(path) in 1..@max_heading_bytes and
+      Regex.match?(@heading_path_regex, path) and
+      (String.contains?(path, "/") or Regex.match?(@heading_extension_regex, path)) and
+      not Regex.match?(@heading_sentence_punctuation_regex, path)
+  end
+
   defp group_matches(matches) do
     {paths, groups} =
       Enum.reduce(matches, {[], %{}}, fn match, {paths, groups} ->
@@ -202,4 +282,20 @@ defmodule CodexPooler.Gateway.RequestCompression.Strategies.SearchResults do
     |> Enum.map(&length(&1.matches))
     |> Enum.sum()
   end
+
+  defp finalize(strategy, original, compressed, counts, opts) do
+    Strategies.finalize(strategy, original, compressed, counts, default_model_opts(opts))
+  end
+
+  defp default_model_opts(opts) when is_list(opts) do
+    if Keyword.has_key?(opts, :model), do: opts, else: Keyword.put(opts, :model, @default_model)
+  end
+
+  defp default_model_opts(opts) when is_map(opts) do
+    if Map.has_key?(opts, :model) or Map.has_key?(opts, "model") do
+      opts
+    else
+      Map.put(opts, :model, @default_model)
+    end
+  end
 end

test/codex_pooler/gateway/request_compression/content_detector_test.exs

@@ -186,6 +186,50 @@ defmodule CodexPooler.Gateway.RequestCompression.ContentDetectorTest do
       assert_noop(:text, ContentDetector.detect(negative))
     end
 
+    test "detects classic grouped and nul-delimited search output" do
+      classic =
+        1..3
+        |> Enum.map_join("\n", &"lib/example_#{&1}.ex:#{&1}: found match #{&1}")
+
+      grouped = """
+      lib/grouped.ex
+      12: first grouped match
+      14:3: second grouped match
+      """
+
+      nul_delimited =
+        1..3
+        |> Enum.map_join("\n", &"lib/nul_result.ex\0#{&1}: nul match #{&1}")
+
+      for body <- [classic, grouped, nul_delimited] do
+        assert %{
+                 kind: :search,
+                 confidence: confidence,
+                 compressible: true,
+                 strategy: :search_results
+               } = ContentDetector.detect(body)
+
+        assert confidence >= 0.6
+      end
+    end
+
+    test "keeps prose headings and malformed nul fragments out of search detection" do
+      prose_heading = """
+      Search results from the last review.
+      1: this is prose, not a grouped file match
+      2: this is also prose
+      """
+
+      malformed_nul = """
+      lib/example.ex\0not-a-match-line
+      lib/example.ex\0: missing line number
+      lib/example.ex\0one: missing numeric line
+      """
+
+      assert_noop(:text, ContentDetector.detect(prose_heading))
+      assert_noop(:text, ContentDetector.detect(malformed_nul))
+    end
+
     test "detects build or log output at and above request-compression confidence" do
       positive = """
       example command output line

test/codex_pooler/gateway/request_compression/maybe_compress_test.exs

@@ -172,6 +172,54 @@ defmodule CodexPooler.Gateway.RequestCompression.MaybeCompressTest do
       refute inspect(metadata) =~ "call_json_document_compression"
     end
 
+    test "rewrites nul-delimited search-result tool-output strings" do
+      omitted_sentinel = "nul search omitted sentinel"
+      original_output = compression_nul_search_fixture(omitted_sentinel)
+
+      body =
+        Jason.encode!(%{
+          "model" => @supported_model,
+          "input" => [
+            %{
+              "type" => "function_call_output",
+              "call_id" => "call_nul_search_compression",
+              "output" => original_output
+            }
+          ]
+        })
+
+      {context, request_options} = request_context(body)
+
+      assert {compressed_body, compressed_options} =
+               RequestCompression.maybe_compress(body, context, request_options)
+
+      assert compressed_body != body
+
+      compressed_output =
+        compressed_body
+        |> Jason.decode!()
+        |> Map.fetch!("input")
+        |> List.first()
+        |> Map.fetch!("output")
+
+      assert compressed_output =~ "[compressed search results:"
+      assert compressed_output =~ "lib/nul_result.ex"
+      refute compressed_output =~ <<0>>
+      refute compressed_output =~ omitted_sentinel
+
+      assert %{
+               "status" => "compressed",
+               "candidate_count" => 1,
+               "compressed_count" => 1,
+               "skipped_count" => 0
+             } = metadata = compressed_options.runtime.payload_compression
+
+      assert "search_results" in metadata["strategies"]
+      assert metadata["original_tokens"] > metadata["compressed_tokens"]
+      refute inspect(metadata) =~ omitted_sentinel
+      refute inspect(metadata) =~ "call_nul_search_compression"
+    end
+
     test "skips when no route model is available" do
       body =
         Jason.encode!(%{
@@ -289,4 +337,15 @@ defmodule CodexPooler.Gateway.RequestCompression.MaybeCompressTest do
     ])
     |> Enum.join("\n")
   end
+
+  defp compression_nul_search_fixture(omitted_sentinel) do
+    1..24
+    |> Enum.map_join("\n", fn
+      11 ->
+        "lib/nul_result.ex\011: nul search result 11 #{omitted_sentinel} with extra context"
+
+      index ->
+        "lib/nul_result.ex\0#{index}: nul search result #{index} with enough context to compress"
+    end)
+  end
 end