380 lines
11 KiB
Go
380 lines
11 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"regexp"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/scriptorium"
|
|
)
|
|
|
|
const scriptoriumProviderName = "scriptorium"
|
|
|
|
type ScriptoriumClientConfig struct {
|
|
ProfileDir string
|
|
ProfileFile string
|
|
Assets *AssetRegistry
|
|
Timeout time.Duration
|
|
HTTPClient *http.Client
|
|
EngineOptions []scriptorium.Option
|
|
Recorder *LLMProfileRecorder
|
|
}
|
|
|
|
type ScriptoriumClient struct {
|
|
engine *scriptorium.Engine
|
|
recorder *LLMProfileRecorder
|
|
}
|
|
|
|
type LLMProfileRecorder struct {
|
|
mu sync.Mutex
|
|
profiles map[string]artifacts.LLMProfileManifest
|
|
}
|
|
|
|
var _ contracts.StructuredLLMClient = (*ScriptoriumClient)(nil)
|
|
var _ contracts.LLMProfileManifestProvider = (*ScriptoriumClient)(nil)
|
|
|
|
func NewScriptoriumClient(cfg ScriptoriumClientConfig) (*ScriptoriumClient, error) {
|
|
if cfg.Assets == nil {
|
|
return nil, fmt.Errorf("scriptorium client assets must not be nil")
|
|
}
|
|
if strings.TrimSpace(cfg.ProfileDir) != "" && strings.TrimSpace(cfg.ProfileFile) != "" {
|
|
return nil, fmt.Errorf("scriptorium profile_dir and profile_file are mutually exclusive")
|
|
}
|
|
options, err := cfg.Assets.ScriptoriumOptions()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if profileFile := strings.TrimSpace(cfg.ProfileFile); profileFile != "" {
|
|
options = append(options, scriptorium.WithProfileFile(profileFile))
|
|
}
|
|
options = append(options, cfg.EngineOptions...)
|
|
|
|
engine, err := scriptorium.NewEngine(scriptorium.Config{
|
|
ProfileDir: strings.TrimSpace(cfg.ProfileDir),
|
|
Timeout: cfg.Timeout,
|
|
HTTPClient: cfg.HTTPClient,
|
|
}, options...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create Scriptorium engine: %w", err)
|
|
}
|
|
recorder := cfg.Recorder
|
|
if recorder == nil {
|
|
recorder = NewLLMProfileRecorder()
|
|
}
|
|
return &ScriptoriumClient{
|
|
engine: engine,
|
|
recorder: recorder,
|
|
}, nil
|
|
}
|
|
|
|
func (c *ScriptoriumClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
if c == nil {
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("scriptorium client must not be nil")
|
|
}
|
|
if c.engine == nil {
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("scriptorium client engine must not be nil")
|
|
}
|
|
if err := validateOutputTarget(out); err != nil {
|
|
return contracts.StructuredCompletionResponse{}, err
|
|
}
|
|
promptID := strings.TrimSpace(req.PromptID)
|
|
if promptID == "" {
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion prompt_id must not be empty")
|
|
}
|
|
|
|
runReq := scriptorium.RunRequest{
|
|
PromptID: promptID,
|
|
PromptVersion: strings.TrimSpace(req.PromptVersion),
|
|
ProfileID: strings.TrimSpace(req.ProfileID),
|
|
Inputs: scriptoriumInputs(req.Inputs),
|
|
Vars: scriptoriumVars(req),
|
|
Metadata: scriptoriumMetadata(req),
|
|
}
|
|
prepared, err := c.engine.Prepare(ctx, runReq)
|
|
if err != nil {
|
|
if ctxErr := ctx.Err(); ctxErr != nil {
|
|
return contracts.StructuredCompletionResponse{}, ctxErr
|
|
}
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("prepare Scriptorium prompt %q: %w", promptID, redactScriptoriumError(err))
|
|
}
|
|
result, err := c.engine.Run(ctx, runReq)
|
|
if err != nil {
|
|
if ctxErr := ctx.Err(); ctxErr != nil {
|
|
return contracts.StructuredCompletionResponse{}, ctxErr
|
|
}
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("run Scriptorium prompt %q: %w", promptID, redactScriptoriumError(err))
|
|
}
|
|
if result == nil {
|
|
return contracts.StructuredCompletionResponse{}, fmt.Errorf("run Scriptorium prompt %q: %w: empty result", promptID, contracts.ErrInvalidStructuredOutput)
|
|
}
|
|
response := c.responseFromResult(result, prepared)
|
|
if result.Validation.Status == scriptorium.ValidationFailed || !result.Validation.IsValid {
|
|
return response, fmt.Errorf("run Scriptorium prompt %q: %w: validation failed: %s", promptID, contracts.ErrInvalidStructuredOutput, strings.Join(result.Validation.Errors, "; "))
|
|
}
|
|
if len(strings.TrimSpace(string(response.Content))) == 0 {
|
|
return response, fmt.Errorf("run Scriptorium prompt %q: %w: empty structured output", promptID, contracts.ErrInvalidStructuredOutput)
|
|
}
|
|
if err := json.Unmarshal(response.Content, out); err != nil {
|
|
return response, fmt.Errorf("decode Scriptorium structured output for prompt %q: %w: %w", promptID, contracts.ErrInvalidStructuredOutput, err)
|
|
}
|
|
return response, nil
|
|
}
|
|
|
|
func (c *ScriptoriumClient) responseFromResult(result *scriptorium.RunResult, prepared *scriptorium.PreparedRun) contracts.StructuredCompletionResponse {
|
|
content := result.Artifact.Body
|
|
if len(content) == 0 {
|
|
content = []byte(result.RawOutput)
|
|
}
|
|
profile := artifacts.LLMProfileManifest{
|
|
ID: strings.TrimSpace(result.SelectedProfileID),
|
|
Provider: scriptoriumProviderName,
|
|
Model: firstNonEmpty(result.ModelName, result.EffectiveModelParams.Model),
|
|
}
|
|
if c.recorder != nil {
|
|
c.recorder.Record(profile)
|
|
}
|
|
return contracts.StructuredCompletionResponse{
|
|
Content: append(json.RawMessage(nil), content...),
|
|
Provider: profile.Provider,
|
|
Model: profile.Model,
|
|
ProfileID: profile.ID,
|
|
PromptTokens: result.Usage.PromptTokens,
|
|
CompletionTokens: result.Usage.CompletionTokens,
|
|
TotalTokens: result.Usage.TotalTokens,
|
|
Debug: scriptoriumDebugMaterial(prepared, result),
|
|
}
|
|
}
|
|
|
|
func scriptoriumDebugMaterial(prepared *scriptorium.PreparedRun, result *scriptorium.RunResult) *contracts.LLMDebugMaterial {
|
|
material := &contracts.LLMDebugMaterial{}
|
|
if prepared != nil {
|
|
material.Prompt = scriptoriumDebugPrompt(prepared)
|
|
}
|
|
if result != nil {
|
|
material.Response = scriptoriumDebugResponse(result)
|
|
}
|
|
if material.Prompt == nil && material.Response == nil {
|
|
return nil
|
|
}
|
|
return material
|
|
}
|
|
|
|
func scriptoriumDebugPrompt(prepared *scriptorium.PreparedRun) *contracts.LLMDebugPrompt {
|
|
if prepared == nil {
|
|
return nil
|
|
}
|
|
messages := make([]contracts.LLMDebugMessage, 0, len(prepared.Messages))
|
|
for _, message := range prepared.Messages {
|
|
messages = append(messages, contracts.LLMDebugMessage{
|
|
Role: message.Role,
|
|
Content: message.Content,
|
|
CacheControl: jsonObject(message.CacheControl),
|
|
})
|
|
}
|
|
return &contracts.LLMDebugPrompt{
|
|
PromptID: prepared.PromptID,
|
|
PromptVersion: prepared.PromptVersion,
|
|
PromptHash: prepared.PromptHash,
|
|
SelectedProfileID: prepared.SelectedProfileID,
|
|
SessionID: prepared.SessionID,
|
|
RenderedPromptHash: prepared.RenderedPromptHash,
|
|
Messages: messages,
|
|
EffectiveModelParams: jsonObject(prepared.EffectiveModelParams),
|
|
OutputContract: jsonObject(prepared.OutputContract),
|
|
StructuredOutput: jsonObject(prepared.StructuredOutput),
|
|
InputHashes: cloneStringMap(prepared.InputHashes),
|
|
}
|
|
}
|
|
|
|
func scriptoriumDebugResponse(result *scriptorium.RunResult) *contracts.LLMDebugResponse {
|
|
if result == nil {
|
|
return nil
|
|
}
|
|
content := result.RawOutput
|
|
if content == "" {
|
|
content = string(result.Artifact.Body)
|
|
}
|
|
return &contracts.LLMDebugResponse{
|
|
Content: content,
|
|
RunID: result.RunID,
|
|
PromptID: result.PromptID,
|
|
PromptVersion: result.PromptVersion,
|
|
PromptHash: result.PromptHash,
|
|
RenderedPromptHash: result.RenderedPromptHash,
|
|
SelectedProfileID: result.SelectedProfileID,
|
|
ModelName: result.ModelName,
|
|
Endpoint: result.Endpoint,
|
|
EffectiveModelParams: jsonObject(result.EffectiveModelParams),
|
|
InputHashes: cloneStringMap(result.InputHashes),
|
|
Validation: jsonObject(result.Validation),
|
|
Usage: contracts.LLMDebugUsage{
|
|
PromptTokens: result.Usage.PromptTokens,
|
|
CompletionTokens: result.Usage.CompletionTokens,
|
|
TotalTokens: result.Usage.TotalTokens,
|
|
CachedTokens: result.Usage.CachedTokens,
|
|
CacheWriteTokens: result.Usage.CacheWriteTokens,
|
|
},
|
|
}
|
|
}
|
|
|
|
func jsonObject(value any) map[string]any {
|
|
if value == nil {
|
|
return nil
|
|
}
|
|
data, err := json.Marshal(value)
|
|
if err != nil || string(data) == "null" {
|
|
return nil
|
|
}
|
|
var out map[string]any
|
|
if err := json.Unmarshal(data, &out); err != nil {
|
|
return nil
|
|
}
|
|
if len(out) == 0 {
|
|
return nil
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneStringMap(values map[string]string) map[string]string {
|
|
if len(values) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]string, len(values))
|
|
for key, value := range values {
|
|
out[key] = value
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (c *ScriptoriumClient) LLMProfileManifests() []artifacts.LLMProfileManifest {
|
|
if c == nil || c.recorder == nil {
|
|
return nil
|
|
}
|
|
return c.recorder.Manifests()
|
|
}
|
|
|
|
func NewLLMProfileRecorder() *LLMProfileRecorder {
|
|
return &LLMProfileRecorder{profiles: map[string]artifacts.LLMProfileManifest{}}
|
|
}
|
|
|
|
func (r *LLMProfileRecorder) Record(profile artifacts.LLMProfileManifest) {
|
|
if r == nil {
|
|
return
|
|
}
|
|
profile.ID = strings.TrimSpace(profile.ID)
|
|
profile.Provider = strings.TrimSpace(profile.Provider)
|
|
profile.Model = strings.TrimSpace(profile.Model)
|
|
key := profile.ID + "\x00" + profile.Provider + "\x00" + profile.Model
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if r.profiles == nil {
|
|
r.profiles = map[string]artifacts.LLMProfileManifest{}
|
|
}
|
|
r.profiles[key] = profile
|
|
}
|
|
|
|
func (r *LLMProfileRecorder) Manifests() []artifacts.LLMProfileManifest {
|
|
if r == nil {
|
|
return nil
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if len(r.profiles) == 0 {
|
|
return nil
|
|
}
|
|
keys := make([]string, 0, len(r.profiles))
|
|
for key := range r.profiles {
|
|
keys = append(keys, key)
|
|
}
|
|
sort.Strings(keys)
|
|
out := make([]artifacts.LLMProfileManifest, 0, len(keys))
|
|
for _, key := range keys {
|
|
out = append(out, r.profiles[key])
|
|
}
|
|
return out
|
|
}
|
|
|
|
func scriptoriumInputs(inputs contracts.LLMInputSet) map[string]scriptorium.ArtifactRef {
|
|
if len(inputs) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]scriptorium.ArtifactRef, len(inputs))
|
|
for key, material := range inputs {
|
|
name := strings.TrimSpace(key)
|
|
if name == "" {
|
|
name = strings.TrimSpace(material.Name)
|
|
}
|
|
if name == "" {
|
|
continue
|
|
}
|
|
body := string(material.Content)
|
|
if body == "" {
|
|
body = " "
|
|
}
|
|
if origin := strings.TrimSpace(material.OriginURI); origin != "" {
|
|
out[name] = scriptorium.InlineWithURI(origin, body)
|
|
} else {
|
|
out[name] = scriptorium.Inline(body)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func scriptoriumVars(req contracts.StructuredCompletionRequest) map[string]string {
|
|
vars := make(map[string]string, len(req.Vars)+1)
|
|
for key, value := range req.Vars {
|
|
name := strings.TrimSpace(key)
|
|
if name == "" || value == nil {
|
|
continue
|
|
}
|
|
vars[name] = fmt.Sprint(value)
|
|
}
|
|
if sessionID := strings.TrimSpace(req.SessionID); sessionID != "" {
|
|
vars["session_id"] = sessionID
|
|
}
|
|
if len(vars) == 0 {
|
|
return nil
|
|
}
|
|
return vars
|
|
}
|
|
|
|
func scriptoriumMetadata(req contracts.StructuredCompletionRequest) map[string]string {
|
|
metadata := map[string]string{}
|
|
if stageName := strings.TrimSpace(req.StageName); stageName != "" {
|
|
metadata["stage_name"] = stageName
|
|
}
|
|
if len(metadata) == 0 {
|
|
return nil
|
|
}
|
|
return metadata
|
|
}
|
|
|
|
var bearerTokenPattern = regexp.MustCompile(`(?i)Bearer\s+[A-Za-z0-9._~+/=-]+`)
|
|
|
|
func redactScriptoriumError(err error) error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
return redactedProviderError{err: err}
|
|
}
|
|
|
|
type redactedProviderError struct {
|
|
err error
|
|
}
|
|
|
|
func (e redactedProviderError) Error() string {
|
|
return bearerTokenPattern.ReplaceAllString(e.err.Error(), "Bearer "+secretReplacement)
|
|
}
|
|
|
|
func (e redactedProviderError) Unwrap() error {
|
|
return e.err
|
|
}
|