6714ef9813
- Updated the CitationGenerator to produce simple integer IDs instead of formatted strings, improving clarity and consistency in citation references. - Enhanced the executeAutoSearch method to save both successful and failed search results, capturing detailed execution data for better traceability. - Introduced a new SearchExecutionResult type to structure search result data, including query, keywords, configuration, duration, and error information. - Updated related tests to reflect changes in citation ID format and ensure proper functionality of the new storage mechanisms. - Revised documentation to clarify the new citation format and search result handling processes.
102 lines
2.7 KiB
Go
102 lines
2.7 KiB
Go
package search
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/yaoapp/yao/agent/search/types"
|
|
)
|
|
|
|
// DefaultCitationPrompt is the default prompt for citation instructions
|
|
const DefaultCitationPrompt = `You have access to reference data in <references> tags. Each <ref> has:
|
|
- id: Citation identifier (integer)
|
|
- type: Data type (web/kb/db)
|
|
- weight: Relevance weight (1.0=highest priority, 0.6=lowest)
|
|
- source: Origin (user=user-provided, hook=assistant-searched, auto=auto-searched)
|
|
|
|
Prioritize higher-weight references when answering.
|
|
|
|
When citing a reference, use this exact HTML format:
|
|
<a class="ref" data-ref-id="{id}" data-ref-type="{type}" href="#ref:{id}">[{id}]</a>
|
|
|
|
Example: According to the product data<a class="ref" data-ref-id="1" data-ref-type="db" href="#ref:1">[1]</a>, the price is $999.`
|
|
|
|
// BuildReferences converts search results to unified Reference format
|
|
func BuildReferences(results []*types.Result) []*types.Reference {
|
|
var refs []*types.Reference
|
|
for _, result := range results {
|
|
if result == nil {
|
|
continue
|
|
}
|
|
for _, item := range result.Items {
|
|
if item == nil {
|
|
continue
|
|
}
|
|
refs = append(refs, &types.Reference{
|
|
ID: item.CitationID,
|
|
Type: item.Type,
|
|
Source: item.Source,
|
|
Weight: item.Weight,
|
|
Score: item.Score,
|
|
Title: item.Title,
|
|
Content: item.Content,
|
|
URL: item.URL,
|
|
})
|
|
}
|
|
}
|
|
return refs
|
|
}
|
|
|
|
// FormatReferencesXML formats references as XML for LLM context
|
|
func FormatReferencesXML(refs []*types.Reference) string {
|
|
if len(refs) == 0 {
|
|
return ""
|
|
}
|
|
|
|
var sb strings.Builder
|
|
sb.WriteString("<references>\n")
|
|
|
|
for _, ref := range refs {
|
|
if ref == nil {
|
|
continue
|
|
}
|
|
sb.WriteString(fmt.Sprintf(`<ref id="%s" type="%s" weight="%.1f" source="%s">`,
|
|
ref.ID, ref.Type, ref.Weight, ref.Source))
|
|
sb.WriteString("\n")
|
|
|
|
if ref.Title != "" {
|
|
sb.WriteString(ref.Title)
|
|
sb.WriteString("\n")
|
|
}
|
|
sb.WriteString(ref.Content)
|
|
if ref.URL != "" {
|
|
sb.WriteString("\nURL: ")
|
|
sb.WriteString(ref.URL)
|
|
}
|
|
sb.WriteString("\n</ref>\n")
|
|
}
|
|
|
|
sb.WriteString("</references>")
|
|
return sb.String()
|
|
}
|
|
|
|
// GetCitationPrompt returns the citation instruction prompt
|
|
func GetCitationPrompt(cfg *types.CitationConfig) string {
|
|
if cfg == nil {
|
|
return DefaultCitationPrompt
|
|
}
|
|
if cfg.CustomPrompt != "" {
|
|
return cfg.CustomPrompt
|
|
}
|
|
return DefaultCitationPrompt
|
|
}
|
|
|
|
// BuildReferenceContext builds the complete reference context for LLM
|
|
func BuildReferenceContext(results []*types.Result, cfg *types.CitationConfig) *types.ReferenceContext {
|
|
refs := BuildReferences(results)
|
|
return &types.ReferenceContext{
|
|
References: refs,
|
|
XML: FormatReferencesXML(refs),
|
|
Prompt: GetCitationPrompt(cfg),
|
|
}
|
|
}
|