Enhance Assistant Search Functionality and Cache Management

- Updated the `shouldAutoSearch` method to include additional parameters for improved intent detection, allowing for better decision-making on whether to execute auto search.
- Introduced a new `checkSearchIntent` method to utilize the `__yao.needsearch` agent for determining the necessity of a search based on user input.
- Implemented a `ClearExcept` method in the cache to selectively clear non-system agents while preserving essential system agents during cache management.
- Updated the `LoadBuiltIn` function to maintain system agents in the cache, ensuring they remain available for use.
- Enhanced test coverage for loading system agents and validating search intent detection, ensuring robustness in the assistant's search capabilities.
- Revised localization files to include new messages for search intent feedback, improving user experience during search operations.
This commit is contained in:
Max 2025-12-16 16:44:35 +08:00
parent c86ccb55b1
commit 750dd311b9
17 changed files with 749 additions and 190 deletions

View file

@ -202,7 +202,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// ================================================ // ================================================
// Execute Auto Search (if enabled) // Execute Auto Search (if enabled)
// ================================================ // ================================================
if ast.shouldAutoSearch(ctx, createResponse) { if ast.shouldAutoSearch(ctx, completionMessages, createResponse, opts) {
refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, opts) refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse, opts)
if refCtx != nil && len(refCtx.References) > 0 { if refCtx != nil && len(refCtx.References) > 0 {
completionMessages = ast.injectSearchContext(completionMessages, refCtx) completionMessages = ast.injectSearchContext(completionMessages, refCtx)

View file

@ -124,6 +124,35 @@ func (c *Cache) Clear() {
c.items = make(map[string]*list.Element) c.items = make(map[string]*list.Element)
} }
// ClearExcept removes items from the cache except those matching the keep function
// keep function returns true for items that should be preserved
func (c *Cache) ClearExcept(keep func(id string) bool) {
c.mu.Lock()
defer c.mu.Unlock()
// Collect items to remove
var toRemove []*list.Element
for element := c.list.Front(); element != nil; element = element.Next() {
item := element.Value.(*cacheItem)
if !keep(item.key) {
toRemove = append(toRemove, element)
}
}
// Remove collected items
for _, element := range toRemove {
item := element.Value.(*cacheItem)
// Unregister scripts before removing
if item.value != nil && len(item.value.Scripts) > 0 {
item.value.UnregisterScripts()
}
c.list.Remove(element)
delete(c.items, item.key)
}
}
// removeOldest removes the least recently used item from the cache // removeOldest removes the least recently used item from the cache
func (c *Cache) removeOldest() { func (c *Cache) removeOldest() {
if element := c.list.Back(); element != nil { if element := c.list.Back(); element != nil {

View file

@ -33,8 +33,10 @@ var globalSearchConfig *searchTypes.Config = nil // global search config from ag
// LoadBuiltIn load the built-in assistants // LoadBuiltIn load the built-in assistants
func LoadBuiltIn() error { func LoadBuiltIn() error {
// Clear the cache // Clear non-system agents from cache (preserve system agents loaded by LoadSystemAgents)
loaded.Clear() loaded.ClearExcept(func(id string) bool {
return strings.HasPrefix(id, "__yao.") // Keep system agents
})
root := `/assistants` root := `/assistants`
app, err := fs.Get("app") app, err := fs.Get("app")

View file

@ -146,9 +146,8 @@ func loadSystemAgent(id, pathPrefix string) (*Assistant, error) {
return nil, fmt.Errorf("failed to parse %s: %w", pkgPath, err) return nil, fmt.Errorf("failed to parse %s: %w", pkgPath, err)
} }
// Set assistant_id and path // Set assistant_id (no path - system agents are loaded from storage, not filesystem)
pkgData["assistant_id"] = id pkgData["assistant_id"] = id
pkgData["path"] = "/" + pathPrefix
// Set type if not specified // Set type if not specified
if _, has := pkgData["type"]; !has { if _, has := pkgData["type"]; !has {

View file

@ -544,6 +544,34 @@ func TestLoadSystemAgents(t *testing.T) {
} }
assert.True(t, found, "System agents should be found in storage") assert.True(t, found, "System agents should be found in storage")
}) })
t.Run("SystemAgentsGetFromStorage", func(t *testing.T) {
// Clear cache to force loading from storage
assistant.GetCache().Clear()
// Test Get for each system agent
systemAgents := []string{
"__yao.keyword",
"__yao.querydsl",
"__yao.title",
"__yao.prompt",
"__yao.needsearch",
"__yao.entity",
}
for _, agentID := range systemAgents {
ast, err := assistant.Get(agentID)
require.NoError(t, err, "Get(%s) should succeed", agentID)
require.NotNil(t, ast, "Get(%s) should return assistant", agentID)
assert.Equal(t, agentID, ast.ID)
assert.True(t, ast.BuiltIn, "%s should be built-in", agentID)
assert.True(t, ast.Readonly, "%s should be readonly", agentID)
assert.Contains(t, ast.Tags, "system", "%s should have system tag", agentID)
assert.Equal(t, "worker", ast.Type, "%s should be worker type", agentID)
assert.NotNil(t, ast.Prompts, "%s should have prompts", agentID)
assert.Greater(t, len(ast.Prompts), 0, "%s should have at least one prompt", agentID)
}
})
} }
// TestValidate tests the assistant Validate method // TestValidate tests the assistant Validate method

View file

@ -1,6 +1,7 @@
package assistant package assistant
import ( import (
"encoding/json"
"fmt" "fmt"
"strings" "strings"
"time" "time"
@ -17,9 +18,17 @@ import (
// shouldAutoSearch determines if auto search should be executed // shouldAutoSearch determines if auto search should be executed
// Returns false if: // Returns false if:
// - opts.Skip.Search is true
// - uses.search is "disabled" // - uses.search is "disabled"
// - assistant has no search configuration // - assistant has no search configuration
func (ast *Assistant) shouldAutoSearch(ctx *context.Context, createResponse *context.HookCreateResponse) bool { // - needsearch intent detection returns false
func (ast *Assistant) shouldAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, opts *context.Options) bool {
// Check if search is skipped via options
if opts != nil && opts.Skip != nil && opts.Skip.Search {
ctx.Logger.Debug("Auto search skipped by opts.Skip.Search")
return false
}
// Get merged uses configuration // Get merged uses configuration
uses := ast.getMergedSearchUses(createResponse) uses := ast.getMergedSearchUses(createResponse)
@ -34,10 +43,194 @@ func (ast *Assistant) shouldAutoSearch(ctx *context.Context, createResponse *con
return false return false
} }
// Check search intent using __yao.needsearch agent
if !ast.checkSearchIntent(ctx, messages) {
ctx.Logger.Info("Auto search skipped: intent detection returned false")
return false
}
// Check if search is enabled (builtin, agent, mcp, or empty means builtin) // Check if search is enabled (builtin, agent, mcp, or empty means builtin)
return true return true
} }
// checkSearchIntent uses __yao.needsearch agent to determine if search is needed
// Returns true if search is needed, false otherwise
func (ast *Assistant) checkSearchIntent(ctx *context.Context, messages []context.Message) bool {
// Get the last user message
var userQuery string
for i := len(messages) - 1; i >= 0; i-- {
if messages[i].Role == "user" {
if content, ok := messages[i].Content.(string); ok {
userQuery = content
break
}
}
}
if userQuery == "" {
return true // No user message, proceed with search
}
// Try to get __yao.needsearch agent
needsearchAst, err := Get("__yao.needsearch")
if err != nil {
ctx.Logger.Debug("__yao.needsearch agent not available: %v, proceeding with search", err)
return true // Agent not available, proceed with search
}
// === Output: Send loading message ===
loadingID := ast.sendIntentLoading(ctx)
// Build messages for intent detection
intentMessages := []context.Message{
{Role: "user", Content: userQuery},
}
// Call the needsearch agent (Stack will auto-track)
// IMPORTANT: Skip search to prevent infinite loop, skip output to prevent JSON showing in UI
opts := &context.Options{
Skip: &context.Skip{
History: true, // Don't save to history
Search: true, // Skip search to prevent infinite loop
Output: true, // Skip output to prevent JSON showing in UI
},
}
result, err := needsearchAst.Stream(ctx, intentMessages, opts)
if err != nil {
ctx.Logger.Debug("__yao.needsearch failed: %v, proceeding with search", err)
// === Output: Send done (error case, proceed with search) ===
ast.sendIntentDone(ctx, loadingID, true, "")
return true // On error, proceed with search
}
// Parse the result
// Next hook returns {data: {need_search: bool, search_types: [], confidence: float}}
if response, ok := result.(*context.Response); ok {
// First try to get from Next hook response
if response.Next != nil {
if nextData, ok := response.Next.(map[string]interface{}); ok {
// Check for data field (from Next hook's {data: result})
var intentData map[string]interface{}
if data, ok := nextData["data"].(map[string]interface{}); ok {
intentData = data
} else {
intentData = nextData
}
if needSearch, ok := intentData["need_search"].(bool); ok {
reason, _ := intentData["reason"].(string)
ctx.Logger.Debug("Search intent (from Next): need_search=%v, reason=%s", needSearch, reason)
ast.sendIntentDone(ctx, loadingID, needSearch, reason)
return needSearch
}
}
}
// Fallback: parse from Completion.Content if Next hook didn't process
if response.Completion != nil {
content, ok := response.Completion.Content.(string)
if !ok || content == "" {
ast.sendIntentDone(ctx, loadingID, true, "")
return true
}
needSearch, reason := parseNeedSearchFromContent(content)
ctx.Logger.Debug("Search intent (from Content): need_search=%v, reason=%s", needSearch, reason)
ast.sendIntentDone(ctx, loadingID, needSearch, reason)
return needSearch
}
}
// Default: proceed with search if we can't parse the result
// === Output: Send done (default case) ===
ast.sendIntentDone(ctx, loadingID, true, "")
return true
}
// parseNeedSearchFromContent parses need_search result from LLM completion content
// Handles JSON wrapped in markdown code blocks
func parseNeedSearchFromContent(content string) (bool, string) {
// Remove markdown code block if present
content = strings.TrimSpace(content)
if strings.HasPrefix(content, "```json") {
content = strings.TrimPrefix(content, "```json")
content = strings.TrimSuffix(content, "```")
content = strings.TrimSpace(content)
} else if strings.HasPrefix(content, "```") {
content = strings.TrimPrefix(content, "```")
content = strings.TrimSuffix(content, "```")
content = strings.TrimSpace(content)
}
// Try to parse JSON
var result map[string]interface{}
if err := json.Unmarshal([]byte(content), &result); err != nil {
// Failed to parse, default to search
return true, ""
}
needSearch, ok := result["need_search"].(bool)
if !ok {
return true, ""
}
reason, _ := result["reason"].(string)
return needSearch, reason
}
// sendIntentLoading sends the initial intent detection loading message
// Returns the message ID for later replacement
func (ast *Assistant) sendIntentLoading(ctx *context.Context) string {
loadingMsg := i18n.T(ctx.Locale, "search.intent.loading")
msg := &message.Message{
Type: "loading",
Props: map[string]any{
"message": loadingMsg,
},
}
// Send and get message ID
msgID, err := ctx.SendStream(msg)
if err != nil {
ctx.Logger.Warn("Failed to send intent loading message: %v", err)
return ""
}
return msgID
}
// sendIntentDone replaces loading with result
// Only marks as done when needSearch is false (no further loading will follow)
// When needSearch is true, the search loading will continue
func (ast *Assistant) sendIntentDone(ctx *context.Context, loadingID string, needSearch bool, reason string) {
if loadingID == "" {
return
}
var resultMsg string
if needSearch {
resultMsg = i18n.T(ctx.Locale, "search.intent.need_search")
} else {
resultMsg = i18n.T(ctx.Locale, "search.intent.no_search")
}
msg := &message.Message{
MessageID: loadingID,
Delta: true,
DeltaAction: message.DeltaReplace,
Type: "loading",
Props: map[string]any{
"message": resultMsg,
"done": true, // Intent detection loading is independent, always close it
},
}
if err := ctx.Send(msg); err != nil {
ctx.Logger.Warn("Failed to send intent done message: %v", err)
}
}
// getMergedSearchUses returns the merged uses configuration for search // getMergedSearchUses returns the merged uses configuration for search
// Priority: createResponse > assistant // Priority: createResponse > assistant
func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResponse) *context.Uses { func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResponse) *context.Uses {

View file

@ -196,6 +196,7 @@ type Skip struct {
Trace bool `json:"trace"` // Skip trace logging Trace bool `json:"trace"` // Skip trace logging
Output bool `json:"output"` // Skip output to client (for internal A2A calls that only need response data) Output bool `json:"output"` // Skip output to client (for internal A2A calls that only need response data)
Keyword bool `json:"keyword"` // Skip keyword extraction for web search (use raw query directly) Keyword bool `json:"keyword"` // Skip keyword extraction for web search (use raw query directly)
Search bool `json:"search"` // Skip auto search (for internal calls like needsearch intent detection)
} }
// MessageMetadata stores metadata for sent messages // MessageMetadata stores metadata for sent messages

View file

@ -107,6 +107,11 @@ func init() {
"search.failed": "Search failed", "search.failed": "Search failed",
"search.no_results": "No references found", "search.no_results": "No references found",
// Search Intent: assistant/search.go - Intent detection messages
"search.intent.loading": "Checking if references are needed...",
"search.intent.need_search": "Searching for references...",
"search.intent.no_search": "No references needed",
// Search: assistant/search.go - Trace labels // Search: assistant/search.go - Trace labels
"search.trace.label": "Search", "search.trace.label": "Search",
"search.trace.description": "Search the web and knowledge base for relevant information", "search.trace.description": "Search the web and knowledge base for relevant information",
@ -191,6 +196,11 @@ func init() {
"search.failed": "搜索失败", "search.failed": "搜索失败",
"search.no_results": "未找到相关资料", "search.no_results": "未找到相关资料",
// Search Intent: assistant/search.go - Intent detection messages
"search.intent.loading": "检查是否需要查询资料...",
"search.intent.need_search": "正在查询相关资料...",
"search.intent.no_search": "无需查询资料",
// Search: assistant/search.go - Trace labels // Search: assistant/search.go - Trace labels
"search.trace.label": "搜索", "search.trace.label": "搜索",
"search.trace.description": "搜索网络和知识库获取相关信息", "search.trace.description": "搜索网络和知识库获取相关信息",
@ -303,6 +313,11 @@ func init() {
"search.failed": "搜索失败", "search.failed": "搜索失败",
"search.no_results": "未找到相关资料", "search.no_results": "未找到相关资料",
// Search Intent: assistant/search.go - Intent detection messages
"search.intent.loading": "检查是否需要查询资料...",
"search.intent.need_search": "正在查询相关资料...",
"search.intent.no_search": "无需查询资料",
// Search: assistant/search.go - Trace labels // Search: assistant/search.go - Trace labels
"search.trace.label": "搜索", "search.trace.label": "搜索",
"search.trace.description": "搜索网络和知识库获取相关信息", "search.trace.description": "搜索网络和知识库获取相关信息",

File diff suppressed because it is too large Load diff

View file

@ -2,6 +2,7 @@
"name": "Keyword Extraction", "name": "Keyword Extraction",
"description": "Extract keywords from text content", "description": "Extract keywords from text content",
"type": "worker", "type": "worker",
"uses": { "search": "disabled" },
"options": { "options": {
"max_tokens": 500, "max_tokens": 500,
"temperature": 0.3 "temperature": 0.3

View file

@ -1,24 +1,21 @@
# Keyword Extraction Agent Prompts
- role: system - role: system
content: | content: |
You are a keyword extraction specialist. Your task is to extract relevant keywords from the provided text. Extract keywords from text content.
## Instructions Task:
1. Analyze the input text carefully 1. Analyze input text
2. Extract the most important and relevant keywords 2. Extract important keywords
3. Return keywords in JSON format 3. Return JSON format
4. Match input language
## Response Format Response Format (JSON only):
Always respond with valid JSON:
```json ```json
{ {"keywords": ["keyword1", "keyword2", ...]}
"keywords": ["keyword1", "keyword2", ...]
}
``` ```
## Guidelines Guidelines:
- Extract 5-15 keywords depending on content length - Extract 5-15 keywords based on content length
- Prioritize nouns, proper nouns, and key concepts - Prioritize nouns, proper nouns, key concepts
- Include both single words and short phrases when relevant - Include single words and short phrases
- Exclude common stop words - Exclude common stop words
- Maintain the original language of the content - Keywords MUST be in the same language as input

View file

@ -0,0 +1,113 @@
/**
* Keyword Extraction Agent - Next Hook
* Parses LLM response and extracts keywords with error tolerance
*/
// @ts-nocheck
/**
* Next hook - processes keyword extraction response
* Uses json.Parse for fault-tolerant JSON parsing
*/
function Next(
ctx: agent.Context,
payload: agent.NextHookPayload
): agent.NextHookResponse | null {
const completion = payload.completion;
// No completion, return null for standard handling
if (!completion || !completion.content) {
return null;
}
// Remove markdown code block if present
let content = completion.content.trim();
if (content.startsWith("```json")) {
content = content.slice(7);
} else if (content.startsWith("```")) {
content = content.slice(3);
}
if (content.endsWith("```")) {
content = content.slice(0, -3);
}
content = content.trim();
// Try to parse JSON from completion content
let keywords: string[] = [];
try {
// Use json.Parse for fault-tolerant parsing (handles broken JSON, JSONC, etc.)
const parsed = Process("json.Parse", content) as {
keywords?: string[];
} | null;
if (parsed && Array.isArray(parsed.keywords)) {
keywords = parsed.keywords.filter(
(k) => typeof k === "string" && k.trim().length > 0
);
}
} catch (e) {
// If json.Parse fails, try to extract keywords from text
keywords = extractKeywordsFromText(content);
}
// If still no keywords, try extracting from raw text
if (keywords.length === 0) {
keywords = extractKeywordsFromText(content);
}
// Return parsed keywords
return {
data: {
keywords: keywords,
},
};
}
/**
* Extract keywords from plain text when JSON parsing fails
* Handles formats like:
* - Comma-separated: "keyword1, keyword2, keyword3"
* - Line-separated: "keyword1\nkeyword2\nkeyword3"
* - Bullet points: "- keyword1\n- keyword2"
* - Numbered: "1. keyword1\n2. keyword2"
*/
function extractKeywordsFromText(text: string): string[] {
const keywords: string[] = [];
// Remove common prefixes/suffixes
let cleaned = text
.replace(/^[\s\S]*?keywords?[\s::]*\[?/i, "") // Remove "keywords:" prefix
.replace(/\][\s\S]*$/, "") // Remove trailing ]
.trim();
// Try line-by-line extraction
const lines = cleaned.split(/[\n\r]+/);
for (const line of lines) {
// Remove bullet points, numbers, quotes
let keyword = line
.replace(/^[\s\-\*\•\d\.]+/, "") // Remove bullets/numbers
.replace(/^["'`]+|["'`]+$/g, "") // Remove quotes
.replace(/,\s*$/, "") // Remove trailing comma
.trim();
// Skip empty or too long
if (keyword.length > 0 && keyword.length < 100) {
// Split by comma if contains multiple
if (keyword.includes(",")) {
const parts = keyword.split(",").map((p) => p.trim());
for (const part of parts) {
if (part.length > 0 && part.length < 100) {
keywords.push(part);
}
}
} else {
keywords.push(keyword);
}
}
}
// Deduplicate
return [...new Set(keywords)];
}

View file

@ -0,0 +1,117 @@
/**
* Need Search Agent - Next Hook
* Parses LLM response and extracts search intent with error tolerance
*/
// @ts-nocheck
interface SearchResult {
need_search: boolean;
search_types: string[];
confidence: number;
}
/**
* Next hook - processes search intent response
* Uses json.Parse for fault-tolerant JSON parsing
*/
function Next(
ctx: agent.Context,
payload: agent.NextHookPayload
): agent.NextHookResponse | null {
const completion = payload.completion;
// No completion, return null for standard handling
if (!completion || !completion.content) {
return null;
}
// Remove markdown code block if present
let content = completion.content.trim();
if (content.startsWith("```json")) {
content = content.slice(7); // Remove ```json
} else if (content.startsWith("```")) {
content = content.slice(3); // Remove ```
}
if (content.endsWith("```")) {
content = content.slice(0, -3); // Remove trailing ```
}
content = content.trim();
// Default result
let result: SearchResult = {
need_search: false,
search_types: [],
confidence: 0,
};
try {
// Use json.Parse for fault-tolerant parsing
const parsed = Process("json.Parse", content) as {
need_search?: boolean;
search_types?: string[];
confidence?: number;
} | null;
if (parsed) {
result.need_search = Boolean(parsed.need_search);
result.search_types = Array.isArray(parsed.search_types)
? parsed.search_types.filter(
(t) =>
typeof t === "string" &&
["web", "kb", "db"].includes(t.toLowerCase())
)
: [];
result.confidence =
typeof parsed.confidence === "number"
? Math.min(1, Math.max(0, parsed.confidence))
: 0.5;
}
} catch (e) {
// If json.Parse fails, try to extract from text
result = extractFromText(content);
}
// Return parsed result
return {
data: result,
};
}
/**
* Extract search intent from plain text when JSON parsing fails
*/
function extractFromText(text: string): SearchResult {
const lower = text.toLowerCase();
// Check for explicit indicators
const needSearch =
lower.includes("true") ||
lower.includes("need") ||
lower.includes("search") ||
lower.includes("web") ||
lower.includes("kb") ||
lower.includes("db");
const noSearch =
lower.includes("false") ||
lower.includes("no search") ||
lower.includes("not need");
// Extract search types
const searchTypes: string[] = [];
if (lower.includes("web")) searchTypes.push("web");
if (lower.includes("kb") || lower.includes("knowledge"))
searchTypes.push("kb");
if (lower.includes("db") || lower.includes("database"))
searchTypes.push("db");
// Determine need_search
const need = noSearch ? false : needSearch && searchTypes.length > 0;
return {
need_search: need,
search_types: need ? searchTypes : [],
confidence: 0.5, // Low confidence for text extraction
};
}

View file

@ -2,6 +2,7 @@
"name": "Prompt Optimizer", "name": "Prompt Optimizer",
"description": "Transform user requirements into effective prompts", "description": "Transform user requirements into effective prompts",
"type": "worker", "type": "worker",
"uses": { "search": "disabled" },
"options": { "options": {
"temperature": 0 "temperature": 0
} }

View file

@ -2,6 +2,7 @@
"name": "QueryDSL Generator", "name": "QueryDSL Generator",
"description": "Generate QueryDSL from natural language", "description": "Generate QueryDSL from natural language",
"type": "worker", "type": "worker",
"uses": { "search": "disabled" },
"options": { "options": {
"max_tokens": 2000, "max_tokens": 2000,
"temperature": 0.2 "temperature": 0.2

View file

@ -2,6 +2,7 @@
"name": "Title Generator", "name": "Title Generator",
"description": "Generate concise titles for conversations", "description": "Generate concise titles for conversations",
"type": "worker", "type": "worker",
"uses": { "search": "disabled" },
"options": { "options": {
"temperature": 0 "temperature": 0
} }

View file

@ -5,12 +5,16 @@
Task: Task:
1. Analyze content and identify main topic 1. Analyze content and identify main topic
2. Create brief, descriptive title 2. Create brief, descriptive title
3. Match input language 3. Title MUST be in the same language as user input
4. Return ONLY the title, no explanation
Output:
- Return ONLY the plain text title
- NO markdown, NO code blocks, NO quotes, NO explanation
- Just the title text itself
Length: Length:
- English: 2-6 words, 15-50 chars - English: 2-6 words, 15-50 chars
- CJK: 2-10 chars - CJK (Chinese/Japanese/Korean): 2-10 chars
- Mixed: max 50 chars - Mixed: max 50 chars
Style: Style:
@ -20,7 +24,14 @@
- Sentence case for English - Sentence case for English
Examples: Examples:
"How to bake cookies?" → Chocolate Chip Cookie Recipe Input: "How to bake cookies?"
"请教如何制作曲奇" → 巧克力曲奇制作 Output: Chocolate Chip Cookie Recipe
"Debug my React component" → React Component Debugging
"帮我调试React组件" → React组件调试 Input: "请教如何制作曲奇"
Output: 巧克力曲奇制作
Input: "Debug my React component"
Output: React Component Debugging
Input: "帮我调试React组件"
Output: React组件调试