internal/httpapi/openai/history/current_input_file.go

package history

import (
	"context"
	"crypto/rand"
	"math/big"
	"strings"
	"time"
	"unicode/utf8"

	"ds2api/internal/auth"
	"ds2api/internal/config"
	dsclient "ds2api/internal/deepseek/client"
	"ds2api/internal/httpapi/openai/shared"
	"ds2api/internal/promptcompat"
)

const (
	currentToolsFilename    = promptcompat.CurrentToolsContextFilename
	currentInputContentType = "text/plain; charset=utf-8"
	currentInputPurpose     = "assistants"
)

type CurrentInputConfigReader interface {
	CurrentInputFileEnabled() bool
	CurrentInputFileMinChars() int
}

type CurrentInputUploader interface {
	UploadFile(ctx context.Context, a *auth.RequestAuth, req dsclient.UploadFileRequest, maxAttempts int) (*dsclient.UploadFileResult, error)
}

type Service struct {
	Store CurrentInputConfigReader
	DS    CurrentInputUploader
}

func (s Service) ApplyCurrentInputFile(ctx context.Context, a *auth.RequestAuth, stdReq promptcompat.StandardRequest) (promptcompat.StandardRequest, error) {
	if stdReq.CurrentInputFileApplied || s.Store == nil || s.DS == nil || !s.Store.CurrentInputFileEnabled() {
		return stdReq, nil
	}
	modelID := strings.TrimSpace(stdReq.ResolvedModel)
	if modelID == "" {
		modelID = strings.TrimSpace(stdReq.RequestedModel)
	}
	modelType, ok := config.GetModelType(modelID)
	if !ok || modelType != "default" {
		return stdReq, nil
	}
	_, latestText := latestUserInputForFile(stdReq.Messages)
	if strings.TrimSpace(latestText) == "" {
		return stdReq, nil
	}
	minChars := s.Store.CurrentInputFileMinChars()
	if minChars > 0 && utf8.RuneCountInString(latestText) < minChars {
		return stdReq, nil
	}
	historyText := promptcompat.BuildOpenAICurrentInputContextTranscript(stdReq.Messages)
	if strings.TrimSpace(historyText) == "" {
		return stdReq, nil
	}
	filename := strings.TrimSpace(stdReq.CurrentInputFilename)
	if filename == "" {
		filename = randomContextFilename()
	}
	toolsText, toolNames := promptcompat.BuildOpenAIToolsContextTranscript(stdReq.ToolsRaw, stdReq.ToolChoice)
	hasToolsFile := strings.TrimSpace(toolsText) != ""

	historyResult, err := s.DS.UploadFile(ctx, a, dsclient.UploadFileRequest{
		Filename:    filename,
		ContentType: currentInputContentType,
		Purpose:     currentInputPurpose,
		ModelType:   modelType,
		Data:        []byte(historyText),
	}, 3)
	if err != nil {
		return stdReq, err
	}
	historyID := ""
	if historyResult != nil {
		historyID = strings.TrimSpace(historyResult.ID)
	}

	toolsID := ""
	if hasToolsFile {
		toolsResult, err := s.DS.UploadFile(ctx, a, dsclient.UploadFileRequest{
			Filename:    currentToolsFilename,
			ContentType: currentInputContentType,
			Purpose:     currentInputPurpose,
			ModelType:   modelType,
			Data:        []byte(toolsText),
		}, 3)
		if err != nil {
			return stdReq, err
		}
		if toolsResult != nil {
			toolsID = strings.TrimSpace(toolsResult.ID)
		}
	}

	promptMessages := make([]any, 0, 2)
	if hasToolsFile {
		toolInstructions, instructionNames := promptcompat.BuildOpenAIToolPromptInstructions(stdReq.ToolsRaw, stdReq.ToolChoice)
		if len(stdReq.ToolNames) == 0 {
			if len(instructionNames) > 0 {
				stdReq.ToolNames = instructionNames
			} else if len(toolNames) > 0 {
				stdReq.ToolNames = toolNames
			}
		}
		systemText := "Tool descriptions and parameter schemas are attached in context_tools.txt. Use only those tools."
		if strings.TrimSpace(toolInstructions) != "" {
			systemText += "\n\n" + toolInstructions
		}
		promptMessages = append(promptMessages, map[string]any{
			"role":    "system",
			"content": systemText,
		})
	} else if len(stdReq.ToolNames) == 0 && len(toolNames) > 0 {
		stdReq.ToolNames = toolNames
	}
	promptMessages = append(promptMessages, map[string]any{
		"role":    "user",
		"content": currentInputFilePrompt(filename, hasToolsFile),
	})
	finalPrompt, _ := promptcompat.BuildOpenAIPrompt(promptMessages, nil, "", stdReq.ToolChoice, stdReq.Thinking)
	if strings.TrimSpace(finalPrompt) == "" {
		finalPrompt = currentInputFilePrompt(filename, hasToolsFile)
	}

	promptTokenText := historyText
	if strings.TrimSpace(toolsText) != "" {
		if !strings.HasSuffix(promptTokenText, "\n") {
			promptTokenText += "\n"
		}
		promptTokenText += toolsText
	}

	stdReq.CurrentInputFileApplied = true
	stdReq.CurrentInputFileID = historyID
	stdReq.CurrentInputFilename = filename
	stdReq.CurrentToolsFileID = toolsID
	stdReq.HistoryText = historyText
	stdReq.Messages = promptMessages
	stdReq.FinalPrompt = finalPrompt
	stdReq.PromptTokenText = promptTokenText
	stdReq.RefFileIDs = prependUniqueRefFileIDs(stdReq.RefFileIDs, historyID, toolsID)
	return stdReq, nil
}

