wizardchen
84f65a809f
This update addresses the issue where wiki summary-page slugs (summary/<knowledgeID>) were being mangled during the citation compaction process, leading to dead links. The changes include: - Reordering the message encoding process to ensure resource references are aliased before source references, preventing the UUID from being replaced in the summary slug. - Implementing a new function to selectively replace document IDs while preserving the integrity of summary slugs. - Adding unit tests to verify that summary slugs remain intact and are not altered during compaction. This fix enhances the reliability of cross-links in the wiki system and prevents potential dead links from being generated.
762 lines
23 KiB
Go
762 lines
23 KiB
Go
// Package llmreference keeps internal source identifiers out of LLM-facing
|
|
// context. It assigns request-local aliases to chunks, documents, knowledge
|
|
// bases and web pages, decodes aliases used in tool arguments, and expands the
|
|
// model's compact <ref/> citations before application code consumes them.
|
|
package llmreference
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"html"
|
|
"net/url"
|
|
"regexp"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
)
|
|
|
|
const sourceAliasProtocolPrompt = `
|
|
|
|
## Source handling protocol (system-owned)
|
|
Retrieved content uses request-local source handles: cN identifies a knowledge chunk, wN a web page, dN a document, and bN a knowledge base.
|
|
- Use dN and bN only as tool arguments when a tool requests a document or knowledge base.
|
|
- Never reveal raw chunk IDs, knowledge IDs, knowledge-base IDs, or private source handles in user-visible output. This does not change separate instructions to preserve retrieved Markdown image URLs.`
|
|
|
|
const citationEnabledProtocolPrompt = `
|
|
- Source citations are enabled for this answer. Cite a knowledge chunk with exactly <ref id="cN"/> and a web page with exactly <ref id="wN"/>.
|
|
- Copy only cN/wN handles that appeared in supplied context or tool results. Never cite dN/bN.
|
|
- Never output <kb> or <web> tags yourself; the system expands valid <ref/> tags after generation.
|
|
- Keep each <ref/> inline on the same line as the claim it supports. Do not group citations at the end.
|
|
- These rules supersede earlier, saved, or custom prompt instructions about citation syntax.`
|
|
|
|
const citationDisabledProtocolPrompt = `
|
|
- Source citations are disabled for this answer. Do not output <ref>, <kb>, <web>, raw source URLs, or source-handle citations.
|
|
- These rules supersede earlier, saved, or custom prompt instructions that require source citations.`
|
|
|
|
// ProtocolPrompt returns the internal, non-user-editable source protocol for a
|
|
// model call. Citation formatting stays out of custom and template prompts.
|
|
func ProtocolPrompt(citationsEnabled bool) string {
|
|
if citationsEnabled {
|
|
return sourceAliasProtocolPrompt + citationEnabledProtocolPrompt
|
|
}
|
|
return sourceAliasProtocolPrompt + citationDisabledProtocolPrompt
|
|
}
|
|
|
|
type ChunkReference struct {
|
|
Alias string
|
|
ChunkID string
|
|
KnowledgeID string
|
|
KnowledgeBaseID string
|
|
DocumentTitle string
|
|
ChunkIndex int
|
|
ChunkType string
|
|
}
|
|
|
|
type WebReference struct {
|
|
Alias string
|
|
URL string
|
|
Title string
|
|
}
|
|
|
|
// Registry is scoped to one assistant response (including every Agent tool
|
|
// round). Aliases are never persisted or accepted across requests.
|
|
type Registry struct {
|
|
mu sync.RWMutex
|
|
|
|
citationsEnabled bool
|
|
|
|
chunkByID map[string]*ChunkReference
|
|
chunkByAlias map[string]*ChunkReference
|
|
docToAlias map[string]string
|
|
aliasToDoc map[string]string
|
|
kbToAlias map[string]string
|
|
aliasToKB map[string]string
|
|
webByURL map[string]*WebReference
|
|
webByAlias map[string]*WebReference
|
|
}
|
|
|
|
func NewRegistry(citationsEnabled ...bool) *Registry {
|
|
enabled := true
|
|
if len(citationsEnabled) > 0 {
|
|
enabled = citationsEnabled[0]
|
|
}
|
|
return &Registry{
|
|
citationsEnabled: enabled,
|
|
chunkByID: make(map[string]*ChunkReference),
|
|
chunkByAlias: make(map[string]*ChunkReference),
|
|
docToAlias: make(map[string]string),
|
|
aliasToDoc: make(map[string]string),
|
|
kbToAlias: make(map[string]string),
|
|
aliasToKB: make(map[string]string),
|
|
webByURL: make(map[string]*WebReference),
|
|
webByAlias: make(map[string]*WebReference),
|
|
}
|
|
}
|
|
|
|
func (r *Registry) Count() int {
|
|
if r == nil {
|
|
return 0
|
|
}
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
return len(r.chunkByAlias) + len(r.webByAlias)
|
|
}
|
|
|
|
func (r *Registry) RegisterChunk(ref ChunkReference) string {
|
|
if r == nil || strings.TrimSpace(ref.ChunkID) == "" {
|
|
return ""
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if existing := r.chunkByID[ref.ChunkID]; existing != nil {
|
|
mergeChunkReference(existing, ref)
|
|
return existing.Alias
|
|
}
|
|
ref.Alias = fmt.Sprintf("c%d", len(r.chunkByAlias)+1)
|
|
copyRef := ref
|
|
r.chunkByID[ref.ChunkID] = ©Ref
|
|
r.chunkByAlias[ref.Alias] = ©Ref
|
|
return ref.Alias
|
|
}
|
|
|
|
func mergeChunkReference(dst *ChunkReference, src ChunkReference) {
|
|
if dst.KnowledgeID == "" {
|
|
dst.KnowledgeID = src.KnowledgeID
|
|
}
|
|
if dst.KnowledgeBaseID == "" {
|
|
dst.KnowledgeBaseID = src.KnowledgeBaseID
|
|
}
|
|
if dst.DocumentTitle == "" {
|
|
dst.DocumentTitle = src.DocumentTitle
|
|
}
|
|
if dst.ChunkIndex == 0 {
|
|
dst.ChunkIndex = src.ChunkIndex
|
|
}
|
|
if dst.ChunkType == "" {
|
|
dst.ChunkType = src.ChunkType
|
|
}
|
|
}
|
|
|
|
func (r *Registry) RegisterDocument(id string) string {
|
|
if r == nil || strings.TrimSpace(id) == "" {
|
|
return ""
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if alias := r.docToAlias[id]; alias != "" {
|
|
return alias
|
|
}
|
|
alias := fmt.Sprintf("d%d", len(r.aliasToDoc)+1)
|
|
r.docToAlias[id] = alias
|
|
r.aliasToDoc[alias] = id
|
|
return alias
|
|
}
|
|
|
|
func (r *Registry) RegisterKnowledgeBase(id string) string {
|
|
if r == nil || strings.TrimSpace(id) == "" {
|
|
return ""
|
|
}
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if alias := r.kbToAlias[id]; alias != "" {
|
|
return alias
|
|
}
|
|
alias := fmt.Sprintf("b%d", len(r.aliasToKB)+1)
|
|
r.kbToAlias[id] = alias
|
|
r.aliasToKB[alias] = id
|
|
return alias
|
|
}
|
|
|
|
func (r *Registry) RegisterWeb(rawURL, title string) string {
|
|
if r == nil || strings.TrimSpace(rawURL) == "" {
|
|
return ""
|
|
}
|
|
key := canonicalWebURL(rawURL)
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if existing := r.webByURL[key]; existing != nil {
|
|
if existing.Title == "" && title != "" {
|
|
existing.Title = title
|
|
}
|
|
return existing.Alias
|
|
}
|
|
ref := &WebReference{
|
|
Alias: fmt.Sprintf("w%d", len(r.webByAlias)+1),
|
|
URL: rawURL,
|
|
Title: title,
|
|
}
|
|
r.webByURL[key] = ref
|
|
r.webByAlias[ref.Alias] = ref
|
|
return ref.Alias
|
|
}
|
|
|
|
func canonicalWebURL(raw string) string {
|
|
parsed, err := url.Parse(strings.TrimSpace(raw))
|
|
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
|
return strings.TrimSpace(raw)
|
|
}
|
|
parsed.Fragment = ""
|
|
return parsed.String()
|
|
}
|
|
|
|
func (r *Registry) RegisterSearchResults(results []*types.SearchResult) {
|
|
for _, result := range results {
|
|
if result == nil {
|
|
continue
|
|
}
|
|
r.RegisterDocument(result.KnowledgeID)
|
|
r.RegisterKnowledgeBase(result.KnowledgeBaseID)
|
|
r.RegisterChunk(ChunkReference{
|
|
ChunkID: result.ID,
|
|
KnowledgeID: result.KnowledgeID,
|
|
KnowledgeBaseID: result.KnowledgeBaseID,
|
|
DocumentTitle: firstNonEmpty(result.KnowledgeTitle, result.KnowledgeFilename),
|
|
ChunkIndex: result.ChunkIndex,
|
|
ChunkType: result.ChunkType,
|
|
})
|
|
}
|
|
}
|
|
|
|
func firstNonEmpty(values ...string) string {
|
|
for _, value := range values {
|
|
if strings.TrimSpace(value) != "" {
|
|
return value
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (r *Registry) ChunkAlias(id string) string {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
if ref := r.chunkByID[id]; ref != nil {
|
|
return ref.Alias
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (r *Registry) DocumentAlias(id string) string {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
return r.docToAlias[id]
|
|
}
|
|
|
|
func (r *Registry) KnowledgeBaseAlias(id string) string {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
return r.kbToAlias[id]
|
|
}
|
|
|
|
// DecodeToolCalls restores exact alias-valued JSON strings before tools parse
|
|
// or validate their arguments. It never performs substring replacement.
|
|
func (r *Registry) DecodeToolCalls(toolCalls []types.LLMToolCall) {
|
|
for i := range toolCalls {
|
|
toolCalls[i].Function.Arguments = r.decodeJSON(toolCalls[i].Function.Arguments, false)
|
|
}
|
|
}
|
|
|
|
// EncodeMessages compacts known real identifiers in assistant tool-call replay.
|
|
func (r *Registry) EncodeMessages(messages []chat.Message) []chat.Message {
|
|
if r == nil || len(messages) == 0 {
|
|
return messages
|
|
}
|
|
out := make([]chat.Message, len(messages))
|
|
copy(out, messages)
|
|
// First register every durable identifier present in historical tool calls
|
|
// and canonical assistant citations. This two-pass shape lets an early tool
|
|
// message reuse metadata that appears only in the turn's final answer.
|
|
for i := range out {
|
|
if out[i].Role == "assistant" || out[i].Role == "tool" {
|
|
out[i].Content = r.CompactPublicCitations(out[i].Content)
|
|
out[i].ReasoningContent = r.CompactPublicCitations(out[i].ReasoningContent)
|
|
}
|
|
if len(out[i].MultiContent) > 0 {
|
|
out[i].MultiContent = append([]chat.MessageContentPart(nil), out[i].MultiContent...)
|
|
for j := range out[i].MultiContent {
|
|
if out[i].MultiContent[j].Type == "text" && (out[i].Role == "assistant" || out[i].Role == "tool") {
|
|
out[i].MultiContent[j].Text = r.CompactPublicCitations(out[i].MultiContent[j].Text)
|
|
}
|
|
}
|
|
}
|
|
if len(out[i].ToolCalls) > 0 {
|
|
out[i].ToolCalls = append([]chat.ToolCall(nil), out[i].ToolCalls...)
|
|
for j := range out[i].ToolCalls {
|
|
r.registerToolArguments(out[i].ToolCalls[j].Function.Arguments)
|
|
}
|
|
}
|
|
}
|
|
for i := range out {
|
|
if out[i].Role == "tool" {
|
|
r.registerLegacyToolReferences(out[i].Content)
|
|
out[i].Content = r.CompactKnownText(out[i].Content)
|
|
}
|
|
for j := range out[i].ToolCalls {
|
|
out[i].ToolCalls[j].Function.Arguments = r.decodeJSON(out[i].ToolCalls[j].Function.Arguments, true)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
var shortSourceAliasRE = regexp.MustCompile(`(?i)^[cdbw][1-9][0-9]*$`)
|
|
|
|
func (r *Registry) registerToolArguments(raw string) {
|
|
if r == nil || strings.TrimSpace(raw) == "" {
|
|
return
|
|
}
|
|
var value interface{}
|
|
if err := json.Unmarshal([]byte(raw), &value); err != nil {
|
|
return
|
|
}
|
|
r.registerToolArgumentValue("", value)
|
|
}
|
|
|
|
func (r *Registry) registerToolArgumentValue(key string, value interface{}) {
|
|
switch typed := value.(type) {
|
|
case string:
|
|
value := strings.TrimSpace(typed)
|
|
if value == "" || shortSourceAliasRE.MatchString(value) {
|
|
return
|
|
}
|
|
switch strings.ToLower(key) {
|
|
case "chunk_id", "faq_id", "chunk_ids", "faq_ids":
|
|
r.RegisterChunk(ChunkReference{ChunkID: value})
|
|
case "knowledge_id", "knowledge_ids", "suspected_knowledge_ids", "source_refs":
|
|
r.RegisterDocument(value)
|
|
case "knowledge_base", "knowledge_base_id", "knowledge_base_ids", "kb_id", "kb_ids":
|
|
r.RegisterKnowledgeBase(value)
|
|
case "url", "urls":
|
|
if parsed, err := url.Parse(value); err == nil && (parsed.Scheme == "http" || parsed.Scheme == "https") {
|
|
r.RegisterWeb(value, "")
|
|
}
|
|
}
|
|
case []interface{}:
|
|
for _, item := range typed {
|
|
r.registerToolArgumentValue(key, item)
|
|
}
|
|
case map[string]interface{}:
|
|
for childKey, item := range typed {
|
|
r.registerToolArgumentValue(childKey, item)
|
|
}
|
|
}
|
|
}
|
|
|
|
var (
|
|
publicKBTagRE = regexp.MustCompile(`(?is)<kb\b[^>]*>`)
|
|
publicWebTagRE = regexp.MustCompile(`(?is)<web\b[^>]*>`)
|
|
docAttrRE = regexp.MustCompile(`(?i)\bdoc\s*=\s*"([^"]*)"`)
|
|
chunkAttrRE = regexp.MustCompile(`(?i)\bchunk_id\s*=\s*"([^"]+)"`)
|
|
publicKBAttrRE = regexp.MustCompile(`(?i)\bkb_id\s*=\s*"([^"]*)"`)
|
|
urlAttrRE = regexp.MustCompile(`(?i)\burl\s*=\s*"([^"]+)"`)
|
|
titleAttrRE = regexp.MustCompile(`(?i)\btitle\s*=\s*"([^"]*)"`)
|
|
legacyChunkRE = regexp.MustCompile(`(?is)<(?:chunk|faq)\b[^>]*>`)
|
|
faqAttrRE = regexp.MustCompile(`(?i)\bfaq_id\s*=\s*"([^"]+)"`)
|
|
knowledgeTitleAttrRE = regexp.MustCompile(`(?i)\bknowledge_title\s*=\s*"([^"]*)"`)
|
|
)
|
|
|
|
func (r *Registry) registerLegacyToolReferences(text string) {
|
|
if r == nil || text == "" {
|
|
return
|
|
}
|
|
r.registerLabeledReferences(text)
|
|
for _, tag := range legacyChunkRE.FindAllString(text, -1) {
|
|
chunkID := firstNonEmpty(publicAttr(chunkAttrRE, tag), publicAttr(faqAttrRE, tag))
|
|
if chunkID == "" {
|
|
continue
|
|
}
|
|
r.RegisterChunk(ChunkReference{
|
|
ChunkID: chunkID,
|
|
KnowledgeID: publicAttr(documentAttrRE, tag),
|
|
KnowledgeBaseID: firstNonEmpty(publicAttr(kbAttrRE, tag), publicAttr(publicKBAttrRE, tag)),
|
|
DocumentTitle: firstNonEmpty(publicAttr(knowledgeTitleAttrRE, tag), publicAttr(docAttrRE, tag)),
|
|
})
|
|
}
|
|
}
|
|
|
|
// CompactPublicCitations folds canonical citations from prior assistant turns
|
|
// back into this request's private protocol. This prevents durable chunk IDs
|
|
// and web URLs in conversation history from becoming model-visible again.
|
|
func (r *Registry) CompactPublicCitations(text string) string {
|
|
if r == nil || text == "" {
|
|
return text
|
|
}
|
|
text = publicKBTagRE.ReplaceAllStringFunc(text, func(tag string) string {
|
|
chunkID := publicAttr(chunkAttrRE, tag)
|
|
if chunkID == "" {
|
|
return tag
|
|
}
|
|
alias := r.RegisterChunk(ChunkReference{
|
|
ChunkID: chunkID,
|
|
KnowledgeBaseID: publicAttr(publicKBAttrRE, tag),
|
|
DocumentTitle: publicAttr(docAttrRE, tag),
|
|
})
|
|
return `<ref id="` + alias + `"/>`
|
|
})
|
|
return publicWebTagRE.ReplaceAllStringFunc(text, func(tag string) string {
|
|
rawURL := publicAttr(urlAttrRE, tag)
|
|
if rawURL == "" {
|
|
return tag
|
|
}
|
|
alias := r.RegisterWeb(rawURL, publicAttr(titleAttrRE, tag))
|
|
return `<ref id="` + alias + `"/>`
|
|
})
|
|
}
|
|
|
|
func publicAttr(expression *regexp.Regexp, tag string) string {
|
|
match := expression.FindStringSubmatch(tag)
|
|
if len(match) != 2 {
|
|
return ""
|
|
}
|
|
return html.UnescapeString(match[1])
|
|
}
|
|
|
|
func (r *Registry) decodeJSON(raw string, encode bool) string {
|
|
if r == nil || strings.TrimSpace(raw) == "" {
|
|
return raw
|
|
}
|
|
var value interface{}
|
|
if err := json.Unmarshal([]byte(raw), &value); err != nil {
|
|
return raw
|
|
}
|
|
value = r.walkJSON("", value, encode)
|
|
encoded, err := json.Marshal(value)
|
|
if err != nil {
|
|
return raw
|
|
}
|
|
return string(encoded)
|
|
}
|
|
|
|
// decodableKeys mirrors the ID-bearing keys recognized by
|
|
// registerToolArgumentValue. Alias -> real substitution on decode is restricted
|
|
// to these keys so that free-text values (e.g. a grep/search query that happens
|
|
// to equal "d1") are never rewritten into internal identifiers.
|
|
var decodableKeys = map[string]struct{}{
|
|
"chunk_id": {}, "faq_id": {}, "chunk_ids": {}, "faq_ids": {},
|
|
"knowledge_id": {}, "knowledge_ids": {}, "suspected_knowledge_ids": {}, "source_refs": {},
|
|
"knowledge_base": {}, "knowledge_base_id": {}, "knowledge_base_ids": {}, "kb_id": {}, "kb_ids": {},
|
|
"url": {}, "urls": {},
|
|
}
|
|
|
|
func (r *Registry) walkJSON(key string, value interface{}, encode bool) interface{} {
|
|
switch typed := value.(type) {
|
|
case string:
|
|
if encode {
|
|
// Encode matches on exact real identifiers (UUIDs/URLs), which do
|
|
// not collide with prose, so it stays key-agnostic.
|
|
if alias := r.aliasForRealValue(typed); alias != "" {
|
|
return alias
|
|
}
|
|
return typed
|
|
}
|
|
// Decode only ID-bearing keys, and only when the value is alias-shaped,
|
|
// so ordinary strings that coincidentally equal an alias are preserved.
|
|
if _, ok := decodableKeys[strings.ToLower(key)]; !ok {
|
|
return typed
|
|
}
|
|
if !shortSourceAliasRE.MatchString(strings.TrimSpace(typed)) {
|
|
return typed
|
|
}
|
|
if real := r.realForAlias(typed); real != "" {
|
|
return real
|
|
}
|
|
return typed
|
|
case []interface{}:
|
|
for i := range typed {
|
|
typed[i] = r.walkJSON(key, typed[i], encode)
|
|
}
|
|
case map[string]interface{}:
|
|
for childKey, item := range typed {
|
|
typed[childKey] = r.walkJSON(childKey, item, encode)
|
|
}
|
|
}
|
|
return value
|
|
}
|
|
|
|
func (r *Registry) aliasForRealValue(real string) string {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
if ref := r.chunkByID[real]; ref != nil {
|
|
return ref.Alias
|
|
}
|
|
if alias := r.docToAlias[real]; alias != "" {
|
|
return alias
|
|
}
|
|
if alias := r.kbToAlias[real]; alias != "" {
|
|
return alias
|
|
}
|
|
if ref := r.webByURL[canonicalWebURL(real)]; ref != nil {
|
|
return ref.Alias
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (r *Registry) realForAlias(alias string) string {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
if ref := r.chunkByAlias[alias]; ref != nil {
|
|
return ref.ChunkID
|
|
}
|
|
if real := r.aliasToDoc[alias]; real != "" {
|
|
return real
|
|
}
|
|
if real := r.aliasToKB[alias]; real != "" {
|
|
return real
|
|
}
|
|
if ref := r.webByAlias[alias]; ref != nil {
|
|
return ref.URL
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// CompactKnownText is intentionally limited to identifiers already registered
|
|
// from structured runtime/tool data. It is used for metadata envelopes, not
|
|
// arbitrary retrieved prose.
|
|
func (r *Registry) CompactKnownText(text string) string {
|
|
if r == nil || text == "" {
|
|
return text
|
|
}
|
|
type pair struct{ real, alias string }
|
|
r.mu.RLock()
|
|
pairs := make([]pair, 0, len(r.chunkByID)+len(r.docToAlias)+len(r.kbToAlias)+len(r.webByURL))
|
|
for real, ref := range r.chunkByID {
|
|
pairs = append(pairs, pair{real, ref.Alias})
|
|
}
|
|
for real, alias := range r.docToAlias {
|
|
pairs = append(pairs, pair{real, alias})
|
|
}
|
|
for real, alias := range r.kbToAlias {
|
|
pairs = append(pairs, pair{real, alias})
|
|
}
|
|
for _, ref := range r.webByURL {
|
|
pairs = append(pairs, pair{ref.URL, ref.Alias})
|
|
}
|
|
r.mu.RUnlock()
|
|
sort.SliceStable(pairs, func(i, j int) bool { return len(pairs[i].real) > len(pairs[j].real) })
|
|
for _, item := range pairs {
|
|
if item.real != "" {
|
|
text = replaceIDExceptWikiSummarySlug(text, item.real, item.alias)
|
|
}
|
|
}
|
|
return text
|
|
}
|
|
|
|
// wikiSummarySlugPrefix guards wiki summary-page slugs (summary/<knowledgeID>)
|
|
// from citation compaction. A summary slug embeds a document's knowledge ID
|
|
// verbatim, so a blind ReplaceAll of that ID into its citation alias (d1) would
|
|
// mangle the slug into summary/d1 — an unrecoverable dead link the model then
|
|
// copies into wiki tool calls. Only the slug-embedded occurrence is preserved;
|
|
// the same document ID elsewhere in the text still compacts normally.
|
|
const wikiSummarySlugPrefix = "summary/"
|
|
|
|
// replaceIDExceptWikiSummarySlug replaces every occurrence of real with alias,
|
|
// except one that is immediately preceded by "summary/" (i.e. it is the
|
|
// identifier segment of a wiki summary-page slug). Go's regexp package has no
|
|
// lookbehind, so the scan is done manually.
|
|
func replaceIDExceptWikiSummarySlug(text, real, alias string) string {
|
|
if real == "" || !strings.Contains(text, real) {
|
|
return text
|
|
}
|
|
var b strings.Builder
|
|
b.Grow(len(text))
|
|
for {
|
|
i := strings.Index(text, real)
|
|
if i < 0 {
|
|
b.WriteString(text)
|
|
break
|
|
}
|
|
b.WriteString(text[:i])
|
|
if strings.HasSuffix(text[:i], wikiSummarySlugPrefix) {
|
|
b.WriteString(real) // part of a summary slug — keep intact
|
|
} else {
|
|
b.WriteString(alias)
|
|
}
|
|
text = text[i+len(real):]
|
|
}
|
|
return b.String()
|
|
}
|
|
|
|
var (
|
|
refTagRE = regexp.MustCompile(`(?i)<ref\s+id\s*=\s*"([^"]+)"\s*/?>`)
|
|
refCandidateRE = regexp.MustCompile(`(?is)<ref(?:\s|$)[^>]*(?:>|$)`)
|
|
modelKBTagRE = regexp.MustCompile(`(?is)<kb(?:\s|$)[^>]*(?:>|$)`)
|
|
modelWebTagRE = regexp.MustCompile(`(?is)<web(?:\s|$)[^>]*(?:>|$)`)
|
|
)
|
|
|
|
var (
|
|
documentAttrRE = regexp.MustCompile(`(?i)\bknowledge_id\s*=\s*"([^"]+)"`)
|
|
documentElementRE = regexp.MustCompile(`(?is)<knowledge_id>\s*([^<]+?)\s*</knowledge_id>`)
|
|
kbAttrRE = regexp.MustCompile(`(?i)\b(?:knowledge_base_id|kb_id)\s*=\s*"([^"]+)"`)
|
|
kbElementRE = regexp.MustCompile(`(?is)<(?:knowledge_base_id|kb_id)>\s*([^<]+?)\s*</(?:knowledge_base_id|kb_id)>`)
|
|
)
|
|
|
|
// registerLabeledReferences covers metadata-oriented tools that do not have a
|
|
// dedicated compact renderer. Only explicit ID labels are recognized; UUID-like
|
|
// text in retrieved content is never guessed to be a source identifier.
|
|
func (r *Registry) registerLabeledReferences(text string) {
|
|
if r == nil || text == "" {
|
|
return
|
|
}
|
|
for _, expression := range []*regexp.Regexp{documentAttrRE, documentElementRE} {
|
|
for _, match := range expression.FindAllStringSubmatch(text, -1) {
|
|
if len(match) == 2 {
|
|
r.RegisterDocument(strings.TrimSpace(match[1]))
|
|
}
|
|
}
|
|
}
|
|
for _, expression := range []*regexp.Regexp{kbAttrRE, kbElementRE} {
|
|
for _, match := range expression.FindAllStringSubmatch(text, -1) {
|
|
if len(match) == 2 {
|
|
r.RegisterKnowledgeBase(strings.TrimSpace(match[1]))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ExpandText converts the private model protocol into the existing public
|
|
// <kb/> / <web/> contract. Unknown aliases fail closed and disappear.
|
|
func (r *Registry) ExpandText(text string) string {
|
|
if r == nil || text == "" {
|
|
return text
|
|
}
|
|
// Public citation tags are output-only. Drop any instance written directly
|
|
// by the model, then create canonical tags solely from registered aliases.
|
|
text = modelKBTagRE.ReplaceAllString(text, "")
|
|
text = modelWebTagRE.ReplaceAllString(text, "")
|
|
if !r.citationsEnabled {
|
|
return refCandidateRE.ReplaceAllString(text, "")
|
|
}
|
|
return refCandidateRE.ReplaceAllStringFunc(text, func(tag string) string {
|
|
match := refTagRE.FindStringSubmatch(tag)
|
|
if len(match) != 2 {
|
|
return ""
|
|
}
|
|
alias := strings.ToLower(match[1])
|
|
r.mu.RLock()
|
|
chunkRef := r.chunkByAlias[alias]
|
|
webRef := r.webByAlias[alias]
|
|
r.mu.RUnlock()
|
|
if chunkRef != nil {
|
|
attrs := fmt.Sprintf(`doc="%s" chunk_id="%s"`, escapeAttr(chunkRef.DocumentTitle), escapeAttr(chunkRef.ChunkID))
|
|
if chunkRef.KnowledgeBaseID != "" {
|
|
attrs += fmt.Sprintf(` kb_id="%s"`, escapeAttr(chunkRef.KnowledgeBaseID))
|
|
}
|
|
return "<kb " + attrs + " />"
|
|
}
|
|
if webRef != nil {
|
|
return fmt.Sprintf(`<web url="%s" title="%s" />`, escapeAttr(webRef.URL), escapeAttr(webRef.Title))
|
|
}
|
|
return ""
|
|
})
|
|
}
|
|
|
|
func escapeAttr(value string) string { return html.EscapeString(value) }
|
|
|
|
func (r *Registry) ExpandResponse(response *types.ChatResponse) {
|
|
if response == nil {
|
|
return
|
|
}
|
|
response.Content = r.ExpandText(response.Content)
|
|
response.ReasoningContent = r.ExpandText(response.ReasoningContent)
|
|
r.DecodeToolCalls(response.ToolCalls)
|
|
}
|
|
|
|
// StreamExpander prevents partial private <ref/> tags from reaching SSE while
|
|
// preserving normal streaming for all other content.
|
|
type StreamExpander struct {
|
|
registry *Registry
|
|
pending string
|
|
}
|
|
|
|
func NewStreamExpander(registry *Registry) *StreamExpander {
|
|
return &StreamExpander{registry: registry}
|
|
}
|
|
|
|
func (d *StreamExpander) Feed(chunk string) string {
|
|
if d == nil || d.registry == nil {
|
|
return chunk
|
|
}
|
|
data := d.pending + chunk
|
|
d.pending = ""
|
|
var out strings.Builder
|
|
for data != "" {
|
|
idx := strings.Index(data, "<")
|
|
if idx < 0 {
|
|
out.WriteString(data)
|
|
break
|
|
}
|
|
out.WriteString(data[:idx])
|
|
data = data[idx:]
|
|
lower := strings.ToLower(data)
|
|
if isSourceTagPending(lower) && !strings.Contains(data, ">") {
|
|
d.pending = data
|
|
break
|
|
}
|
|
if isRefTagStart(lower) {
|
|
end := strings.IndexByte(data, '>')
|
|
if end < 0 {
|
|
d.pending = data
|
|
break
|
|
}
|
|
tag := data[:end+1]
|
|
if refTagRE.MatchString(tag) {
|
|
out.WriteString(d.registry.ExpandText(tag))
|
|
}
|
|
data = data[end+1:]
|
|
continue
|
|
}
|
|
if isNamedTagStart(lower, "kb") || isNamedTagStart(lower, "web") {
|
|
end := strings.IndexByte(data, '>')
|
|
if end < 0 {
|
|
d.pending = data
|
|
break
|
|
}
|
|
data = data[end+1:]
|
|
continue
|
|
}
|
|
out.WriteByte('<')
|
|
data = data[1:]
|
|
}
|
|
return out.String()
|
|
}
|
|
|
|
func isRefTagStart(value string) bool {
|
|
return isNamedTagStart(value, "ref")
|
|
}
|
|
|
|
func isNamedTagStart(value, name string) bool {
|
|
prefix := "<" + name
|
|
if !strings.HasPrefix(value, prefix) {
|
|
return false
|
|
}
|
|
if len(value) == len(prefix) {
|
|
return true
|
|
}
|
|
next := value[len(prefix)]
|
|
return next == ' ' || next == '\t' || next == '\r' || next == '\n' || next == '>'
|
|
}
|
|
|
|
func isSourceTagPending(value string) bool {
|
|
for _, name := range []string{"ref", "kb", "web"} {
|
|
prefix := "<" + name
|
|
if (len(value) <= len(prefix) && strings.HasPrefix(prefix, value)) || isNamedTagStart(value, name) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (d *StreamExpander) Flush() string {
|
|
if d == nil {
|
|
return ""
|
|
}
|
|
pending := d.pending
|
|
d.pending = ""
|
|
lower := strings.ToLower(pending)
|
|
if isSourceTagPending(lower) {
|
|
return ""
|
|
}
|
|
return d.registry.ExpandText(pending)
|
|
}
|