444 lines
13 KiB
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
|
|
}
|