oci-portal/internal/service/aigateway_extras.go

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"oci-portal/internal/aiwire"
	"oci-portal/internal/model"
	"oci-portal/internal/oci"
)

// moderationsMaxInputs 限制单次审核条数,防止批量滥用拖垮渠道配额。
const moderationsMaxInputs = 8

// SpeechBodyNormalize 校验并规范化 TTS 请求体:model / input 必填;
// xAI 上游把 language 当必填(实测缺失返回 422),缺省注入 "auto"。
// 其余字段(voice / response_format / extra_body 平铺项)原样保留。
func SpeechBodyNormalize(raw []byte) (string, []byte, error) {
	dec := json.NewDecoder(strings.NewReader(string(raw)))
	dec.UseNumber()
	var body map[string]any
	if err := dec.Decode(&body); err != nil {
		return "", nil, fmt.Errorf("解析请求体: %w", err)
	}
	modelName, _ := body["model"].(string)
	input, _ := body["input"].(string)
	if strings.TrimSpace(modelName) == "" || strings.TrimSpace(input) == "" {
		return "", nil, fmt.Errorf("model 与 input 不能为空")
	}
	if lang, ok := body["language"].(string); !ok || strings.TrimSpace(lang) == "" {
		body["language"] = "auto"
	}
	out, err := json.Marshal(body)
	if err != nil {
		return "", nil, fmt.Errorf("重组请求体: %w", err)
	}
	return modelName, out, nil
}

// ttsDefaultModel 是 /ai/v1/tts 缺省注入的模型(xAI 官方格式无 model 字段)。
const ttsDefaultModel = "xai.grok-tts"

// TtsBodyConvert 把 xAI 官方 TTS 格式转换为上游 OpenAI 兼容 audio/speech 形态:
// text→input、voice_id→voice,text 与 language 必填(对齐 xAI 官方);
// model 为网关扩展字段,缺省注入 ttsDefaultModel;其余字段原样保留。
func TtsBodyConvert(raw []byte) (string, []byte, error) {
	dec := json.NewDecoder(strings.NewReader(string(raw)))
	dec.UseNumber()
	var body map[string]any
	if err := dec.Decode(&body); err != nil {
		return "", nil, fmt.Errorf("解析请求体: %w", err)
	}
	text, _ := body["text"].(string)
	language, _ := body["language"].(string)
	if strings.TrimSpace(text) == "" || strings.TrimSpace(language) == "" {
		return "", nil, fmt.Errorf("text 与 language 不能为空")
	}
	modelName, _ := body["model"].(string)
	if strings.TrimSpace(modelName) == "" {
		modelName = ttsDefaultModel
		body["model"] = modelName
	}
	delete(body, "text")
	body["input"] = text
	if voice, ok := body["voice_id"].(string); ok && strings.TrimSpace(voice) != "" {
		body["voice"] = voice
	}
	delete(body, "voice_id")
	out, err := json.Marshal(body)
	if err != nil {
		return "", nil, fmt.Errorf("重组请求体: %w", err)
	}
	return modelName, out, nil
}

// Speech 编排 TTS 调用:按 TTS 能力选渠道,可重试错误换渠道(上限 3 次)。
func (s *AiGatewayService) Speech(ctx context.Context, modelName string, body []byte, group string) ([]byte, string, ChatMeta, error) {
	meta := ChatMeta{}
	excluded := map[uint]bool{}
	var lastErr error
	for attempt := 0; attempt < 3; attempt++ {
		cand, err := s.pick(ctx, modelName, group, "TTS", excluded)
		if err != nil {
			return nil, "", meta, firstErr(lastErr, err)
		}
		meta.ChannelID, meta.ChannelName = cand.ch.ID, cand.ch.Name
		cred, err := s.configs.credentialsByID(ctx, cand.ch.OciConfigID)
		if err != nil {
			return nil, "", meta, err
		}
		audio, contentType, err := s.client.GenAiCompatSpeech(ctx, cred, cand.ch.Region, body)
		if err == nil {
			s.markSuccess(ctx, cand.ch.ID)
			return audio, contentType, meta, nil
		}
		if done := s.retryStep(ctx, cand.ch.ID, err, excluded, &meta, &lastErr); done {
			return nil, "", meta, err
		}
	}
	return nil, "", meta, lastErr
}

