audita/internal/framework/llm/adapter.go

package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"sync"
	"time"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/outcomes"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/v3/internal/prompts"
	pk "go.maximumdirect.net/promptkit"
)

// CallHooks receive one logical start and finish record and the execution
// events for this call. Call records can contain transcript content.
type CallHooks struct {
	OnCall          func(structuredoutput.CallRecord)
	OnEvent         pk.ExecutionObserver
	OnEventWithCall func(uint64, structuredoutput.DiagnosticContext, pk.ExecutionEvent)
}

type RoleClient struct {
	runtime         *Runtime
	role            string
	profileID       string
	settings        config.LLM
	hooks           CallHooks
	credentialMu    sync.Mutex
	credentialNames map[string]struct{}
}

var _ structuredoutput.Client = (*RoleClient)(nil)

func (r *Runtime) NewRoleClient(role, profileID string, settings config.LLM, hooks CallHooks) (*RoleClient, error) {
	if r == nil || r.Engine == nil {
		return nil, fmt.Errorf("LLM engine is unavailable")
	}
	if profileID == "" {
		return nil, fmt.Errorf("profile selection is required")
	}
	return &RoleClient{runtime: r, role: role, profileID: profileID, settings: settings, hooks: hooks}, nil
}

func (c *RoleClient) CompleteStructured(parent context.Context, req structuredoutput.Request, out any) (response structuredoutput.Response, err error) {
	ctx, cancel := context.WithTimeout(parent, c.settings.OperationTimeout)
	defer cancel()
	var terminalMu sync.Mutex
	var terminal pk.ExecutionEvent
	terminalAvailable := false
	defer func() {
		terminalMu.Lock()
		if terminalAvailable {
			response.Execution = outcomes.ExecutionCounts{
				Available:         true,
				ProviderAttempts:  int(terminal.Attempts),
				TransientRetries:  int(terminal.TransientRetries),
				StructuralRepairs: int(terminal.RepairCalls),
				RetriesExhausted:  terminal.RetriesExhausted,
			}
			if terminal.StatusCode != 0 {
				status := terminal.StatusCode
				response.Execution.HTTPStatus = &status
			}
		}
		terminalMu.Unlock()
	}()
	var record structuredoutput.CallRecord
	if c.hooks.OnCall != nil || c.hooks.OnEventWithCall != nil {
		record.CallID = c.runtime.nextCallID.Add(1)
		record.RequestedPromptID = req.PromptID
		record.Context = copyContext(req.Context)
		record.Context.Role = c.role
	}
	if onCall := c.hooks.OnCall; onCall != nil {
		record.Started = time.Now().UTC()
		record.Prompt = structuredoutput.PromptMetadata{PromptID: req.PromptID, ProfileID: c.profileID}
		onCall(record)
		defer func() {
			record.Finished = time.Now().UTC()
			if err != nil {
				var failure *structuredoutput.Failure
				if errors.As(err, &failure) && failure.Reason == outcomes.ReasonInvalidOutputExhausted {
					record.Failure = structuredoutput.FailureMalformedOutput
				} else {
					record.Failure = structuredoutput.FailureRuntime
				}
			}
			onCall(record)
		}()
	}
	if parent.Err() != nil {
		return response, classifiedFailure(parent, ctx, parent.Err(), nil)
	}
	if err = validateOutputTarget(out); err != nil {
		return response, &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonSetupFailure, Message: "invalid structured output target", Cause: err}
	}
	contract, contractErr := OutputContract(req.Schema, c.settings.RepairAttempts)
	if contractErr != nil {
		return response, &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonSetupFailure, Message: "unknown structured response schema", Cause: contractErr}
	}
	vars := prompts.WithSharedVariables(req.Variables)
	for name, value := range prompts.SharedVariables() {
		vars[name] = value
	}
	baseID, baseVars := req.PromptID, vars
	var suffix []pk.RenderedMessage
	if req.Replay != nil {
		if req.Replay.Exchange == nil {
			return response, &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonSetupFailure, Message: "repair exchange is unavailable"}
		}
		instruction, renderErr := c.runtime.Engine.Prepare(ctx, pk.RunRequest{PromptID: req.PromptID, ProfileID: c.profileID, Vars: vars})
		if renderErr != nil {
			return response, classifiedFailure(parent, ctx, renderErr, nil)
		}
		response.RepairInstruction = convertMessages(instruction.Messages)
		exchange := req.Replay.Exchange
		baseID, baseVars = exchange.PromptID, exchange.Variables
		suffix = append(suffix, pk.RenderedMessage{Role: pk.RoleAssistant, Content: exchange.FinalContent})
		for _, message := range req.Replay.Turns {
			native := pk.RenderedMessage{Role: message.Role, Content: message.Content}
			if cache := message.CacheControl; cache != nil {
				native.CacheControl = &pk.CacheControl{Type: pk.CacheControlType(cache.Type), TTL: cache.TTL}
			}
			suffix = append(suffix, native)
		}
		suffix = append(suffix, instruction.Messages...)
	}
	handle, prepErr := c.runtime.Engine.PrepareExecution(ctx, pk.RunRequest{
		PromptID: baseID, ProfileID: c.profileID, Vars: baseVars, AppendedMessages: suffix,
		TransientRetries: &c.settings.TransientRetries,
		Validation:       contract,
		Observer: func(event pk.ExecutionEvent) {
			if event.Kind == pk.ExecutionFinished {
				terminalMu.Lock()
				terminal = event
				terminalAvailable = true
				terminalMu.Unlock()
			}
			if c.hooks.OnEvent != nil {
				c.hooks.OnEvent(event)
			}
			if c.hooks.OnEventWithCall != nil {
				c.hooks.OnEventWithCall(record.CallID, record.Context, event)
			}
		},
	})
	if prepErr != nil {
		return response, classifiedFailure(parent, ctx, prepErr, nil)
	}
	defer handle.Discard()
	details := handle.Details()
	if req.Replay != nil {
		prefix := req.Replay.Exchange.Messages
		if len(details.Messages) < len(prefix) || !reflect.DeepEqual(convertMessages(details.Messages[:len(prefix)]), prefix) {
			return response, &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonInternalFailure, Message: "repair conversation prefix changed"}
		}
	}
	if name := details.EffectiveModelParams.APIKeyEnv; name != "" {
		c.credentialMu.Lock()
		if c.credentialNames == nil {
			c.credentialNames = make(map[string]struct{})
		}
		c.credentialNames[name] = struct{}{}
		c.credentialMu.Unlock()
	}
	response.Prompt = promptMetadata(details)
	record.Prompt = response.Prompt
	response.Model = details.EffectiveModelParams.Model
	response.Prepared = preparedSnapshot(details)
	record.Prepared = response.Prepared
	result, runErr := c.runtime.Engine.RunPrepared(ctx, handle)
	if result != nil {
		response.FinalContent = result.RawOutput
		record.FinalContent = result.RawOutput
		response.Model = result.ModelName
		response.PromptTokens = result.Usage.PromptTokens
		response.CompletionTokens = result.Usage.CompletionTokens
		response.TotalTokens = result.Usage.TotalTokens
	}
	if runErr != nil {
		return response, classifiedFailure(parent, ctx, runErr, result)
	}
	if result == nil {
		return response, &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonInternalFailure, Message: "structured model call returned no result"}
	}
	if !result.Validation.IsValid {
		return response, classifiedFailure(parent, ctx, nil, result)
	}
	if decodeErr := json.Unmarshal([]byte(result.RawOutput), out); decodeErr != nil {
		var typeError *json.UnmarshalTypeError
		if errors.As(decodeErr, &typeError) {
			return response, &structuredoutput.Failure{Class: structuredoutput.FailureRecoverable, Reason: outcomes.ReasonInvalidOutput, Message: "structured output contains an unusable value", Cause: decodeErr}
		}
		return response, &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonInternalFailure, Message: "validated structured output could not be decoded", Cause: decodeErr}
	}
	if req.CaptureExchange && req.Replay == nil {
		frozen := make(map[string]string, len(vars))
		for key, value := range vars {
			frozen[key] = value
		}
		response.Exchange = &structuredoutput.Exchange{PromptID: req.PromptID, Variables: frozen, Messages: convertMessages(details.Messages), FinalContent: result.RawOutput}
	}
	response.Content = json.RawMessage(result.RawOutput)
	response.Provider = "openai-compatible"
	return response, nil
}

