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

package app

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

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

const (
	defaultKnowledgeArkBaseURL             = "https://ark.cn-beijing.volces.com/api/v3"
	defaultKnowledgeURLParseTimeout        = 60 * time.Second
	defaultKnowledgeURLMarkdownTimeout     = 90 * time.Second
	defaultKnowledgeURLMarkdownChunkChars  = 12000
	defaultKnowledgeURLMarkdownOutputToken = 12000
)

type arkToolExecuteRequest struct {
	ActionName string         `json:"action_name"`
	ToolName   string         `json:"tool_name"`
	Timeout    int            `json:"timeout,omitempty"`
	Parameters map[string]any `json:"parameters,omitempty"`
}

type arkToolExecuteResponse struct {
	StatusCode int                `json:"status_code"`
	Data       arkToolExecuteData `json:"data"`
}

type arkToolExecuteData struct {
	ArkWebDataList []arkWebDataItem `json:"ark_web_data_list"`
}

type arkWebDataItem struct {
	URL          string `json:"url"`
	Title        string `json:"title"`
	Content      string `json:"content"`
	ErrorMessage string `json:"error_message"`
	ErrorCode    string `json:"error_code"`
}

func (s *KnowledgeService) parseWebsiteKnowledge(ctx context.Context, rawURL string) (*knowledgeParsedContent, error) {
	item, err := s.executeLinkReader(ctx, rawURL)
	if err != nil {
		return nil, err
	}

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

	title := strings.TrimSpace(item.Title)
	markdown, err := s.formatWebsiteMarkdown(ctx, title, text, rawURL)
	if err != nil {
		if s.logger != nil {
			s.logger.Warn("knowledge website markdown formatting failed, fallback to local markdown",
				zap.String("url", rawURL),
				zap.Error(err),
			)
		}
		markdown = buildKnowledgeMarkdownFromText(title, text)
	}

	return &knowledgeParsedContent{
		Text:     text,
		Markdown: markdown,
	}, nil
}

func (s *KnowledgeService) executeLinkReader(ctx context.Context, rawURL string) (*arkWebDataItem, error) {
	baseURL := strings.TrimRight(strings.TrimSpace(s.arkBaseURL), "/")
	if baseURL == "" {
		baseURL = defaultKnowledgeArkBaseURL
	}

	timeout := defaultKnowledgeURLParseTimeout
	callCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	body, err := json.Marshal(arkToolExecuteRequest{
		ActionName: "LinkReader",
		ToolName:   "LinkReader",
		Timeout:    int(timeout / time.Second),
		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(callCtx, http.MethodPost, baseURL+"/tools/execute", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create webpage parser request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+s.arkAPIKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := s.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))
	}
	return &item, nil
}

func (s *KnowledgeService) formatWebsiteMarkdown(ctx context.Context, title, content, rawURL string) (string, error) {
	content = strings.TrimSpace(content)
	if content == "" {
		return "", fmt.Errorf("website content is empty")
	}

	chunks := splitKnowledgeMarkdownChunks(content, defaultKnowledgeURLMarkdownChunkChars)
	if len(chunks) == 0 {
		return "", fmt.Errorf("website content is empty after chunking")
	}

	sections := make([]string, 0, len(chunks)+1)
	for index, chunk := range chunks {
		result, err := s.llmClient.Generate(ctx, llm.GenerateRequest{
			Model:           s.urlMarkdownModel,
			Prompt:          prompts.KnowledgeWebsiteMarkdownPrompt(title, rawURL, chunk, index, len(chunks)),
			Timeout:         defaultKnowledgeURLMarkdownTimeout,
			MaxOutputTokens: defaultKnowledgeURLMarkdownOutputToken,
		}, nil)
		if err != nil {
			return "", err
		}

		section := sanitizeLLMMarkdown(strings.TrimSpace(result.Content))
		if section == "" {
			continue
		}
		if index > 0 && !strings.HasPrefix(section, "#") {
			section = fmt.Sprintf("## 第 %d 部分\n\n%s", index+1, section)
		}
		sections = append(sections, section)
	}

	if len(sections) == 0 {
		return "", fmt.Errorf("website formatter returned empty markdown")
	}

	markdown := strings.TrimSpace(strings.Join(sections, "\n\n"))
	if title != "" && !strings.HasPrefix(markdown, "# ") {
		markdown = "# " + title + "\n\n" + markdown
	}
	return markdown, nil
}

func splitKnowledgeMarkdownChunks(content string, maxChars int) []string {
	if maxChars <= 0 {
		maxChars = defaultKnowledgeURLMarkdownChunkChars
	}

	lines := strings.Split(content, "\n")
	chunks := make([]string, 0, len(lines))
	current := make([]string, 0, 16)
	currentLen := 0

	flush := func() {
		if len(current) == 0 {
			return
		}
		chunks = append(chunks, strings.TrimSpace(strings.Join(current, "\n")))
		current = current[:0]
		currentLen = 0
	}

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		if len([]rune(line)) > maxChars {
			flush()
			for _, segment := range splitKnowledgeLongLine(line, maxChars) {
				if strings.TrimSpace(segment) != "" {
					chunks = append(chunks, segment)
				}
			}
			continue
		}

		nextLen := currentLen + len([]rune(line))
		if len(current) > 0 {
			nextLen++
		}
		if nextLen > maxChars {
			flush()
		}

		current = append(current, line)
		currentLen += len([]rune(line))
		if len(current) > 1 {
			currentLen++
		}
	}

	flush()
	return chunks
}

func splitKnowledgeLongLine(line string, maxChars int) []string {
	runes := []rune(line)
	if len(runes) <= maxChars {
		return []string{line}
	}

	segments := make([]string, 0, (len(runes)/maxChars)+1)
	for start := 0; start < len(runes); start += maxChars {
		end := start + maxChars
		if end > len(runes) {
			end = len(runes)
		}
		segments = append(segments, strings.TrimSpace(string(runes[start:end])))
	}
	return segments
}

func sanitizeLLMMarkdown(content string) string {
	content = strings.TrimSpace(content)
	if !strings.HasPrefix(content, "```") {
		return content
	}

	lines := strings.Split(content, "\n")
	if len(lines) < 3 {
		return strings.Trim(content, "`")
	}
	if !strings.HasPrefix(lines[0], "```") || strings.TrimSpace(lines[len(lines)-1]) != "```" {
		return content
	}
	return strings.TrimSpace(strings.Join(lines[1:len(lines)-1], "\n"))
}