func (s Service) ReuploadAppliedCurrentInputFile(ctx context.Context, a *auth.RequestAuth, stdReq promptcompat.StandardRequest) (promptcompat.StandardRequest, error) {
	if !stdReq.CurrentInputFileApplied || s.Store == nil || s.DS == nil {
		return stdReq, nil
	}
	modelID := strings.TrimSpace(stdReq.ResolvedModel)
	if modelID == "" {
		modelID = strings.TrimSpace(stdReq.RequestedModel)
	}
	modelType, ok := config.GetModelType(modelID)
	if !ok || modelType != "default" {
		return stdReq, nil
	}
	historyText := stdReq.HistoryText
	if strings.TrimSpace(historyText) == "" {
		return stdReq, nil
	}
	filename := strings.TrimSpace(stdReq.CurrentInputFilename)
	if filename == "" {
		filename = randomContextFilename()
	}
	toolsText, _ := promptcompat.BuildOpenAIToolsContextTranscript(stdReq.ToolsRaw, stdReq.ToolChoice)
	hasToolsFile := strings.TrimSpace(toolsText) != ""

	newHistory, err := s.DS.UploadFile(ctx, a, dsclient.UploadFileRequest{
		Filename:    filename,
		ContentType: currentInputContentType,
		Purpose:     currentInputPurpose,
		ModelType:   modelType,
		Data:        []byte(historyText),
	}, 3)
	if err != nil {
		return stdReq, err
	}
	newHistoryID := ""
	if newHistory != nil {
		newHistoryID = strings.TrimSpace(newHistory.ID)
	}

	newToolsID := ""
	if hasToolsFile {
		newTools, err := s.DS.UploadFile(ctx, a, dsclient.UploadFileRequest{
			Filename:    currentToolsFilename,
			ContentType: currentInputContentType,
			Purpose:     currentInputPurpose,
			ModelType:   modelType,
			Data:        []byte(toolsText),
		}, 3)
		if err != nil {
			return stdReq, err
		}
		if newTools != nil {
			newToolsID = strings.TrimSpace(newTools.ID)
		}
	}

	stdReq.RefFileIDs = replaceGeneratedCurrentInputRefs(stdReq.RefFileIDs, stdReq.CurrentInputFileID, stdReq.CurrentToolsFileID, newHistoryID, newToolsID)
	stdReq.CurrentInputFileID = newHistoryID
	stdReq.CurrentInputFilename = filename
	stdReq.CurrentToolsFileID = newToolsID
	return stdReq, nil
}