// Rerank 编排文档重排:按 RERANK 能力选渠道,装配 Jina 风格响应。
func (s *AiGatewayService) Rerank(ctx context.Context, req aiwire.RerankRequest, group string) (*aiwire.RerankResponse, ChatMeta, error) {
	meta := ChatMeta{}
	excluded := map[uint]bool{}
	var lastErr error
	for attempt := 0; attempt < 3; attempt++ {
		cand, err := s.pick(ctx, req.Model, group, "RERANK", excluded)
		if err != nil {
			return nil, meta, firstErr(lastErr, err)
		}
		meta.ChannelID, meta.ChannelName = cand.ch.ID, cand.ch.Name
		cred, err := s.configs.credentialsByID(ctx, cand.ch.OciConfigID)
		if err != nil {
			return nil, meta, err
		}
		ranks, err := s.client.GenAiRerank(ctx, cred, cand.ch.Region, cand.modelOcid, req.Query, req.Documents, req.TopN)
		if err == nil {
			s.markSuccess(ctx, cand.ch.ID)
			return rerankResponse(req, ranks), meta, nil
		}
		if done := s.retryStep(ctx, cand.ch.ID, err, excluded, &meta, &lastErr); done {
			return nil, meta, err
		}
	}
	return nil, meta, lastErr
}

// retryStep 统一处理换渠道重试:不可重试返回 true 终止,可重试记罚并排除渠道。
func (s *AiGatewayService) retryStep(ctx context.Context, chID uint, err error, excluded map[uint]bool, meta *ChatMeta, lastErr *error) bool {
	retry, penalize := switchable(err)
	if !retry {
		return true
	}
	if penalize {
		s.markFailure(ctx, chID)
	}
	excluded[chID] = true
	meta.Retries++
	*lastErr = err
	return false
}

// rerankResponse 把上游排序结果装配为对外响应;return_documents 时回填原文。
func rerankResponse(req aiwire.RerankRequest, ranks []oci.RerankRank) *aiwire.RerankResponse {
	out := &aiwire.RerankResponse{Model: req.Model, Results: make([]aiwire.RerankResult, 0, len(ranks))}
	withDoc := req.ReturnDocuments != nil && *req.ReturnDocuments
	for _, r := range ranks {
		if r.Index < 0 || r.Index >= len(req.Documents) {
			continue
		}
		item := aiwire.RerankResult{Index: r.Index, RelevanceScore: r.Score}
		if withDoc {
			item.Document = &aiwire.RerankDocument{Text: req.Documents[r.Index]}
		}
		out.Results = append(out.Results, item)
	}
	return out
}

// ModerationInputs 解析 moderations 的 input 字段:string 或 []string,
// 条数与空值校验(上限 moderationsMaxInputs)。
func ModerationInputs(raw json.RawMessage) ([]string, error) {
	var single string
	if err := json.Unmarshal(raw, &single); err == nil {
		if strings.TrimSpace(single) == "" {
			return nil, fmt.Errorf("input 不能为空")
		}
		return []string{single}, nil
	}
	var many []string
	if err := json.Unmarshal(raw, &many); err != nil {
		return nil, fmt.Errorf("input 需为字符串或字符串数组")
	}
	if len(many) == 0 || len(many) > moderationsMaxInputs {
		return nil, fmt.Errorf("input 条数需在 1~%d 之间", moderationsMaxInputs)
	}
	for _, item := range many {
		if strings.TrimSpace(item) == "" {
			return nil, fmt.Errorf("input 含空条目")
		}
	}
	return many, nil
}

