geo/server/internal/tenant/app/kol_assist_generate_context.go

package app

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"

	"github.com/geo-platform/tenant-api/internal/shared/config"
	"go.uber.org/zap"
)

const kolAssistReferenceContentMaxRunes = 6000

var kolAssistURLPattern = regexp.MustCompile(`https?://[^\s]+`)

type kolAssistURLResolver struct {
	arkBaseURL string
	arkAPIKey  string
	httpClient *http.Client
	logger     *zap.Logger
}

type kolAssistReferenceMaterial struct {
	URL     string
	Title   string
	Content string
}

func newKolAssistURLResolver(llmCfg config.LLMConfig, logger *zap.Logger) *kolAssistURLResolver {
	baseURL := strings.TrimSpace(llmCfg.BaseURL)
	if baseURL == "" {
		baseURL = defaultKnowledgeArkBaseURL
	}

	return &kolAssistURLResolver{
		arkBaseURL: strings.TrimRight(baseURL, "/"),
		arkAPIKey:  strings.TrimSpace(llmCfg.APIKey),
		httpClient: &http.Client{Timeout: defaultKnowledgeURLParseTimeout},
		logger:     logger,
	}
}

func (r *kolAssistURLResolver) extract(ctx context.Context, rawURL string) (*kolAssistReferenceMaterial, error) {
	if r == nil || r.httpClient == nil || strings.TrimSpace(r.arkAPIKey) == "" {
		return nil, fmt.Errorf("url resolver is unavailable")
	}

	body, err := json.Marshal(arkToolExecuteRequest{
		ActionName: "LinkReader",
		ToolName:   "LinkReader",
		Timeout:    int(defaultKnowledgeURLParseTimeout.Seconds()),
		Parameters: map[string]any{
			"url_list": []string{rawURL},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("marshal webpage parser request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.arkBaseURL+"/tools/execute", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create webpage parser request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+r.arkAPIKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := r.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("execute webpage parser: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= http.StatusBadRequest {
		data, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return nil, fmt.Errorf("webpage parser failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
	}

	var parsed arkToolExecuteResponse
	if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
		return nil, fmt.Errorf("decode webpage parser response: %w", err)
	}
	if parsed.StatusCode != 0 && parsed.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("webpage parser returned status_code=%d", parsed.StatusCode)
	}
	if len(parsed.Data.ArkWebDataList) == 0 {
		return nil, fmt.Errorf("webpage parser returned no content")
	}

	item := parsed.Data.ArkWebDataList[0]
	if strings.TrimSpace(item.ErrorCode) != "" || strings.TrimSpace(item.ErrorMessage) != "" {
		return nil, fmt.Errorf("webpage parser returned error: %s %s", strings.TrimSpace(item.ErrorCode), strings.TrimSpace(item.ErrorMessage))
	}

	content := truncateKolAssistReferenceContent(normalizeKnowledgeText(item.Content))
	if content == "" {
		return nil, fmt.Errorf("webpage parser returned empty content")
	}

	return &kolAssistReferenceMaterial{
		URL:     strings.TrimSpace(rawURL),
		Title:   strings.TrimSpace(item.Title),
		Content: content,
	}, nil
}

func (s *KolAssistService) enrichGenerateDescription(ctx context.Context, description string) string {
	description = strings.TrimSpace(description)
	urls := extractKolAssistURLs(description)
	if len(urls) == 0 || s == nil || s.urlResolver == nil {
		return description
	}

	references := make([]kolAssistReferenceMaterial, 0, len(urls))
	for _, rawURL := range urls {
		material, err := s.urlResolver.extract(ctx, rawURL)
		if err != nil {
			if s.logger != nil {
				s.logger.Warn("kol assist url extraction failed",
					zap.String("url", rawURL),
					zap.Error(err),
				)
			}
			continue
		}
		references = append(references, *material)
	}

	if len(references) == 0 {
		return description
	}

	var builder strings.Builder
	builder.WriteString(description)
	builder.WriteString("\n\n以下是系统从参考网页提取到的正文内容。你必须基于这些内容学习文章的方法、结构、标题策略和风格，但不要复述具体事实，不要直接照抄原文。\n")
	for index, item := range references {
		builder.WriteString(fmt.Sprintf("\n[参考网页 %d]\nURL: %s\n", index+1, item.URL))
		if item.Title != "" {
			builder.WriteString("标题: ")
			builder.WriteString(item.Title)
			builder.WriteString("\n")
		}
		builder.WriteString("正文摘录:\n")
		builder.WriteString(item.Content)
		builder.WriteString("\n")
	}

	return builder.String()
}

func extractKolAssistURLs(input string) []string {
	matches := kolAssistURLPattern.FindAllString(input, -1)
	if len(matches) == 0 {
		return nil
	}

	result := make([]string, 0, len(matches))
	seen := make(map[string]struct{}, len(matches))
	for _, match := range matches {
		candidate := strings.TrimSpace(match)
		candidate = strings.Trim(candidate, `"'<>[](){}`)
		candidate = strings.TrimRight(candidate, ".,;，。；、)]}》」”")
		if candidate == "" {
			continue
		}
		if _, exists := seen[candidate]; exists {
			continue
		}
		seen[candidate] = struct{}{}
		result = append(result, candidate)
	}
	return result
}

func truncateKolAssistReferenceContent(content string) string {
	content = strings.TrimSpace(content)
	if content == "" {
		return ""
	}

	runes := []rune(content)
	if len(runes) <= kolAssistReferenceContentMaxRunes {
		return content
	}
	return strings.TrimSpace(string(runes[:kolAssistReferenceContentMaxRunes]))
}