llmgw/internal/guard/doc.go

package guard

import (
	"encoding/json"
)

// doc is a request body opened just far enough to read and rewrite its text:
// chat messages (string content, or text parts), a completion's prompt, an
// embedding's input, and the tools offered. Everything else stays as the
// client sent it.
type doc struct {
	top  map[string]json.RawMessage
	msgs []map[string]json.RawMessage
}

// segment is one piece of text in a request, and how to put it back changed.
type segment struct {
	role string
	text string
	set  func(string)
}

func parseDoc(body []byte) (*doc, error) {
	d := &doc{}
	if err := json.Unmarshal(body, &d.top); err != nil {
		return nil, err
	}
	if raw, ok := d.top["messages"]; ok {
		if err := json.Unmarshal(raw, &d.msgs); err != nil {
			return nil, err
		}
	}
	return d, nil
}

func (d *doc) encode() []byte {
	if d.msgs != nil {
		d.top["messages"], _ = json.Marshal(d.msgs)
	}
	b, _ := json.Marshal(d.top)
	return b
}

func str(raw json.RawMessage) (string, bool) {
	var s string
	if json.Unmarshal(raw, &s) != nil {
		return "", false
	}
	return s, true
}

func rawString(s string) json.RawMessage {
	b, _ := json.Marshal(s)
	return b
}

// segments are every piece of text the request carries, in order.
func (d *doc) segments() []segment {
	var out []segment
	for i := range d.msgs {
		m := d.msgs[i]
		role, _ := str(m["role"])
		content := m["content"]
		if s, ok := str(content); ok {
			out = append(out, segment{role: role, text: s, set: func(v string) { m["content"] = rawString(v) }})
			continue
		}
		var parts []map[string]json.RawMessage
		if json.Unmarshal(content, &parts) == nil {
			for j := range parts {
				p := parts[j]
				if t, ok := str(p["text"]); ok {
					out = append(out, segment{role: role, text: t, set: func(v string) {
						p["text"] = rawString(v)
						m["content"], _ = json.Marshal(parts)
					}})
				}
			}
		}
	}
	for _, key := range []string{"prompt", "input"} {
		raw, ok := d.top[key]
		if !ok {
			continue
		}
		if s, ok := str(raw); ok {
			out = append(out, segment{role: "user", text: s, set: func(v string) { d.top[key] = rawString(v) }})
			continue
		}
		var list []json.RawMessage
		if json.Unmarshal(raw, &list) == nil {
			for j := range list {
				if s, ok := str(list[j]); ok {
					out = append(out, segment{role: "user", text: s, set: func(v string) {
						list[j] = rawString(v)
						d.top[key], _ = json.Marshal(list)
					}})
				}
			}
		}
	}
	return out
}

// tool is one tool a request offers.
type tool struct {
	name       string
	parameters json.RawMessage
}

func (d *doc) tools() []tool {
	var list []struct {
		Function struct {
			Name       string          `json:"name"`
			Parameters json.RawMessage `json:"parameters"`
		} `json:"function"`
	}
	if raw, ok := d.top["tools"]; ok {
		_ = json.Unmarshal(raw, &list)
	}
	out := make([]tool, len(list))
	for i, t := range list {
		out[i] = tool{name: t.Function.Name, parameters: t.Function.Parameters}
	}
	return out
}

// keepTools drops the offered tools keep says no to.
func (d *doc) keepTools(keep func(name string) bool) {
	raw, ok := d.top["tools"]
	if !ok {
		return
	}
	var list []json.RawMessage
	if json.Unmarshal(raw, &list) != nil {
		return
	}
	var kept []json.RawMessage
	for _, t := range list {
		var one struct {
			Function struct {
				Name string `json:"name"`
			} `json:"function"`
		}
		_ = json.Unmarshal(t, &one)
		if keep(one.Function.Name) {
			kept = append(kept, t)
		}
	}
	if len(kept) == 0 {
		delete(d.top, "tools")
		delete(d.top, "tool_choice")
		return
	}
	d.top["tools"], _ = json.Marshal(kept)
}

// pin puts a system message first, and with drop removes the caller's own.
func (d *doc) pin(prompt string, drop bool) bool {
	if d.msgs == nil || prompt == "" && !drop {
		return false
	}
	var msgs []map[string]json.RawMessage
	if prompt != "" {
		msgs = append(msgs, map[string]json.RawMessage{"role": rawString("system"), "content": rawString(prompt)})
	}
	for _, m := range d.msgs {
		role, _ := str(m["role"])
		if drop && (role == "system" || role == "developer") {
			continue
		}
		msgs = append(msgs, m)
	}
	d.msgs = msgs
	return true
}

// capOutput puts a limit on the answer where the request has none, and
// lowers one above it.
func (d *doc) capOutput(limit int) bool {
	if limit <= 0 {
		return false
	}
	changed := false
	found := false
	for _, k := range []string{"max_tokens", "max_completion_tokens"} {
		raw, ok := d.top[k]
		if !ok {
			continue
		}
		var n int
		if json.Unmarshal(raw, &n) == nil && n > 0 {
			found = true
			if n > limit {
				d.top[k], _ = json.Marshal(limit)
				changed = true
			}
		}
	}
	if !found {
		d.top["max_tokens"], _ = json.Marshal(limit)
		changed = true
	}
	return changed
}

// hasImages reports whether any message carries an image part.
func (d *doc) hasImages() bool {
	for _, m := range d.msgs {
		var parts []struct {
			Type string `json:"type"`
		}
		if json.Unmarshal(m["content"], &parts) == nil {
			for _, p := range parts {
				if p.Type == "image_url" || p.Type == "input_image" || p.Type == "image" {
					return true
				}
			}
		}
	}
	return false
}

// responseSchema is the JSON schema the request asks the answer to fit, if
// it asks for one.
func (d *doc) responseSchema() json.RawMessage {
	var rf struct {
		Type       string `json:"type"`
		JSONSchema struct {
			Schema json.RawMessage `json:"schema"`
		} `json:"json_schema"`
	}
	if raw, ok := d.top["response_format"]; !ok || json.Unmarshal(raw, &rf) != nil || rf.Type != "json_schema" {
		return nil
	}
	return rf.JSONSchema.Schema
}