func latestUserInputForFile(messages []any) (int, string) {
	for i := len(messages) - 1; i >= 0; i-- {
		msg, ok := messages[i].(map[string]any)
		if !ok {
			continue
		}
		role := strings.ToLower(strings.TrimSpace(shared.AsString(msg["role"])))
		if role != "user" {
			continue
		}
		text := promptcompat.NormalizeOpenAIContentForPrompt(msg["content"])
		if strings.TrimSpace(text) == "" {
			return -1, ""
		}
		return i, text
	}
	return -1, ""
}

func currentInputFilePrompt(filename string, hasToolsFile bool) string {
	base := []string{
		"已附上下文文件「%s」。请直接回答最新问题，不要提及文件、上下文或提示词本身。",
		"请从文件 %s 继续处理本轮任务，给出答案，避免说明你在使用上下文文件。",
		"上下文已存入 %s，请将其视为唯一上下文并直接答复；不要复述提示或提到文件。",
	}
	prompt := pickPromptVariant(base)
	text := strings.ReplaceAll(prompt, "%s", filename)
	if hasToolsFile {
		text += " 可用工具与参数说明在 context_tools.txt 中，仅可依据该文件调用工具。"
	}
	return text
}

func pickPromptVariant(variants []string) string {
	if len(variants) == 0 {
		return ""
	}
	idx := 0
	if len(variants) > 1 {
		max := big.NewInt(int64(len(variants)))
		if v, err := rand.Int(rand.Reader, max); err == nil {
			idx = int(v.Int64())
		} else {
			idx = int(time.Now().UnixNano()) % len(variants)
			if idx < 0 {
				idx = 0
			}
		}
	}
	return variants[idx]
}

func randomContextFilename() string {
	const letters = "abcdefghijklmnopqrstuvwxyz0123456789"
	length := 12
	max := big.NewInt(int64(len(letters)))
	b := make([]byte, 0, length)
	for i := 0; i < length; i++ {
		idx, err := rand.Int(rand.Reader, max)
		if err != nil {
			idx = big.NewInt(int64(time.Now().UnixNano() % int64(len(letters))))
		}
		b = append(b, letters[idx.Int64()])
	}
	return "ctx-" + string(b) + ".txt"
}

func prependUniqueRefFileIDs(existing []string, fileIDs ...string) []string {
	out := make([]string, 0, len(existing)+len(fileIDs))
	seen := map[string]struct{}{}
	for _, fileID := range fileIDs {
		trimmed := strings.TrimSpace(fileID)
		if trimmed == "" {
			continue
		}
		key := strings.ToLower(trimmed)
		if _, ok := seen[key]; ok {
			continue
		}
		out = append(out, trimmed)
		seen[key] = struct{}{}
	}
	for _, id := range existing {
		trimmed := strings.TrimSpace(id)
		if trimmed == "" {
			continue
		}
		key := strings.ToLower(trimmed)
		if _, ok := seen[key]; ok {
			continue
		}
		out = append(out, trimmed)
		seen[key] = struct{}{}
	}
	return out
}

func replaceGeneratedCurrentInputRefs(existing []string, oldHistoryID, oldToolsID, newHistoryID, newToolsID string) []string {
	filtered := make([]string, 0, len(existing))
	old := map[string]struct{}{}
	for _, id := range []string{oldHistoryID, oldToolsID} {
		trimmed := strings.ToLower(strings.TrimSpace(id))
		if trimmed != "" {
			old[trimmed] = struct{}{}
		}
	}
	for _, id := range existing {
		trimmed := strings.TrimSpace(id)
		if trimmed == "" {
			continue
		}
		if _, ok := old[strings.ToLower(trimmed)]; ok {
			continue
		}
		filtered = append(filtered, trimmed)
	}
	return prependUniqueRefFileIDs(filtered, newHistoryID, newToolsID)
}