MuyueWorkspace/internal/rag/inject.go

package rag

import (
	"fmt"
	"strings"
)

func BuildContextBlock(results []SearchResult, maxTokens int) string {
	if len(results) == 0 {
		return ""
	}
	if maxTokens <= 0 {
		maxTokens = 4000
	}
	maxChars := maxTokens * 4

	var b strings.Builder
	b.WriteString("<rag_context>\n")
	b.WriteString("The following context was retrieved from indexed documents to help answer the user's question.\n\n")

	for i, r := range results {
		entry := fmt.Sprintf("--- Source: %s (relevance: %.2f) ---\n%s\n\n", r.DocumentName, r.Score, r.Content)
		if b.Len()+len(entry) > maxChars {
			break
		}
		b.WriteString(entry)
		_ = i
	}

	b.WriteString("</rag_context>\n")
	return b.String()
}

func ExtractRAGQueries(message string) (queries []string, cleaned string) {
	cleaned = message
	parts := strings.Split(message, "#")
	if len(parts) <= 1 {
		return nil, message
	}

	var queryParts []string
	var textParts []string

	for i, part := range parts {
		if i == 0 {
			textParts = append(textParts, part)
			continue
		}
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}

		firstSpace := strings.IndexByte(part, ' ')
		newline := strings.IndexByte(part, '\n')

		end := len(part)
		if firstSpace > 0 && (newline < 0 || firstSpace < newline) {
			end = firstSpace
		} else if newline > 0 {
			end = newline
		}

		query := strings.TrimSpace(part[:end])
		if query != "" {
			queryParts = append(queryParts, query)
		}
		if end < len(part) {
			textParts = append(textParts, part[end:])
		}
	}

	if len(queryParts) > 0 {
		cleaned = strings.Join(textParts, " ")
		cleaned = strings.TrimSpace(cleaned)
	}

	return queryParts, cleaned
}