@@ -12,11 +12,16 @@ const defaultFillRateThreshold = 0.1
1212// minFillRateRows is the minimum number of items required to apply fill-rate filtering
1313const minFillRateRows = 3
1414
15+ // maxFlattenDepth is the maximum nesting depth that flatten will recurse into.
16+ // Deeper nested maps are silently dropped.
17+ const maxFlattenDepth = 2
18+
1519// preservedFields is a set of keys that are exempt from all destructive strategies except whitespace normalization.
1620// Keys are matched against post-flatten map keys, so for nested fields like "user.html_url", the dotted key must be
1721// added explicitly. Empty collections are still dropped. Wins over collectionFieldExtractors.
1822var preservedFields = map [string ]bool {
1923 "html_url" : true ,
24+ "draft" : true ,
2025}
2126
2227// collectionFieldExtractors controls how array fields are handled instead of being summarized as "[N items]".
@@ -32,17 +37,24 @@ var collectionFieldExtractors = map[string][]string{
3237
3338// MarshalItems is the single entry point for response optimization.
3439// Handles two shapes: plain JSON arrays and wrapped objects with metadata.
35- func MarshalItems (data any ) ([]byte , error ) {
40+ // An optional maxDepth controls how many nesting levels flatten will recurse
41+ // into; it defaults to maxFlattenDepth when omitted.
42+ func MarshalItems (data any , maxDepth ... int ) ([]byte , error ) {
43+ depth := maxFlattenDepth
44+ if len (maxDepth ) > 0 {
45+ depth = maxDepth [0 ]
46+ }
47+
3648 raw , err := json .Marshal (data )
3749 if err != nil {
3850 return nil , fmt .Errorf ("failed to marshal data: %w" , err )
3951 }
4052
4153 switch raw [0 ] {
4254 case '[' :
43- return optimizeArray (raw )
55+ return optimizeArray (raw , depth )
4456 case '{' :
45- return optimizeObject (raw )
57+ return optimizeObject (raw , depth )
4658 default :
4759 return raw , nil
4860 }
@@ -51,13 +63,13 @@ func MarshalItems(data any) ([]byte, error) {
5163// OptimizeItems runs the full optimization pipeline on a slice of items:
5264// flatten, remove URLs, remove zero-values, normalize whitespace,
5365// summarize collections, and fill-rate filtering.
54- func OptimizeItems (items []map [string ]any ) []map [string ]any {
66+ func OptimizeItems (items []map [string ]any , depth int ) []map [string ]any {
5567 if len (items ) == 0 {
5668 return items
5769 }
5870
5971 for i , item := range items {
60- flattenedItem := flatten (item )
72+ flattenedItem := flattenTo (item , depth )
6173 items [i ] = optimizeItem (flattenedItem )
6274 }
6375
@@ -68,29 +80,26 @@ func OptimizeItems(items []map[string]any) []map[string]any {
6880 return items
6981}
7082
71- // flatten promotes values from one-level-deep nested maps into the parent
72- // using dot-notation keys ("user.login", "user.id"). Nested maps and arrays
73- // within those nested maps are dropped.
74- func flatten (item map [string ]any ) map [string ]any {
83+ // flattenTo recursively promotes values from nested maps into the parent
84+ // using dot-notation keys ("user.login", "commit.author.date"). Arrays
85+ // within nested maps are preserved at their dotted key position.
86+ // Recursion stops at the given maxDepth; deeper nested maps are dropped.
87+ func flattenTo (item map [string ]any , maxDepth int ) map [string ]any {
7588 result := make (map [string ]any , len (item ))
76- for key , value := range item {
77- nested , ok := value .(map [string ]any )
78- if ! ok {
79- result [key ] = value
80- continue
81- }
89+ flattenInto (item , "" , result , 1 , maxDepth )
90+ return result
91+ }
8292
83- for nk , nv := range nested {
84- switch nv .(type ) {
85- case map [string ]any , []any :
86- // skip complex nested values
87- default :
88- result [key + "." + nk ] = nv
89- }
93+ // flattenInto is the recursive worker for flattenTo.
94+ func flattenInto (item map [string ]any , prefix string , result map [string ]any , depth int , maxDepth int ) {
95+ for key , value := range item {
96+ fullKey := prefix + key
97+ if nested , ok := value .(map [string ]any ); ok && depth < maxDepth {
98+ flattenInto (nested , fullKey + "." , result , depth + 1 , maxDepth )
99+ } else if ! ok {
100+ result [fullKey ] = value
90101 }
91102 }
92-
93- return result
94103}
95104
96105// filterByFillRate drops keys that appear on less than the threshold proportion of items.
@@ -126,16 +135,16 @@ func filterByFillRate(items []map[string]any, threshold float64) []map[string]an
126135}
127136
128137// optimizeArray is the entry point for optimizing a raw JSON array.
129- func optimizeArray (raw []byte ) ([]byte , error ) {
138+ func optimizeArray (raw []byte , depth int ) ([]byte , error ) {
130139 var items []map [string ]any
131140 if err := json .Unmarshal (raw , & items ); err != nil {
132141 return nil , fmt .Errorf ("failed to unmarshal JSON: %w" , err )
133142 }
134- return json .Marshal (OptimizeItems (items ))
143+ return json .Marshal (OptimizeItems (items , depth ))
135144}
136145
137146// optimizeObject is the entry point for optimizing a raw JSON object.
138- func optimizeObject (raw []byte ) ([]byte , error ) {
147+ func optimizeObject (raw []byte , depth int ) ([]byte , error ) {
139148 var wrapper map [string ]any
140149 if err := json .Unmarshal (raw , & wrapper ); err != nil {
141150 return nil , fmt .Errorf ("failed to unmarshal JSON: %w" , err )
@@ -161,7 +170,7 @@ func optimizeObject(raw []byte) ([]byte, error) {
161170 items = append (items , m )
162171 }
163172 }
164- wrapper [dataKey ] = OptimizeItems (items )
173+ wrapper [dataKey ] = OptimizeItems (items , depth )
165174
166175 return json .Marshal (wrapper )
167176}
0 commit comments