yao/agent/search/reference.go
Max d264c7a784 Implement search configuration management and testing enhancements
- Introduced a new search configuration structure, allowing for detailed settings for web, knowledge base, database, citation, and weights.
- Updated the `Assistant` model to include a `Search` field, enabling assistant-specific search configurations.
- Enhanced the loading and merging logic for search configurations, ensuring global defaults can be overridden by assistant-specific settings.
- Added comprehensive tests for loading, saving, and updating assistants with search configurations, verifying the integrity of search settings.
- Updated documentation in DESIGN.md to reflect the new search configuration hierarchy and usage, clarifying the interaction between global and assistant-level settings.
2025-12-13 11:19:20 +08:00

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
- 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="ref_001" data-ref-type="db" href="#ref:ref_001">[ref_001]</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),
}
}