// Moderations 编排内容审核:guardrails 为服务级 API 无模型维度,
// 按分组选启用渠道(整批同渠道),可重试错误换渠道(上限 3 次)。
func (s *AiGatewayService) Moderations(ctx context.Context, id string, inputs []string, group string) (*aiwire.ModerationsResponse, ChatMeta, error) {
	meta := ChatMeta{}
	excluded := map[uint]bool{}
	var lastErr error
	for attempt := 0; attempt < 3; attempt++ {
		ch, err := s.pickGuardChannel(ctx, group, excluded)
		if err != nil {
			return nil, meta, firstErr(lastErr, err)
		}
		meta.ChannelID, meta.ChannelName = ch.ID, ch.Name
		resp, err := s.moderateOnce(ctx, ch, id, inputs)
		if err == nil {
			s.markSuccess(ctx, ch.ID)
			return resp, meta, nil
		}
		if done := s.retryStep(ctx, ch.ID, err, excluded, &meta, &lastErr); done {
			return nil, meta, err
		}
	}
	return nil, meta, lastErr
}

// pickGuardChannel 按分组挑一个启用且未熔断的渠道(权重加成,不查模型缓存)。
func (s *AiGatewayService) pickGuardChannel(ctx context.Context, group string, excluded map[uint]bool) (model.AiChannel, error) {
	q := s.db.WithContext(ctx).Where("enabled = ?", true)
	if group != "" {
		q = q.Where("channel_group = ?", group)
	}
	var channels []model.AiChannel
	if err := q.Find(&channels).Error; err != nil {
		return model.AiChannel{}, err
	}
	now := time.Now()
	avail := channels[:0]
	for _, ch := range channels {
		if excluded[ch.ID] || (ch.DisabledUntil != nil && ch.DisabledUntil.After(now)) {
			continue
		}
		avail = append(avail, ch)
	}
	if len(avail) == 0 {
		return model.AiChannel{}, ErrAiNoChannel
	}
	return weightedPick(topPriority(avail)), nil
}

// moderateOnce 在单渠道上逐条审核并装配 OpenAI moderations 外壳。
func (s *AiGatewayService) moderateOnce(ctx context.Context, ch model.AiChannel, id string, inputs []string) (*aiwire.ModerationsResponse, error) {
	cred, err := s.configs.credentialsByID(ctx, ch.OciConfigID)
	if err != nil {
		return nil, err
	}
	out := &aiwire.ModerationsResponse{ID: id, Model: "oci-guardrails",
		Results: make([]aiwire.ModerationResult, 0, len(inputs))}
	for _, text := range inputs {
		outcome, err := s.client.GenAiApplyGuardrails(ctx, cred, ch.Region, text)
		if err != nil {
			return nil, err
		}
		out.Results = append(out.Results, moderationResult(outcome))
	}
	return out, nil
}

// moderationFlagThreshold 是判定违规的得分阈值(上游 CM / PI 为二值得分)。
const moderationFlagThreshold = 0.5

// moderationResult 把 guardrails 结果映射为 OpenAI moderations 条目:
// flagged 由内容审核与提示注入判定,PII 命中只作扩展信息不参与 flagged。
func moderationResult(outcome *oci.GuardrailsOutcome) aiwire.ModerationResult {
	res := aiwire.ModerationResult{Categories: map[string]bool{}, CategoryScores: map[string]float64{}}
	for _, cat := range outcome.Categories {
		key := strings.ToLower(cat.Name)
		res.Categories[key] = cat.Score >= moderationFlagThreshold
		res.CategoryScores[key] = cat.Score
		if res.Categories[key] {
			res.Flagged = true
		}
	}
	if outcome.PromptInjectionScore != nil {
		hit := *outcome.PromptInjectionScore >= moderationFlagThreshold
		res.Categories["prompt_injection"] = hit
		res.CategoryScores["prompt_injection"] = *outcome.PromptInjectionScore
		if hit {
			res.Flagged = true
		}
	}
	for _, p := range outcome.Pii {
		res.Pii = append(res.Pii, aiwire.ModerationPii{Text: p.Text, Label: p.Label,
			Score: p.Score, Offset: p.Offset, Length: p.Length})
	}
	return res
}