Assistant/cmd/cognitive-worker/model_qualification.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"slices"
	"sync"
	"time"

	"git.cin.su/m/Assistant/adapters/models"
	"git.cin.su/m/Assistant/contracts/authority"
	"git.cin.su/m/Assistant/contracts/capability"
	"git.cin.su/m/Assistant/contracts/cognitive"
	"git.cin.su/m/Assistant/contracts/identity"
	"git.cin.su/m/Assistant/platform/ids"
)

type qualificationObserver interface {
	ObserveQualification(context.Context, capability.ModelQualificationObservation) error
}

type httpQualificationObserver struct{ resolver *httpBindingResolver }

func (o httpQualificationObserver) ObserveQualification(ctx context.Context, value capability.ModelQualificationObservation) error {
	if o.resolver == nil {
		return fmt.Errorf("qualification observer resolver is nil")
	}
	return o.resolver.post(ctx, "/v1/cognitive-registry/model-qualification-observations", value, nil)
}

type modelQualificationRequest struct {
	Prompt         string                       `json:"prompt"`
	MaxOutputBytes int                          `json:"max_output_bytes"`
	ModelIDs       []capability.ProviderModelID `json:"model_ids,omitempty"`
}

type modelQualificationResult struct {
	ProviderModelID capability.ProviderModelID     `json:"provider_model_id"`
	Status          string                         `json:"status"`
	FinishReason    cognitive.FinishReason         `json:"finish_reason,omitempty"`
	LatencyMillis   int64                          `json:"latency_millis"`
	OutputObserved  bool                           `json:"output_observed"`
	QualityAssessed bool                           `json:"quality_assessed"`
	Category        capability.ModelRatingCategory `json:"category"`
	Error           string                         `json:"error,omitempty"`
}

type modelQualificationResponse struct {
	Provider             capability.ProviderInstanceID `json:"provider"`
	ObservedAtUnixMillis int64                         `json:"observed_at_unix_millis"`
	Results              []modelQualificationResult    `json:"results"`
}

func modelQualificationHandler(providers map[string]providerRegistration, owner string, observers ...qualificationObserver) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		name := r.PathValue("provider")
		registration, ok := providers[name]
		if !ok || registration.Discovery == nil || registration.Provider == nil {
			http.Error(w, "provider qualification not found", http.StatusNotFound)
			return
		}
		var request modelQualificationRequest
		if !decodeJSON(w, r, 16<<10, &request) || request.Prompt == "" {
			return
		}
		if request.MaxOutputBytes < 1 {
			request.MaxOutputBytes = 512
		}
		if request.MaxOutputBytes > 4096 {
			http.Error(w, "max_output_bytes exceeds qualification bound", http.StatusBadRequest)
			return
		}
		modelsFound, probe := registration.Discovery.Discover(r.Context())
		if probe.Health != capability.HealthAvailable || len(modelsFound) == 0 {
			http.Error(w, "provider unavailable for qualification", http.StatusServiceUnavailable)
			return
		}
		if len(request.ModelIDs) > 0 {
			modelsFound = slices.DeleteFunc(modelsFound, func(model models.DiscoveredModel) bool {
				return !slices.Contains(request.ModelIDs, model.ProviderModelID)
			})
		}
		if len(modelsFound) == 0 {
			http.Error(w, "requested model not discovered", http.StatusNotFound)
			return
		}
		if len(modelsFound) > 16 {
			modelsFound = modelsFound[:16]
		}
		results := qualifyModels(r.Context(), registration, owner, request.Prompt, request.MaxOutputBytes, modelsFound)
		if len(observers) > 0 && observers[0] != nil {
			if err := persistQualificationResults(r.Context(), observers[0], registration, results); err != nil {
				http.Error(w, "persist model qualification observation", http.StatusBadGateway)
				return
			}
		}
		w.Header().Set("Content-Type", "application/json")
		_ = jsonEncode(w, modelQualificationResponse{Provider: registration.ProviderInstanceID, ObservedAtUnixMillis: time.Now().UnixMilli(), Results: results})
	})
}

func qualifyModels(ctx context.Context, registration providerRegistration, owner, prompt string, maxOutputBytes int, discovered []models.DiscoveredModel) []modelQualificationResult {
	results := make([]modelQualificationResult, len(discovered))
	var wg sync.WaitGroup
	for index, model := range discovered {
		wg.Add(1)
		go func(index int, model models.DiscoveredModel) {
			defer wg.Done()
			started := time.Now()
			request := cognitive.AdvisoryRequest{Version: cognitive.CreatorAdvisoryV1, RequestID: cognitive.RequestID(fmt.Sprintf("qualification:%s:%d", model.ProviderModelID, started.UnixNano())), ActorID: identity.ActorID("qualification-runtime"), PrincipalID: identity.PrincipalID(owner), ScopeID: authority.ScopeID("qualification-scope"), Message: prompt, Locale: "en-US", MaxOutputBytes: maxOutputBytes, ProviderBindingID: "qualification-binding-v1", ProviderBindingRevision: 1, ProviderInstanceID: registration.ProviderInstanceID, ModelDescriptorID: capability.ModelDescriptorID("qualification-" + string(model.ProviderModelID)), ProviderModelID: model.ProviderModelID}
			response, err := registration.Provider.Advise(ctx, request)
			result := modelQualificationResult{ProviderModelID: model.ProviderModelID, LatencyMillis: time.Since(started).Milliseconds(), QualityAssessed: false, Category: capability.ModelRatingUnrated}
			if err != nil {
				result.Status = "UNAVAILABLE"
				result.FinishReason = cognitive.FinishProviderUnavailable
				result.Error = err.Error()
				results[index] = result
				return
			}
			result.FinishReason = response.FinishReason
			result.OutputObserved = response.AdvisoryText != ""
			if response.FinishReason == cognitive.FinishCompleted && result.OutputObserved {
				result.Status = "REACHABLE"
			} else {
				result.Status = "RESPONDED_WITHOUT_QUALIFICATION"
			}
			results[index] = result
		}(index, model)
	}
	wg.Wait()
	return results
}

func persistQualificationResults(ctx context.Context, observer qualificationObserver, registration providerRegistration, results []modelQualificationResult) error {
	for _, result := range results {
		observationID, err := ids.New()
		if err != nil {
			return err
		}
		status := capability.ModelQualificationRespondedWithoutRating
		switch result.Status {
		case "REACHABLE":
			status = capability.ModelQualificationReachable
		case "UNAVAILABLE":
			status = capability.ModelQualificationUnavailable
		}
		observation := capability.ModelQualificationObservation{
			ID: capability.ModelQualificationObservationID(observationID), ProviderInstanceID: registration.ProviderInstanceID, ProviderModelID: result.ProviderModelID,
			Status: status, FinishReason: string(result.FinishReason), LatencyMillis: result.LatencyMillis,
			OutputObserved: result.OutputObserved, QualityAssessed: false, Category: capability.ModelRatingUnrated,
			ObservedAtUnixMillis: time.Now().UnixMilli(),
		}
		if err := observation.Validate(); err != nil {
			return err
		}
		if err := observer.ObserveQualification(ctx, observation); err != nil {
			return err
		}
	}
	return nil
}

func jsonEncode(w http.ResponseWriter, value any) error {
	return json.NewEncoder(w).Encode(value)
}