xinmi/internal/workbench/auxiliary_usage.go

package workbench

import (
	"context"
	"time"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/models/asr"
	"github.com/Tencent/WeKnora/internal/models/embedding"
	"github.com/Tencent/WeKnora/internal/models/rerank"
	"github.com/Tencent/WeKnora/internal/models/vlm"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
)

// Auxiliary providers do not expose billable usage through their interfaces.
// Record measured units without inventing token counts or debit transactions.
func (s *Service) recordAuxiliaryUsage(ctx context.Context, modelID, name, action string, started time.Time, units JSON, failure error) {
	tenant, ok := types.TenantIDFromContext(ctx)
	if !ok || tenant == 0 {
		return
	}
	user, _ := types.UserIDFromContext(ctx)
	a, _ := ctx.Value(meteringContextKey{}).(meteringAttribution)
	units["model_name"], units["duration_ms"], units["billing_status"] = name, time.Since(started).Milliseconds(), "provider_usage_unavailable"
	row := Usage{ID: uuid.NewString(), CallID: uuid.NewString(), TenantID: tenant, OwnerID: user, AgentID: a.AgentID, JobID: a.JobID, ModelID: modelID, Action: action, Status: "succeeded", Metadata: units, CreatedAt: time.Now()}
	if failure != nil {
		row.Status = "failed"
		category, message := jobErrorDetails(failure)
		units["error_class"], units["error_message"] = category, message
	}
	finish, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
	defer cancel()
	if err := s.settleBalance(finish, nil, row); err != nil {
		logJobFailure(finish, nil, "auxiliary_usage", err)
	}
	if failure != nil {
		_, message := jobErrorDetails(failure)
		if err := s.recordOperation(finish, OperationLog{Kind: "system", Level: "error", Action: action, Message: message, TraceID: row.CallID, Metadata: JSON{"model_id": modelID}}); err != nil {
			logJobFailure(finish, nil, "auxiliary_error", err)
		}
	}
}

func textUnits(texts []string) JSON {
	characters := 0
	for _, text := range texts {
		characters += utf8.RuneCountInString(text)
	}
	return JSON{"input_texts": len(texts), "input_characters": characters}
}

func (s *Service) WrapEmbedding(model embedding.Embedder) embedding.Embedder {
	return &observedEmbedding{Embedder: model, service: s}
}

type observedEmbedding struct {
	embedding.Embedder
	service *Service
}

func (m *observedEmbedding) Embed(ctx context.Context, text string) (result []float32, err error) {
	started := time.Now()
	defer func() {
		m.service.recordAuxiliaryUsage(ctx, m.GetModelID(), m.GetModelName(), "知识向量化", started, textUnits([]string{text}), err)
	}()
	return m.Embedder.Embed(ctx, text)
}
func (m *observedEmbedding) BatchEmbed(ctx context.Context, texts []string) (result [][]float32, err error) {
	started := time.Now()
	defer func() {
		m.service.recordAuxiliaryUsage(ctx, m.GetModelID(), m.GetModelName(), "批量知识向量化", started, textUnits(texts), err)
	}()
	return m.Embedder.BatchEmbed(ctx, texts)
}
func (m *observedEmbedding) BatchEmbedWithPool(ctx context.Context, _ embedding.Embedder, texts []string) (result [][]float32, err error) {
	started := time.Now()
	defer func() {
		units := textUnits(texts)
		units["returned_vectors"] = len(result)
		m.service.recordAuxiliaryUsage(ctx, m.GetModelID(), m.GetModelName(), "批量知识向量化", started, units, err)
	}()
	return m.Embedder.BatchEmbedWithPool(ctx, m.Embedder, texts)
}

func (s *Service) WrapVLM(model vlm.VLM) vlm.VLM { return &observedVLM{VLM: model, service: s} }

type observedVLM struct {
	vlm.VLM
	service *Service
}

func (m *observedVLM) Predict(ctx context.Context, images [][]byte, prompt string) (result string, err error) {
	started := time.Now()
	defer func() {
		units := textUnits([]string{prompt})
		units["input_images"], units["output_characters"] = len(images), utf8.RuneCountInString(result)
		m.service.recordAuxiliaryUsage(ctx, m.GetModelID(), m.GetModelName(), "图片识别", started, units, err)
	}()
	return m.VLM.Predict(ctx, images, prompt)
}

func (s *Service) WrapASR(model asr.ASR) asr.ASR { return &observedASR{ASR: model, service: s} }

type observedASR struct {
	asr.ASR
	service *Service
}

func (m *observedASR) Transcribe(ctx context.Context, audio []byte, filename string) (result *asr.TranscriptionResult, err error) {
	started := time.Now()
	defer func() {
		units := JSON{"input_audio_bytes": len(audio)}
		if result != nil {
			units["output_characters"] = utf8.RuneCountInString(result.Text)
		}
		m.service.recordAuxiliaryUsage(ctx, m.GetModelID(), m.GetModelName(), "语音转写", started, units, err)
	}()
	return m.ASR.Transcribe(ctx, audio, filename)
}

func (s *Service) WrapRerank(model rerank.Reranker) rerank.Reranker {
	return &observedRerank{Reranker: model, service: s}
}

type observedRerank struct {
	rerank.Reranker
	service *Service
}

func (m *observedRerank) Rerank(ctx context.Context, query string, documents []string) (result []rerank.RankResult, err error) {
	started := time.Now()
	defer func() {
		units := textUnits(documents)
		units["query_characters"] = utf8.RuneCountInString(query)
		m.service.recordAuxiliaryUsage(ctx, m.GetModelID(), m.GetModelName(), "知识检索精排", started, units, err)
	}()
	return m.Reranker.Rerank(ctx, query, documents)
}