WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/workbench/auxiliary_usage.go

@@ -0,0 +1,139 @@
+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)
+}