// OutputContract maps an Audita domain schema to the native request contract.
func OutputContract(schema responseschema.Key, repairs int) (*pk.OutputContract, error) {
	path, ok := responseschema.AssetPath(schema)
	if !ok {
		return nil, fmt.Errorf("schema %q", schema)
	}
	return &pk.OutputContract{Format: pk.FormatJSON, ValidationMode: pk.ValidationJSONSchema, SchemaPath: path, RepairAttempts: repairs}, nil
}

// classifiedFailure deliberately checks setup defects before broad PromptKit
// validation/generation wrappers. Provider text is kept only in the typed cause.
func classifiedFailure(parent, local context.Context, cause error, result *pk.RunResult) *structuredoutput.Failure {
	failure := &structuredoutput.Failure{Class: structuredoutput.FailureFatal, Reason: outcomes.ReasonInternalFailure, Message: "structured model call failed", Cause: cause}
	if parent.Err() != nil {
		failure.Class = structuredoutput.FailureCanceled
		failure.Reason = outcomes.ReasonCanceled
		failure.Message = "processing canceled"
		failure.Cause = parent.Err()
		return failure
	}
	if isSetupFailure(cause) {
		failure.Reason = outcomes.ReasonSetupFailure
		failure.Message = "structured model setup failed"
		return failure
	}
	if errors.Is(local.Err(), context.DeadlineExceeded) || errors.Is(cause, context.DeadlineExceeded) {
		failure.Class = structuredoutput.FailureRecoverable
		failure.Reason = outcomes.ReasonOperationTimeout
		failure.Message = "model operation timed out"
		return failure
	}
	if errors.Is(cause, pk.ErrCapacityExceeded) {
		failure.Class = structuredoutput.FailureRecoverable
		failure.Reason = outcomes.ReasonCapacityExceeded
		failure.Message = "model capacity was unavailable"
		return failure
	}
	if errors.Is(cause, pk.ErrTransport) {
		failure.Class = structuredoutput.FailureRecoverable
		failure.Reason = outcomes.ReasonTransportFailed
		failure.Message = "model transport failed"
		return failure
	}
	if errors.Is(cause, pk.ErrMalformedResponse) {
		failure.Class = structuredoutput.FailureRecoverable
		failure.Reason = outcomes.ReasonProviderResponseInvalid
		failure.Message = "model provider response was invalid"
		return failure
	}
	var provider *pk.GenerationError
	if errors.As(cause, &provider) && provider.StatusCode() != 0 {
		failure.Class = structuredoutput.FailureRecoverable
		failure.Reason = outcomes.ReasonProviderFailed
		failure.Message = "model provider failed"
		return failure
	}
	if result != nil && !result.Validation.IsValid {
		failure.Class = structuredoutput.FailureRecoverable
		failure.Kind = structuredoutput.FailureMalformedOutput
		failure.Reason = outcomes.ReasonInvalidOutputExhausted
		failure.Message = "structured output remained invalid"
		return failure
	}
	if errors.Is(cause, pk.ErrValidation) {
		failure.Reason = outcomes.ReasonSetupFailure
		failure.Message = "structured validation contract failed"
	}
	return failure
}

