127 lines
3.5 KiB
Go
127 lines
3.5 KiB
Go
package llm
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"regexp"
|
|
"strings"
|
|
)
|
|
|
|
const redactedSecret = "[REDACTED]"
|
|
|
|
var stageSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_.-]+`)
|
|
|
|
// InteractionArtifacts contains file paths for written prompt/response
|
|
// diagnostics.
|
|
type InteractionArtifacts struct {
|
|
RequestMetadataPath string
|
|
RequestPayloadPath string
|
|
ResponsePayloadPath string
|
|
ErrorPayloadPath string
|
|
}
|
|
|
|
// DiagnosticsWriter writes generic machine-readable LLM interaction artifacts.
|
|
type DiagnosticsWriter struct {
|
|
dir string
|
|
secrets []string
|
|
}
|
|
|
|
// NewDiagnosticsWriter creates a diagnostics writer rooted at dir.
|
|
func NewDiagnosticsWriter(dir string, secrets []string) *DiagnosticsWriter {
|
|
filtered := make([]string, 0, len(secrets))
|
|
for _, secret := range secrets {
|
|
secret = strings.TrimSpace(secret)
|
|
if secret != "" {
|
|
filtered = append(filtered, secret)
|
|
}
|
|
}
|
|
|
|
return &DiagnosticsWriter{
|
|
dir: dir,
|
|
secrets: filtered,
|
|
}
|
|
}
|
|
|
|
// WriteInteraction writes request/response/error artifacts for a generic LLM
|
|
// interaction stage.
|
|
func (w *DiagnosticsWriter) WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (InteractionArtifacts, error) {
|
|
if strings.TrimSpace(w.dir) == "" {
|
|
return InteractionArtifacts{}, fmt.Errorf("diagnostics directory must not be empty")
|
|
}
|
|
if err := os.MkdirAll(w.dir, 0o755); err != nil {
|
|
return InteractionArtifacts{}, fmt.Errorf("create diagnostics directory: %w", err)
|
|
}
|
|
|
|
base := sanitizeStageName(stage)
|
|
paths := InteractionArtifacts{
|
|
RequestMetadataPath: filepath.Join(w.dir, base+"-request-metadata.json"),
|
|
RequestPayloadPath: filepath.Join(w.dir, base+"-request-payload.json"),
|
|
ResponsePayloadPath: filepath.Join(w.dir, base+"-response-payload.json"),
|
|
}
|
|
|
|
if err := w.writeJSON(paths.RequestMetadataPath, requestMetadata); err != nil {
|
|
return InteractionArtifacts{}, err
|
|
}
|
|
if err := w.writeJSON(paths.RequestPayloadPath, requestPayload); err != nil {
|
|
return InteractionArtifacts{}, err
|
|
}
|
|
if err := w.writeJSON(paths.ResponsePayloadPath, responsePayload); err != nil {
|
|
return InteractionArtifacts{}, err
|
|
}
|
|
if errorPayload != nil {
|
|
paths.ErrorPayloadPath = filepath.Join(w.dir, base+"-error-payload.json")
|
|
if err := w.writeJSON(paths.ErrorPayloadPath, errorPayload); err != nil {
|
|
return InteractionArtifacts{}, err
|
|
}
|
|
}
|
|
|
|
return paths, nil
|
|
}
|
|
|
|
func sanitizeStageName(stage string) string {
|
|
stage = strings.TrimSpace(stage)
|
|
if stage == "" {
|
|
return "llm"
|
|
}
|
|
safe := stageSanitizer.ReplaceAllString(stage, "_")
|
|
safe = strings.Trim(safe, "._-")
|
|
if safe == "" {
|
|
return "llm"
|
|
}
|
|
return safe
|
|
}
|
|
|
|
func (w *DiagnosticsWriter) writeJSON(path string, payload any) error {
|
|
raw, err := json.MarshalIndent(payload, "", " ")
|
|
if err != nil {
|
|
return fmt.Errorf("marshal json artifact %q: %w", filepath.Base(path), err)
|
|
}
|
|
redacted := redactJSONSecrets(raw, w.secrets)
|
|
if !json.Valid(redacted) {
|
|
return fmt.Errorf("redacted json artifact %q is invalid", filepath.Base(path))
|
|
}
|
|
redacted = append(redacted, '\n')
|
|
if err := os.WriteFile(path, redacted, 0o644); err != nil {
|
|
return fmt.Errorf("write json artifact %q: %w", filepath.Base(path), err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func redactJSONSecrets(raw []byte, secrets []string) []byte {
|
|
if len(secrets) == 0 {
|
|
return raw
|
|
}
|
|
|
|
result := string(raw)
|
|
for _, secret := range secrets {
|
|
if secret == "" {
|
|
continue
|
|
}
|
|
result = strings.ReplaceAll(result, "Bearer "+secret, "Bearer "+redactedSecret)
|
|
result = strings.ReplaceAll(result, secret, redactedSecret)
|
|
}
|
|
return []byte(result)
|
|
}
|