Files
notarius/internal/framework/llm/promptkit_client.go

444 lines
13 KiB
Go

package llm
import (
"context"
"encoding/json"
"errors"
"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/promptkit"
)
const promptKitProviderName = "promptkit"
type PromptKitLocalBackendConfig struct {
Endpoint string
ConcurrencyLimit int
}
type PromptKitClientConfig struct {
ProfileDir string
ProfileFile string
LocalBackend *PromptKitLocalBackendConfig
Assets *AssetRegistry
Timeout time.Duration
HTTPClient *http.Client
EngineOptions []promptkit.Option
Recorder *LLMProfileRecorder
ReasoningEffort *string
}
type PromptKitClient struct {
engine *promptkit.Engine
recorder *LLMProfileRecorder
profileDir string
profileFile string
localEndpoint string
fallbackProfileDigest string
reasoningEffort *string
}
type LLMProfileRecorder struct {
mu sync.Mutex
profiles map[string]artifacts.LLMProfileManifest
}
var _ contracts.StructuredLLMClient = (*PromptKitClient)(nil)
var _ contracts.LLMProfileManifestProvider = (*PromptKitClient)(nil)
func PromptKitLocalBackendOption(cfg PromptKitLocalBackendConfig) promptkit.Option {
return promptkit.WithBackend(promptkit.LocalBackend(cfg.Endpoint, cfg.ConcurrencyLimit))
}
func NewPromptKitClient(cfg PromptKitClientConfig) (*PromptKitClient, error) {
if cfg.Assets == nil {
return nil, fmt.Errorf("PromptKit client assets must not be nil")
}
profileSource, profileOptions, err := promptKitProfileSourceEngineOptions(PromptKitProfileSourceConfig{
ProfileDir: cfg.ProfileDir,
ProfileFile: cfg.ProfileFile,
LocalBackend: cfg.LocalBackend,
})
if err != nil {
return nil, err
}
options, fallbackProfileDigest, err := cfg.Assets.promptKitOptions()
if err != nil {
return nil, err
}
options = append(options, profileOptions...)
options = append(options, cfg.EngineOptions...)
engine, err := promptkit.NewEngine(promptkit.Config{
ProfileDir: profileSource.ProfileDir,
Timeout: cfg.Timeout,
HTTPClient: cfg.HTTPClient,
}, options...)
if err != nil {
return nil, fmt.Errorf("create PromptKit engine: %w", err)
}
recorder := cfg.Recorder
if recorder == nil {
recorder = NewLLMProfileRecorder()
}
var reasoningEffort *string
if cfg.ReasoningEffort != nil {
value := *cfg.ReasoningEffort
reasoningEffort = &value
}
return &PromptKitClient{
engine: engine,
recorder: recorder,
profileDir: profileSource.ProfileDir,
profileFile: profileSource.ProfileFile,
localEndpoint: profileSource.localEndpoint(),
fallbackProfileDigest: fallbackProfileDigest,
reasoningEffort: reasoningEffort,
}, nil
}
func (c *PromptKitClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
if c == nil {
return contracts.StructuredCompletionResponse{}, fmt.Errorf("PromptKit client must not be nil")
}
if c.engine == nil {
return contracts.StructuredCompletionResponse{}, fmt.Errorf("PromptKit 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")
}
sessionID := strings.TrimSpace(req.SessionID)
var execution *promptkit.ExecutionTargetOverride
if c.reasoningEffort != nil {
reasoningEffort := *c.reasoningEffort
execution = &promptkit.ExecutionTargetOverride{
ReasoningEffort: &reasoningEffort,
}
}
runReq := promptkit.RunRequest{
PromptID: promptID,
PromptVersion: strings.TrimSpace(req.PromptVersion),
ProfileID: strings.TrimSpace(req.ProfileID),
SessionID: sessionID,
Inputs: promptKitInputs(req.Inputs),
Vars: promptKitVars(req, sessionID),
Execution: execution,
}
prepared, err := c.engine.PrepareExecution(ctx, runReq)
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return contracts.StructuredCompletionResponse{}, ctxErr
}
return contracts.StructuredCompletionResponse{}, fmt.Errorf("prepare PromptKit prompt %q: %w", promptID, redactPromptKitError(err))
}
defer prepared.Discard()
preparedDetails := prepared.Details()
result, err := c.engine.RunPrepared(ctx, prepared)
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return contracts.StructuredCompletionResponse{}, ctxErr
}
if errors.Is(err, promptkit.ErrCapacityExceeded) {
var capacityErr *promptkit.CapacityError
if errors.As(err, &capacityErr) && strings.TrimSpace(capacityErr.BackendID) != "" {
return contracts.StructuredCompletionResponse{}, fmt.Errorf(
"run PromptKit prompt %q on backend %q: %w: %v",
promptID,
strings.TrimSpace(capacityErr.BackendID),
contracts.ErrLLMCapacityExceeded,
redactPromptKitError(err),
)
}
return contracts.StructuredCompletionResponse{}, fmt.Errorf(
"run PromptKit prompt %q: %w: %v",
promptID,
contracts.ErrLLMCapacityExceeded,
redactPromptKitError(err),
)
}
return contracts.StructuredCompletionResponse{}, fmt.Errorf("run PromptKit prompt %q: %w", promptID, redactPromptKitError(err))
}
if result == nil {
return contracts.StructuredCompletionResponse{}, fmt.Errorf("run PromptKit prompt %q: %w: empty result", promptID, contracts.ErrInvalidStructuredOutput)
}
response := c.responseFromResult(result, &preparedDetails)
if result.Validation.Status == promptkit.ValidationFailed || !result.Validation.IsValid {
return response, fmt.Errorf("run PromptKit 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 PromptKit prompt %q: %w: empty structured output", promptID, contracts.ErrInvalidStructuredOutput)
}
if err := json.Unmarshal(response.Content, out); err != nil {
return response, fmt.Errorf("decode PromptKit structured output for prompt %q: %w: %w", promptID, contracts.ErrInvalidStructuredOutput, err)
}
return response, nil
}
func (c *PromptKitClient) responseFromResult(result *promptkit.RunResult, prepared *promptkit.PreparedRun) contracts.StructuredCompletionResponse {
content := result.Artifact.Body
if len(content) == 0 {
content = []byte(result.RawOutput)
}
profile := artifacts.LLMProfileManifest{
ID: strings.TrimSpace(result.SelectedProfileID),
Provider: promptKitProviderName,
Model: firstNonEmpty(result.ModelName, result.EffectiveModelParams.Model),
BackendID: strings.TrimSpace(result.SelectedBackendID),
ReasoningEffort: strings.TrimSpace(result.EffectiveModelParams.ReasoningEffort),
}
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: promptKitDebugMaterial(prepared, result),
}
}
func promptKitDebugMaterial(prepared *promptkit.PreparedRun, result *promptkit.RunResult) *contracts.LLMDebugMaterial {
material := &contracts.LLMDebugMaterial{}
if prepared != nil {
material.Prompt = promptKitDebugPrompt(prepared)
}
if result != nil {
material.Response = promptKitDebugResponse(result)
}
if material.Prompt == nil && material.Response == nil {
return nil
}
return material
}
func promptKitDebugPrompt(prepared *promptkit.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,
SelectedBackendID: prepared.SelectedBackendID,
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 promptKitDebugResponse(result *promptkit.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 *PromptKitClient) LLMProfileManifests() []artifacts.LLMProfileManifest {
if c == nil || c.recorder == nil {
return nil
}
return c.recorder.Manifests()
}
func (c *PromptKitClient) LLMCheckpointFingerprints() ([]CheckpointFingerprint, error) {
if c == nil {
return nil, nil
}
fingerprint, err := promptKitProfileFingerprint(c.profileDir, c.profileFile, c.fallbackProfileDigest)
if err != nil {
return nil, err
}
fingerprints := []CheckpointFingerprint{fingerprint}
if c.localEndpoint != "" {
fingerprints = append(fingerprints, promptKitLocalBackendFingerprint(c.localEndpoint))
}
return fingerprints, nil
}
func NewLLMProfileRecorder() *LLMProfileRecorder {
return &LLMProfileRecorder{profiles: map[string]artifacts.LLMProfileManifest{}}
}
func (r *LLMProfileRecorder) Record(profile artifacts.LLMProfileManifest) {
if r == nil {
return
}
profile = profile.Normalized()
key := profile.IdentityKey()
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 promptKitInputs(inputs contracts.LLMInputSet) map[string]promptkit.ArtifactRef {
if len(inputs) == 0 {
return nil
}
out := make(map[string]promptkit.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] = promptkit.InlineWithURI(origin, body)
} else {
out[name] = promptkit.Inline(body)
}
}
return out
}
func promptKitVars(req contracts.StructuredCompletionRequest, sessionID string) 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 != "" {
vars["session_id"] = sessionID
}
if len(vars) == 0 {
return nil
}
return vars
}
var bearerTokenPattern = regexp.MustCompile(`(?i)Bearer\s+[A-Za-z0-9._~+/=-]+`)
func redactPromptKitError(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
}