func isSetupFailure(err error) bool {
	return errors.Is(err, pk.ErrInvalidConfig) || errors.Is(err, pk.ErrInvalidRequest) ||
		errors.Is(err, pk.ErrPromptNotFound) || errors.Is(err, pk.ErrProfileNotFound) ||
		errors.Is(err, pk.ErrPromptLoad) || errors.Is(err, pk.ErrProfileLoad) ||
		errors.Is(err, pk.ErrArtifactLoad) || errors.Is(err, pk.ErrPromptRender)
}

func promptMetadata(details pk.PreparedRun) structuredoutput.PromptMetadata {
	return structuredoutput.PromptMetadata{
		PromptID:       details.PromptID,
		DefinitionHash: details.PromptHash, RenderedHash: details.RenderedPromptHash,
		ProfileID: details.SelectedProfileID, BackendID: details.SelectedBackendID,
		Model: details.EffectiveModelParams.Model,
	}
}

func preparedSnapshot(details pk.PreparedRun) *structuredoutput.PreparedSnapshot {
	messages := make([]structuredoutput.Message, len(details.Messages))
	for i, message := range details.Messages {
		messages[i] = structuredoutput.Message{Role: string(message.Role), Content: message.Content}
	}
	return &structuredoutput.PreparedSnapshot{Messages: messages, Prompt: promptMetadata(details)}
}

// ConfiguredSecrets supplies artifact redaction with the engine's referenced credentials.
func (c *RoleClient) ConfiguredSecrets(ctx context.Context) []string {
	c.credentialMu.Lock()
	names := make([]string, 0, len(c.credentialNames))
	for name := range c.credentialNames {
		names = append(names, name)
	}
	c.credentialMu.Unlock()
	return c.runtime.ConfiguredSecrets(ctx, names...)
}

func convertMessages(messages []pk.RenderedMessage) []structuredoutput.Message {
	result := make([]structuredoutput.Message, len(messages))
	for i, message := range messages {
		result[i] = structuredoutput.Message{Role: string(message.Role), Content: message.Content}
		if cache := message.CacheControl; cache != nil {
			result[i].CacheControl = &structuredoutput.CacheControl{Type: string(cache.Type), TTL: cache.TTL}
		}
	}
	return result
}