opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance Assistant Search Functionality and Cache Management

Browse source 
- 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
Show stats Download patch file Download diff file Expand all files Collapse all files

2
agent/assistant/agent.go
View file

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

29
agent/assistant/cache.go
View file

@ -124,6 +124,35 @@ func (c *Cache) Clear() {
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
func (c *Cache) removeOldest() {
if element := c.list.Back(); element != nil {

6
agent/assistant/load.go
View file

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

3
agent/assistant/load_system.go
View file

@ -146,9 +146,8 @@ func loadSystemAgent(id, pathPrefix string) (*Assistant, error) {
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["path"] = "/" + pathPrefix
// Set type if not specified
if _, has := pkgData["type"]; !has {

28
agent/assistant/load_test.go
View file

@ -544,6 +544,34 @@ func TestLoadSystemAgents(t *testing.T) {
}
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

195
agent/assistant/search.go
View file

@ -1,6 +1,7 @@
package assistant
import (
"encoding/json"
"fmt"
"strings"
"time"
@ -17,9 +18,17 @@ import (
// shouldAutoSearch determines if auto search should be executed
// Returns false if:
// - opts.Skip.Search is true
// - uses.search is "disabled"
// - 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
uses := ast.getMergedSearchUses(createResponse)
@ -34,10 +43,194 @@ func (ast *Assistant) shouldAutoSearch(ctx *context.Context, createResponse *con
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)
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
// Priority: createResponse > assistant
func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResponse) *context.Uses {

1
agent/context/types.go
View file

@ -196,6 +196,7 @@ type Skip struct {
Trace bool `json:"trace"` // Skip trace logging
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)
Search bool `json:"search"` // Skip auto search (for internal calls like needsearch intent detection)
}
// MessageMetadata stores metadata for sent messages

15
agent/i18n/builtin.go
View file

@ -107,6 +107,11 @@ func init() {
"search.failed": "Search failed",
"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.trace.label": "Search",
"search.trace.description": "Search the web and knowledge base for relevant information",
@ -191,6 +196,11 @@ func init() {
"search.failed": "搜索失败",
"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.trace.label": "搜索",
"search.trace.description": "搜索网络和知识库获取相关信息",
@ -303,6 +313,11 @@ func init() {
"search.failed": "搜索失败",
"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.trace.label": "搜索",
"search.trace.description": "搜索网络和知识库获取相关信息",

372
data/bindata.go
View file

File diff suppressed because it is too large Load diff

1
yao/assistants/keyword/package.yao
View file

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

29
yao/assistants/keyword/prompts.yml
View file

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

113
yao/assistants/keyword/src/index.ts Normal file
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)];
}

117
yao/assistants/needsearch/src/index.ts Normal file
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
};
}

1
yao/assistants/prompt/package.yao
View file

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

1
yao/assistants/querydsl/package.yao
View file

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

1
yao/assistants/title/package.yao
View file

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

25
yao/assistants/title/prompts.yml
View file

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