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

Merge pull request #1389 from trheyi/main

Browse source 
Refactor KB Search Handler and Update Test Cases
This commit is contained in:
Max 2025-12-20 17:04:15 +08:00 • committed by GitHub
commit 212815ebf6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 1279 additions and 160 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
.github/workflows/pr-test.yml vendored
View file

@ -336,6 +336,11 @@ jobs:
- name: Run KB Tests (kb)
run: make unit-test-kb
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
- name: "Comment on PR - KB Tests Done"
uses: actions/github-script@v7
with:
@ -528,6 +533,11 @@ jobs:
- name: Run AI Tests (agent, aigc)
run: make unit-test-ai
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
- name: "Comment on PR - AI Tests Done"
uses: actions/github-script@v7
with:

10
.github/workflows/unit-test.yml vendored
View file

@ -295,6 +295,11 @@ jobs:
- name: Run KB Tests (kb)
run: make unit-test-kb
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
# =============================================================================
# AI Tests (agent, aigc) - Run once with SQLite
# =============================================================================
@ -425,6 +430,11 @@ jobs:
- name: Run AI Tests (agent, aigc)
run: make unit-test-ai
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
# =============================================================================
# Benchmark & Memory Leak Tests - Run with MySQL8.0 and SQLite3
# =============================================================================

8
Makefile
View file

@ -76,7 +76,7 @@ unit-test-core:
# AI Unit Test (agent, aigc)
.PHONY: unit-test-ai
unit-test-ai:
echo "mode: count" > coverage-ai.out
echo "mode: count" > coverage.out
for d in $(TESTFOLDER_AI); do \
$(GO) test -tags $(TESTTAGS) -v -timeout=20m -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal' $$d > tmp.out; \
cat tmp.out; \
@ -100,7 +100,7 @@ unit-test-ai:
exit 1; \
fi; \
if [ -f profile.out ]; then \
cat profile.out | grep -v "mode:" >> coverage-ai.out; \
cat profile.out | grep -v "mode:" >> coverage.out; \
rm profile.out; \
fi; \
done
@ -108,7 +108,7 @@ unit-test-ai:
# KB Unit Test (kb)
.PHONY: unit-test-kb
unit-test-kb:
echo "mode: count" > coverage-kb.out
echo "mode: count" > coverage.out
for d in $(TESTFOLDER_KB); do \
$(GO) test -tags $(TESTTAGS) -v -timeout=20m -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestSearchCleanup' $$d > tmp.out; \
cat tmp.out; \
@ -132,7 +132,7 @@ unit-test-kb:
exit 1; \
fi; \
if [ -f profile.out ]; then \
cat profile.out | grep -v "mode:" >> coverage-kb.out; \
cat profile.out | grep -v "mode:" >> coverage.out; \
rm profile.out; \
fi; \
done

29
agent/assistant/search.go
View file

@ -801,15 +801,25 @@ func (ast *Assistant) buildSearchRequests(ctx *context.Context, query string, co
threshold = config.KB.Threshold
}
}
requests = append(requests, &searchTypes.Request{
// Filter collections by authorization (Collection-level permission check)
allowedCollections := FilterKBCollectionsByAuth(ctx, ast.KB.Collections)
if len(allowedCollections) == 0 {
ctx.Logger.Info("No accessible KB collections after auth filter")
} else {
// Build KB request
kbReq := &searchTypes.Request{
Type: searchTypes.SearchTypeKB,
Query: query, // KB uses original query for semantic search
Source: searchTypes.SourceAuto,
Limit: limit,
Collections: ast.KB.Collections,
Collections: allowedCollections,
Threshold: threshold,
Graph: config != nil && config.KB != nil && config.KB.Graph,
})
}
requests = append(requests, kbReq)
}
}
// DB search - check if DB is configured and allowed by intent
@ -818,13 +828,22 @@ func (ast *Assistant) buildSearchRequests(ctx *context.Context, query string, co
if config != nil && config.DB != nil && config.DB.MaxResults > 0 {
limit = config.DB.MaxResults
}
requests = append(requests, &searchTypes.Request{
// Build DB request with auth where clauses
dbReq := &searchTypes.Request{
Type: searchTypes.SearchTypeDB,
Query: query, // DB uses original query for QueryDSL generation
Source: searchTypes.SourceAuto,
Limit: limit,
Models: ast.DB.Models,
})
}
// Apply authorization where clauses
if authWheres := BuildDBAuthWheres(ctx); authWheres != nil {
dbReq.Wheres = authWheres
}
requests = append(requests, dbReq)
}
return requests, extractedKeywords

103
agent/assistant/search_auth_db.go Normal file
View file

@ -0,0 +1,103 @@
package assistant
import (
"github.com/yaoapp/gou/query/gou"
"github.com/yaoapp/yao/agent/context"
)
// BuildDBAuthWheres builds where clauses for DB search based on authorization
// This applies permission-based filtering to database queries
// Returns gou.Where clauses to filter records by authorization scope
func BuildDBAuthWheres(ctx *context.Context) []gou.Where {
if ctx == nil || ctx.Authorized == nil {
return nil
}
authInfo := ctx.Authorized
// No constraints, no filter needed
if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
return nil
}
var wheres []gou.Where
// Team only - User can access:
// 1. Public records (public = true)
// 2. Records in their team where:
// - They created the record (__yao_created_by matches)
// - OR the record is shared with team (share = "team")
if authInfo.Constraints.TeamOnly && authInfo.TeamID != "" {
wheres = append(wheres, gou.Where{
Wheres: []gou.Where{
// Public records
{Condition: gou.Condition{
Field: &gou.Expression{Field: "public"},
Value: true,
OP: "=",
OR: true,
}},
// Team records
{
Wheres: []gou.Where{
{Condition: gou.Condition{
Field: &gou.Expression{Field: "__yao_team_id"},
Value: authInfo.TeamID,
OP: "=",
}},
{Wheres: []gou.Where{
{Condition: gou.Condition{
Field: &gou.Expression{Field: "__yao_created_by"},
Value: authInfo.UserID,
OP: "=",
}},
{Condition: gou.Condition{
Field: &gou.Expression{Field: "share"},
Value: "team",
OP: "=",
OR: true,
}},
}},
},
},
},
})
return wheres
}
// Owner only - User can access:
// 1. Public records (public = true)
// 2. Records they created where:
// - __yao_team_id is null (not team records)
// - __yao_created_by matches their user ID
if authInfo.Constraints.OwnerOnly && authInfo.UserID != "" {
wheres = append(wheres, gou.Where{
Wheres: []gou.Where{
// Public records
{Condition: gou.Condition{
Field: &gou.Expression{Field: "public"},
Value: true,
OP: "=",
OR: true,
}},
// Owner records
{
Wheres: []gou.Where{
{Condition: gou.Condition{
Field: &gou.Expression{Field: "__yao_team_id"},
OP: "null",
}},
{Condition: gou.Condition{
Field: &gou.Expression{Field: "__yao_created_by"},
Value: authInfo.UserID,
OP: "=",
}},
},
},
},
})
return wheres
}
return wheres
}

631
agent/assistant/search_auth_integration_test.go Normal file
View file

@ -0,0 +1,631 @@
package assistant_test
import (
"context"
"fmt"
"testing"
"time"
"github.com/stretchr/testify/assert"
graphragtypes "github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/yao/agent/assistant"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search"
searchTypes "github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
"github.com/yaoapp/yao/kb"
"github.com/yaoapp/yao/kb/api"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
)
// ========== Test Constants ==========
const (
// Collection IDs for auth testing
AuthTestCollectionTeam1 = "auth_test_team1"
AuthTestCollectionTeam2 = "auth_test_team2"
AuthTestCollectionPublic = "auth_test_public"
// Test users and teams
TestUserA = "user_a"
TestUserB = "user_b"
TestTeam1 = "team_1"
TestTeam2 = "team_2"
)
// ========== Setup Test ==========
// TestAuthSearchSetup creates test collections with different permissions.
// Run once before running auth tests:
//
// go test -v -run "TestAuthSearchSetup" ./agent/assistant/...
func TestAuthSearchSetup(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
if kb.API == nil {
t.Fatal("KB API not initialized")
}
ctx := context.Background()
// Check if collections already exist
team1Exists := collectionReady(ctx, AuthTestCollectionTeam1, 2)
team2Exists := collectionReady(ctx, AuthTestCollectionTeam2, 2)
publicExists := collectionReady(ctx, AuthTestCollectionPublic, 2)
if team1Exists && team2Exists && publicExists {
t.Log("✓ All auth test collections already exist")
t.Log(" Run TestAuthSearchCleanup to recreate")
return
}
// Cleanup existing
t.Log("Cleaning up existing collections...")
cleanupAuthCollections(ctx, t)
time.Sleep(1 * time.Second)
// Create Team1 collection (owned by UserA, Team1)
t.Log("Creating Team1 collection...")
createAuthCollection(ctx, t, AuthTestCollectionTeam1, TestUserA, TestTeam1, false, "team")
addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc1", "Team1 private document about quantum physics and relativity theory.")
addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc2", "Team1 shared document about machine learning and neural networks.")
// Create Team2 collection (owned by UserB, Team2)
t.Log("Creating Team2 collection...")
createAuthCollection(ctx, t, AuthTestCollectionTeam2, TestUserB, TestTeam2, false, "team")
addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc1", "Team2 private document about deep learning algorithms.")
addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc2", "Team2 shared document about computer vision techniques.")
// Create Public collection
t.Log("Creating Public collection...")
createAuthCollection(ctx, t, AuthTestCollectionPublic, TestUserA, TestTeam1, true, "")
addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc1", "Public document about artificial intelligence and robotics.")
addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc2", "Public document about natural language processing.")
// Wait for indexing
t.Log("Waiting for indexing...")
time.Sleep(2 * time.Second)
t.Log("✓ Auth test setup complete!")
}
// TestAuthSearchCleanup removes auth test collections.
func TestAuthSearchCleanup(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
if kb.API == nil {
t.Fatal("KB API not initialized")
}
ctx := context.Background()
cleanupAuthCollections(ctx, t)
t.Log("✓ Auth test cleanup complete!")
}
// ========== KB Collection-Level Auth Filter Tests ==========
// Note: KB permission filtering works at the Collection level.
// The Collection metadata contains __yao_team_id, __yao_created_by, public, share fields.
// FilterKBCollectionsByAuth filters collections based on user authorization.
func TestKBCollectionAuthFilter(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Ensure test data exists (auto-creates if not)
ensureAuthTestData(t)
t.Run("TeamMemberCanAccessTeamCollection", func(t *testing.T) {
// UserA from Team1 should access Team1 collection
ctx := createAuthContext(TestUserA, TestTeam1, true, false)
collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
assert.Contains(t, allowed, AuthTestCollectionTeam1, "Team1 member should access Team1 collection")
t.Logf(" Allowed collections: %v", allowed)
})
t.Run("TeamMemberCannotAccessOtherTeamCollection", func(t *testing.T) {
// UserA from Team1 should NOT access Team2 collection
ctx := createAuthContext(TestUserA, TestTeam1, true, false)
collections := []string{AuthTestCollectionTeam2}
allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
assert.NotContains(t, allowed, AuthTestCollectionTeam2, "Team1 member should NOT access Team2 collection")
t.Logf(" Allowed collections: %v (expected empty)", allowed)
})
t.Run("OwnerCanAccessOwnCollection", func(t *testing.T) {
// UserA with OwnerOnly should access collections they created
ctx := createAuthContext(TestUserA, "", false, true)
collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
assert.Contains(t, allowed, AuthTestCollectionTeam1, "Owner should access own collection")
assert.NotContains(t, allowed, AuthTestCollectionTeam2, "Owner should NOT access other's collection")
t.Logf(" Allowed collections: %v", allowed)
})
t.Run("PublicCollectionAccessibleToAll", func(t *testing.T) {
// Note: The 'public' field in Metadata is not automatically saved to the database
// by the current KB API. This test documents the expected behavior.
// When public=true is properly set in DB, this should pass.
// First, check the collection metadata
bgCtx := context.Background()
collection, err := kb.API.GetCollection(bgCtx, AuthTestCollectionPublic)
assert.NoError(t, err)
// Check if public is set correctly
publicVal := collection["public"]
t.Logf(" Public collection public field: %v (type: %T)", publicVal, publicVal)
// If public is not set (0 or false), the test documents current behavior
// The collection should be accessible via owner check since UserA created it
ctx := createAuthContext(TestUserA, TestTeam1, false, true) // Owner check
collections := []string{AuthTestCollectionPublic}
allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
assert.Contains(t, allowed, AuthTestCollectionPublic, "Owner should access their collection")
t.Logf(" Allowed collections (owner check): %v", allowed)
})
t.Run("NoConstraintsMeansFullAccess", func(t *testing.T) {
// User with no constraints should access all collections
ctx := createAuthContext(TestUserA, TestTeam1, false, false)
collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
assert.Len(t, allowed, 3, "No constraints should allow all collections")
t.Logf(" Allowed collections: %v", allowed)
})
t.Run("NilContextMeansFullAccess", func(t *testing.T) {
collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
allowed := assistant.FilterKBCollectionsByAuth(nil, collections)
assert.Len(t, allowed, 2, "Nil context should allow all collections")
t.Logf(" Allowed collections: %v", allowed)
})
}
// ========== DB Auth Wheres Tests ==========
func TestDBAuthWheresFilter(t *testing.T) {
// Note: This test doesn't need KB, just tests the BuildDBAuthWheres function
t.Run("TeamOnlyGeneratesCorrectWheres", func(t *testing.T) {
ctx := createAuthContext(TestUserA, TestTeam1, true, false)
wheres := assistant.BuildDBAuthWheres(ctx)
assert.NotNil(t, wheres)
assert.Len(t, wheres, 1)
// Verify structure: should have 2 top-level conditions (public OR team filter)
where := wheres[0]
assert.Len(t, where.Wheres, 2, "Should have 2 conditions: public OR team")
// First condition: public = true (OR)
publicCond := where.Wheres[0]
assert.NotNil(t, publicCond.Condition.Field)
assert.Equal(t, "public", publicCond.Condition.Field.Field)
assert.Equal(t, true, publicCond.Condition.Value)
assert.True(t, publicCond.Condition.OR)
// Second condition: team filter with nested conditions
teamCond := where.Wheres[1]
assert.Len(t, teamCond.Wheres, 2, "Team filter should have team_id and (created_by OR share)")
// Team ID check
teamIDCond := teamCond.Wheres[0]
assert.Equal(t, "__yao_team_id", teamIDCond.Condition.Field.Field)
assert.Equal(t, TestTeam1, teamIDCond.Condition.Value)
// Created by OR share = team
ownerOrShareCond := teamCond.Wheres[1]
assert.Len(t, ownerOrShareCond.Wheres, 2)
assert.Equal(t, "__yao_created_by", ownerOrShareCond.Wheres[0].Condition.Field.Field)
assert.Equal(t, TestUserA, ownerOrShareCond.Wheres[0].Condition.Value)
assert.Equal(t, "share", ownerOrShareCond.Wheres[1].Condition.Field.Field)
assert.Equal(t, "team", ownerOrShareCond.Wheres[1].Condition.Value)
assert.True(t, ownerOrShareCond.Wheres[1].Condition.OR)
t.Logf(" TeamOnly: Verified team_id=%s, created_by=%s", TestTeam1, TestUserA)
})
t.Run("OwnerOnlyGeneratesCorrectWheres", func(t *testing.T) {
ctx := createAuthContext(TestUserA, "", false, true)
wheres := assistant.BuildDBAuthWheres(ctx)
assert.NotNil(t, wheres)
assert.Len(t, wheres, 1)
// Verify structure: should have 2 top-level conditions (public OR owner filter)
where := wheres[0]
assert.Len(t, where.Wheres, 2, "Should have 2 conditions: public OR owner")
// First condition: public = true (OR)
publicCond := where.Wheres[0]
assert.NotNil(t, publicCond.Condition.Field)
assert.Equal(t, "public", publicCond.Condition.Field.Field)
assert.Equal(t, true, publicCond.Condition.Value)
assert.True(t, publicCond.Condition.OR)
// Second condition: owner filter with nested conditions
ownerCond := where.Wheres[1]
assert.Len(t, ownerCond.Wheres, 2, "Owner filter should have team_id IS NULL and created_by")
// Team ID is null check
teamNullCond := ownerCond.Wheres[0]
assert.Equal(t, "__yao_team_id", teamNullCond.Condition.Field.Field)
assert.Equal(t, "null", teamNullCond.Condition.OP)
// Created by check
createdByCond := ownerCond.Wheres[1]
assert.Equal(t, "__yao_created_by", createdByCond.Condition.Field.Field)
assert.Equal(t, TestUserA, createdByCond.Condition.Value)
t.Logf(" OwnerOnly: Verified created_by=%s, team_id IS NULL", TestUserA)
})
t.Run("NoConstraintsReturnsNil", func(t *testing.T) {
ctx := createAuthContext(TestUserA, TestTeam1, false, false)
wheres := assistant.BuildDBAuthWheres(ctx)
assert.Nil(t, wheres, "No constraints should return nil")
t.Log(" No constraints: nil wheres (no filter)")
})
t.Run("EmptyTeamIDReturnsNil", func(t *testing.T) {
ctx := createAuthContext(TestUserA, "", true, false)
wheres := assistant.BuildDBAuthWheres(ctx)
assert.Nil(t, wheres, "Empty TeamID with TeamOnly should return nil")
t.Log(" Empty TeamID with TeamOnly: nil wheres")
})
t.Run("EmptyUserIDReturnsNil", func(t *testing.T) {
ctx := createAuthContext("", TestTeam1, false, true)
wheres := assistant.BuildDBAuthWheres(ctx)
assert.Nil(t, wheres, "Empty UserID with OwnerOnly should return nil")
t.Log(" Empty UserID with OwnerOnly: nil wheres")
})
t.Run("NilContextReturnsNil", func(t *testing.T) {
wheres := assistant.BuildDBAuthWheres(nil)
assert.Nil(t, wheres, "Nil context should return nil")
t.Log(" Nil context: nil wheres")
})
t.Run("NilAuthorizedReturnsNil", func(t *testing.T) {
ctx := &agentContext.Context{Authorized: nil}
wheres := assistant.BuildDBAuthWheres(ctx)
assert.Nil(t, wheres, "Nil Authorized should return nil")
t.Log(" Nil Authorized: nil wheres")
})
}
// ========== KB Search Integration Tests ==========
func TestKBSearchIntegration(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Ensure test data exists (auto-creates if not)
ensureAuthTestData(t)
t.Run("TeamMemberSearchOnlyFindsTeamData", func(t *testing.T) {
// UserA from Team1 searches - should ONLY find Team1 data
ctx := createAuthContext(TestUserA, TestTeam1, true, false)
// Filter collections first
allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
// Should only allow Team1
assert.Contains(t, allowed, AuthTestCollectionTeam1)
assert.NotContains(t, allowed, AuthTestCollectionTeam2)
assert.Len(t, allowed, 1, "Should only have 1 allowed collection")
// Search on allowed collections
result := executeKBSearchOnCollections(t, allowed, "quantum physics deep learning")
assert.Greater(t, len(result.Items), 0, "Should find Team1 documents")
// Verify ALL results are from Team1 collection only
for _, item := range result.Items {
assert.Equal(t, AuthTestCollectionTeam1, item.Collection,
"All results should be from Team1 collection, got: %s", item.Collection)
}
t.Logf(" ✓ Team1 member found %d items, all from Team1 collection", len(result.Items))
})
t.Run("TeamMemberCannotAccessOtherTeamData", func(t *testing.T) {
// UserA from Team1 tries to access Team2 - should be blocked
ctx := createAuthContext(TestUserA, TestTeam1, true, false)
// Try to filter Team2 collection
collections := []string{AuthTestCollectionTeam2}
allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
// Should be empty - no access
assert.Empty(t, allowed, "Team1 member should NOT have access to Team2 collection")
t.Log(" ✓ Team1 member correctly blocked from Team2 collection")
})
t.Run("OwnerSearchOnlyFindsOwnData", func(t *testing.T) {
// UserA with OwnerOnly - should only find collections they created
ctx := createAuthContext(TestUserA, "", false, true)
// Filter all collections
allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
// UserA created Team1 and Public, not Team2
assert.Contains(t, allowed, AuthTestCollectionTeam1, "Owner should access Team1 (created by UserA)")
assert.Contains(t, allowed, AuthTestCollectionPublic, "Owner should access Public (created by UserA)")
assert.NotContains(t, allowed, AuthTestCollectionTeam2, "Owner should NOT access Team2 (created by UserB)")
// Search and verify results
result := executeKBSearchOnCollections(t, allowed, "quantum artificial intelligence")
assert.Greater(t, len(result.Items), 0, "Should find owner's documents")
// Verify NO results from Team2
for _, item := range result.Items {
assert.NotEqual(t, AuthTestCollectionTeam2, item.Collection,
"Should NOT have results from Team2, got: %s", item.Collection)
}
t.Logf(" ✓ Owner found %d items, none from Team2", len(result.Items))
})
t.Run("NoConstraintsSearchFindsAllData", func(t *testing.T) {
// User with no constraints - should find all data
ctx := createAuthContext(TestUserA, TestTeam1, false, false)
// Filter all collections
allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
// Should have access to all
assert.Len(t, allowed, 3, "No constraints should allow all collections")
// Search and verify results from multiple collections
result := executeKBSearchOnCollections(t, allowed, "quantum deep learning artificial")
// Should find results from multiple collections
collectionsFound := make(map[string]bool)
for _, item := range result.Items {
collectionsFound[item.Collection] = true
}
assert.Greater(t, len(collectionsFound), 1, "Should find results from multiple collections")
t.Logf(" ✓ No constraints: found %d items from %d collections", len(result.Items), len(collectionsFound))
})
t.Run("SearchResultsMatchCollectionFilter", func(t *testing.T) {
// Verify that search results ONLY come from allowed collections
ctx := createAuthContext(TestUserB, TestTeam2, true, false)
// UserB from Team2 - should only access Team2
allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
assert.Contains(t, allowed, AuthTestCollectionTeam2, "Team2 member should access Team2")
assert.NotContains(t, allowed, AuthTestCollectionTeam1, "Team2 member should NOT access Team1")
// Search
result := executeKBSearchOnCollections(t, allowed, "deep learning computer vision")
// Verify results
if len(result.Items) > 0 {
for _, item := range result.Items {
// Results should only be from allowed collections
assert.Contains(t, allowed, item.Collection,
"Result from %s should be in allowed list %v", item.Collection, allowed)
}
t.Logf(" ✓ Team2 member found %d items, all from allowed collections", len(result.Items))
} else {
t.Log(" ✓ Team2 member found 0 items (collection may be empty)")
}
})
}
// ========== Helper Functions ==========
// ensureAuthTestData checks if auth test data exists, creates if not.
// This is a utility function that can be called from any test.
// Note: testutils.Prepare must be called before this function.
func ensureAuthTestData(t *testing.T) {
if kb.API == nil {
t.Fatal("KB API not initialized - call testutils.Prepare first")
}
ctx := context.Background()
// Check if all collections already exist with required documents
team1Ready := collectionReady(ctx, AuthTestCollectionTeam1, 2)
team2Ready := collectionReady(ctx, AuthTestCollectionTeam2, 2)
publicReady := collectionReady(ctx, AuthTestCollectionPublic, 2)
if team1Ready && team2Ready && publicReady {
// All data exists, skip creation
return
}
// Data not ready, create it
t.Log("Auth test data not found, creating...")
createAuthTestData(t, ctx)
}
// createAuthTestData creates the test collections and documents
func createAuthTestData(t *testing.T, ctx context.Context) {
// Cleanup existing
t.Log("Cleaning up existing collections...")
cleanupAuthCollections(ctx, t)
time.Sleep(1 * time.Second)
// Create Team1 collection (owned by UserA, Team1)
t.Log("Creating Team1 collection...")
createAuthCollection(ctx, t, AuthTestCollectionTeam1, TestUserA, TestTeam1, false, "team")
addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc1", "Team1 private document about quantum physics and relativity theory.")
addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc2", "Team1 shared document about machine learning and neural networks.")
// Create Team2 collection (owned by UserB, Team2)
t.Log("Creating Team2 collection...")
createAuthCollection(ctx, t, AuthTestCollectionTeam2, TestUserB, TestTeam2, false, "team")
addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc1", "Team2 private document about deep learning algorithms.")
addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc2", "Team2 shared document about computer vision techniques.")
// Create Public collection
t.Log("Creating Public collection...")
createAuthCollection(ctx, t, AuthTestCollectionPublic, TestUserA, TestTeam1, true, "")
addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc1", "Public document about artificial intelligence and robotics.")
addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc2", "Public document about natural language processing.")
// Wait for indexing
t.Log("Waiting for indexing...")
time.Sleep(2 * time.Second)
t.Log("✓ Auth test data created!")
}
func createAuthContext(userID, teamID string, teamOnly, ownerOnly bool) *agentContext.Context {
return &agentContext.Context{
Authorized: &oauthtypes.AuthorizedInfo{
UserID: userID,
TeamID: teamID,
Constraints: oauthtypes.DataConstraints{
TeamOnly: teamOnly,
OwnerOnly: ownerOnly,
},
},
}
}
func collectionReady(ctx context.Context, collectionID string, minDocs int) bool {
collection, err := kb.API.GetCollection(ctx, collectionID)
if err != nil || collection == nil {
return false
}
docs, err := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{
Page: 1,
PageSize: 20,
CollectionID: collectionID,
})
if err != nil || docs == nil {
return false
}
return len(docs.Data) >= minDocs
}
func cleanupAuthCollections(ctx context.Context, t *testing.T) {
collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
for _, id := range collections {
if result, err := kb.API.RemoveCollection(ctx, id); err == nil && result.Removed {
t.Logf(" Removed: %s", id)
}
}
}
func createAuthCollection(ctx context.Context, t *testing.T, id, userID, teamID string, public bool, share string) {
params := &api.CreateCollectionParams{
ID: id,
Metadata: map[string]interface{}{
"name": id,
"public": public,
"share": share,
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Locale: "en",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
AuthScope: map[string]interface{}{
"__yao_created_by": userID,
"__yao_team_id": teamID,
},
}
_, err := kb.API.CreateCollection(ctx, params)
if err != nil {
t.Fatalf("Failed to create collection %s: %v", id, err)
}
t.Logf(" ✓ Created: %s", id)
}
func addAuthDocument(ctx context.Context, t *testing.T, collectionID, title, content string) {
params := &api.AddTextParams{
CollectionID: collectionID,
Text: content,
DocID: fmt.Sprintf("%s__%s", collectionID, sanitizeForID(title)),
Metadata: map[string]interface{}{
"title": title,
},
Chunking: &api.ProviderConfigParams{
ProviderID: "__yao.structured",
OptionID: "standard",
},
Embedding: &api.ProviderConfigParams{
ProviderID: "__yao.openai",
OptionID: "text-embedding-3-small",
},
}
_, err := kb.API.AddText(ctx, params)
if err != nil {
t.Logf(" Warning: Failed to add document '%s': %v", title, err)
return
}
t.Logf(" ✓ Added: %s", title)
}
func sanitizeForID(s string) string {
result := ""
for _, c := range s {
if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') {
result += string(c)
} else if c == ' ' {
result += "_"
}
}
return result
}
func executeKBSearch(t *testing.T, collectionID, query string, metadata map[string]interface{}) *searchTypes.Result {
return executeKBSearchOnCollections(t, []string{collectionID}, query)
}
func executeKBSearchOnCollections(t *testing.T, collections []string, query string) *searchTypes.Result {
if len(collections) == 0 {
return &searchTypes.Result{Items: []*searchTypes.ResultItem{}}
}
cfg := &searchTypes.Config{
KB: &searchTypes.KBConfig{
Collections: collections,
Threshold: 0.3,
},
}
searcher := search.New(cfg, nil)
req := &searchTypes.Request{
Type: searchTypes.SearchTypeKB,
Query: query,
Collections: collections,
Threshold: 0.3,
Limit: 20,
Source: searchTypes.SourceAuto,
}
result, err := searcher.Search(nil, req)
if err != nil {
t.Fatalf("Search failed: %v", err)
}
return result
}

123
agent/assistant/search_auth_kb.go Normal file
View file

@ -0,0 +1,123 @@
package assistant
import (
"context"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/kb"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
)
// FilterKBCollectionsByAuth filters collections based on user authorization.
// Returns only collections that the user has permission to access.
// Permission is determined by Collection's metadata (public, share, __yao_team_id, __yao_created_by).
func FilterKBCollectionsByAuth(ctx *agentContext.Context, collections []string) []string {
if ctx == nil || ctx.Authorized == nil {
return collections // No auth context, return all
}
authInfo := ctx.Authorized
// No constraints, return all collections
if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
return collections
}
// Check KB API
if kb.API == nil {
return collections // KB not initialized, return all
}
var allowed []string
bgCtx := context.Background()
for _, collectionID := range collections {
// Get collection metadata
collection, err := kb.API.GetCollection(bgCtx, collectionID)
if err != nil {
continue // Skip if can't get collection
}
if hasCollectionAccess(authInfo, collection) {
allowed = append(allowed, collectionID)
}
}
return allowed
}
// hasCollectionAccess checks if user has access to a collection based on its metadata.
func hasCollectionAccess(authInfo *oauthtypes.AuthorizedInfo, collection map[string]interface{}) bool {
if authInfo == nil {
return true
}
// No constraints, allow access
if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
return true
}
// Check public access (handle different types: bool, int, float64)
if isPublicValue(collection["public"]) {
return true
}
// Get metadata for permission fields
metadata, _ := collection["metadata"].(map[string]interface{})
if metadata == nil {
metadata = collection
}
// Team only check
if authInfo.Constraints.TeamOnly && authInfo.TeamID != "" {
teamID, _ := metadata["__yao_team_id"].(string)
if teamID == "" {
teamID, _ = collection["__yao_team_id"].(string)
}
if teamID == authInfo.TeamID {
createdBy, _ := metadata["__yao_created_by"].(string)
if createdBy == "" {
createdBy, _ = collection["__yao_created_by"].(string)
}
share, _ := metadata["share"].(string)
if share == "" {
share, _ = collection["share"].(string)
}
if createdBy == authInfo.UserID || share == "team" {
return true
}
}
}
// Owner only check
if authInfo.Constraints.OwnerOnly && authInfo.UserID != "" {
createdBy, _ := metadata["__yao_created_by"].(string)
if createdBy == "" {
createdBy, _ = collection["__yao_created_by"].(string)
}
if createdBy == authInfo.UserID {
return true
}
}
return false
}
// isPublicValue checks if a value represents "public" access
func isPublicValue(v interface{}) bool {
switch val := v.(type) {
case bool:
return val
case int:
return val == 1
case int64:
return val == 1
case float64:
return val == 1
case string:
return val == "true" || val == "1"
}
return false
}

110
agent/search/handlers/kb/handler.go
View file

@ -1,12 +1,16 @@
package kb
import (
"context"
"fmt"
"time"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/kb"
kbapi "github.com/yaoapp/yao/kb/api"
)
// Handler implements KB search
// Handler implements KB search using the KB API
type Handler struct {
config *types.KBConfig // KB search configuration
}
@ -22,7 +26,6 @@ func (h *Handler) Type() types.SearchType {
}
// Search executes vector search and optional graph association
// TODO: Implement actual vector search and graph association logic
func (h *Handler) Search(req *types.Request) (*types.Result, error) {
start := time.Now()
@ -39,6 +42,19 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) {
}, nil
}
// Check if KB API is available
if kb.API == nil {
return &types.Result{
Type: types.SearchTypeKB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
Error: "knowledge base not initialized",
}, nil
}
// Get collections from request or config
collections := req.Collections
if len(collections) == 0 && h.config != nil {
@ -72,24 +88,96 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) {
limit = 10 // default
}
// TODO: Implement actual vector search
// 1. Generate embedding for query using collection's embedding config
// 2. Search each collection with vector similarity
// 3. If req.Graph is true, perform graph association
// 4. Merge and return results
// Determine search mode
mode := kbapi.SearchModeVector
if req.Graph {
mode = kbapi.SearchModeExpand
}
if h.config != nil && h.config.Graph {
mode = kbapi.SearchModeExpand
}
// For now, return empty result (skeleton)
result := &types.Result{
// Build KB API queries - one per collection
var queries []kbapi.Query
for _, collectionID := range collections {
queries = append(queries, kbapi.Query{
CollectionID: collectionID,
Input: req.Query,
Mode: mode,
Threshold: threshold,
PageSize: limit,
Metadata: req.Metadata,
})
}
// Execute search using KB API
ctx := context.Background()
searchResult, err := kb.API.Search(ctx, queries)
if err != nil {
return &types.Result{
Type: types.SearchTypeKB,
Query: req.Query,
Source: req.Source,
Items: []*types.ResultItem{},
Total: 0,
Duration: time.Since(start).Milliseconds(),
Error: fmt.Sprintf("search failed: %v", err),
}, nil
}
// Store threshold in result metadata for debugging
_ = threshold
// Convert segments to result items
// Note: MinScore filtering is already done by KB API, no need to filter again
items := make([]*types.ResultItem, 0, len(searchResult.Segments))
for _, seg := range searchResult.Segments {
item := &types.ResultItem{
Type: types.SearchTypeKB,
Source: req.Source,
Score: seg.Score,
Content: seg.Text,
DocumentID: seg.DocumentID,
Collection: seg.CollectionID,
Metadata: seg.Metadata,
}
// Extract title from metadata if available
if seg.Metadata != nil {
if title, ok := seg.Metadata["title"].(string); ok {
item.Title = title
}
}
items = append(items, item)
}
// Convert graph data if available
var graphNodes []*types.GraphNode
if searchResult.Graph != nil {
for _, node := range searchResult.Graph.Nodes {
// Extract name from properties if available
name := ""
if node.Properties != nil {
if n, ok := node.Properties["name"].(string); ok {
name = n
}
}
graphNodes = append(graphNodes, &types.GraphNode{
ID: node.ID,
Type: node.EntityType,
Name: name,
Metadata: node.Properties,
})
}
}
result := &types.Result{
Type: types.SearchTypeKB,
Query: req.Query,
Source: req.Source,
Items: items,
Total: len(items),
Duration: time.Since(start).Milliseconds(),
GraphNodes: graphNodes,
}
return result, nil
}

241
agent/search/handlers/kb/handler_test.go
View file

@ -31,13 +31,12 @@ func TestHandler_Type(t *testing.T) {
assert.Equal(t, types.SearchTypeKB, h.Type())
}
func TestHandler_Search(t *testing.T) {
func TestHandler_Search_Validation(t *testing.T) {
tests := []struct {
name string
config *types.KBConfig
req *types.Request
expectError string
expectItems int
}{
{
name: "empty query",
@ -47,85 +46,15 @@ func TestHandler_Search(t *testing.T) {
Query: "",
},
expectError: "query is required",
expectItems: 0,
},
{
name: "no collections in request or config",
name: "no collections - KB not initialized",
config: nil,
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
},
expectError: "",
expectItems: 0,
},
{
name: "collections from config",
config: &types.KBConfig{
Collections: []string{"docs"},
Threshold: 0.7,
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "collections from request",
config: nil,
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs", "faq"},
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with threshold from request",
config: &types.KBConfig{
Collections: []string{"docs"},
Threshold: 0.7,
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Threshold: 0.9,
Collections: []string{"docs"},
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with graph enabled",
config: &types.KBConfig{
Collections: []string{"docs"},
Graph: true,
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs"},
Graph: true,
},
expectError: "",
expectItems: 0, // skeleton returns empty
},
{
name: "with limit",
config: &types.KBConfig{
Collections: []string{"docs"},
},
req: &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs"},
Limit: 5,
},
expectError: "",
expectItems: 0, // skeleton returns empty
expectError: "knowledge base not initialized",
},
}
@ -138,7 +67,6 @@ func TestHandler_Search(t *testing.T) {
assert.NotNil(t, result)
assert.Equal(t, types.SearchTypeKB, result.Type)
assert.Equal(t, tt.req.Query, result.Query)
assert.Equal(t, tt.expectItems, len(result.Items))
if tt.expectError != "" {
assert.Equal(t, tt.expectError, result.Error)
@ -168,3 +96,166 @@ func TestHandler_Search_SourcePreserved(t *testing.T) {
assert.Equal(t, source, result.Source)
}
}
func TestHandler_Search_CollectionsFromConfig(t *testing.T) {
cfg := &types.KBConfig{
Collections: []string{"docs", "faq"},
Threshold: 0.7,
}
h := NewHandler(cfg)
// Request without collections should use config collections
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.NotNil(t, result)
// Without KB initialized, we get "knowledge base not initialized" error
assert.Equal(t, "knowledge base not initialized", result.Error)
}
func TestHandler_Search_CollectionsFromRequest(t *testing.T) {
h := NewHandler(nil)
// Request with collections
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs", "faq"},
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.NotNil(t, result)
// Without KB initialized, we get "knowledge base not initialized" error
assert.Equal(t, "knowledge base not initialized", result.Error)
}
func TestHandler_Search_ThresholdHandling(t *testing.T) {
tests := []struct {
name string
configThreshold float64
reqThreshold float64
}{
{
name: "threshold from request",
configThreshold: 0.7,
reqThreshold: 0.9,
},
{
name: "threshold from config",
configThreshold: 0.8,
reqThreshold: 0,
},
{
name: "default threshold",
configThreshold: 0,
reqThreshold: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var cfg *types.KBConfig
if tt.configThreshold > 0 {
cfg = &types.KBConfig{
Collections: []string{"docs"},
Threshold: tt.configThreshold,
}
}
h := NewHandler(cfg)
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Threshold: tt.reqThreshold,
Collections: []string{"docs"},
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.NotNil(t, result)
})
}
}
func TestHandler_Search_GraphMode(t *testing.T) {
tests := []struct {
name string
configGraph bool
reqGraph bool
}{
{
name: "graph from request",
configGraph: false,
reqGraph: true,
},
{
name: "graph from config",
configGraph: true,
reqGraph: false,
},
{
name: "no graph",
configGraph: false,
reqGraph: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cfg := &types.KBConfig{
Collections: []string{"docs"},
Graph: tt.configGraph,
}
h := NewHandler(cfg)
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs"},
Graph: tt.reqGraph,
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.NotNil(t, result)
})
}
}
func TestHandler_Search_LimitHandling(t *testing.T) {
h := NewHandler(&types.KBConfig{Collections: []string{"docs"}})
tests := []struct {
name string
limit int
}{
{
name: "custom limit",
limit: 5,
},
{
name: "default limit",
limit: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req := &types.Request{
Type: types.SearchTypeKB,
Query: "test query",
Collections: []string{"docs"},
Limit: tt.limit,
}
result, err := h.Search(req)
assert.NoError(t, err)
assert.NotNil(t, result)
})
}
}

22
agent/search/jsapi_test.go
View file

@ -8,6 +8,7 @@ import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search"
"github.com/yaoapp/yao/agent/search/types"
"github.com/yaoapp/yao/agent/testutils"
)
func TestNewJSAPI(t *testing.T) {
@ -51,6 +52,9 @@ func TestJSAPI_Web_WithOptions(t *testing.T) {
}
func TestJSAPI_KB(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
api := search.NewJSAPI(nil, &types.Config{
KB: &types.KBConfig{Collections: []string{"docs"}},
}, nil)
@ -66,6 +70,9 @@ func TestJSAPI_KB(t *testing.T) {
}
func TestJSAPI_KB_WithOptions(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
api := search.NewJSAPI(nil, &types.Config{
KB: &types.KBConfig{Collections: []string{"docs"}},
}, nil)
@ -87,6 +94,9 @@ func TestJSAPI_KB_WithOptions(t *testing.T) {
}
func TestJSAPI_DB(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
api := search.NewJSAPI(nil, &types.Config{
DB: &types.DBConfig{Models: []string{"product"}},
}, &search.Uses{QueryDSL: "builtin"})
@ -102,6 +112,9 @@ func TestJSAPI_DB(t *testing.T) {
}
func TestJSAPI_DB_WithOptions(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
api := search.NewJSAPI(nil, &types.Config{
DB: &types.DBConfig{Models: []string{"product"}},
}, &search.Uses{QueryDSL: "builtin"})
@ -122,6 +135,9 @@ func TestJSAPI_DB_WithOptions(t *testing.T) {
}
func TestJSAPI_All(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
api := search.NewJSAPI(nil, &types.Config{
KB: &types.KBConfig{Collections: []string{"docs"}},
DB: &types.DBConfig{Models: []string{"product"}},
@ -155,6 +171,9 @@ func TestJSAPI_All(t *testing.T) {
}
func TestJSAPI_Any(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
api := search.NewJSAPI(nil, &types.Config{
KB: &types.KBConfig{Collections: []string{"docs"}},
DB: &types.DBConfig{Models: []string{"product"}},
@ -186,6 +205,9 @@ func TestJSAPI_Any(t *testing.T) {
}
func TestJSAPI_Race(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
api := search.NewJSAPI(nil, &types.Config{
KB: &types.KBConfig{Collections: []string{"docs"}},
DB: &types.DBConfig{Models: []string{"product"}},

59
agent/search/search.go
View file

@ -151,14 +151,24 @@ func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]*
wg.Add(1)
go func(idx int, r *types.Request) {
defer wg.Done()
result, _ := s.Search(ctx, r)
// Check if done before starting
select {
case <-done:
return
default:
}
result, _ := s.Search(ctx, r)
// Try to send result
select {
case <-done:
// Already found a successful result
case resultChan <- struct {
idx int
result *types.Result
}{idx, result}:
case <-done:
// Already found a successful result, discard this one
}
}(i, req)
}
@ -170,25 +180,23 @@ func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]*
}()
// Collect results until we find one with items (success)
var mu sync.Mutex
var foundSuccess bool
for res := range resultChan {
mu.Lock()
results[res.idx] = res.result
// Check if this result has items (success = has results and no error)
if res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" {
mu.Unlock()
close(done) // Signal other goroutines to stop sending
return results, nil
if !foundSuccess && res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" {
foundSuccess = true
close(done) // Signal other goroutines to stop
}
mu.Unlock()
}
// No successful result found, return all results
// All goroutines have completed (resultChan is closed)
return results, nil
}
// parallelRace returns as soon as any search completes (like Promise.race)
// Returns immediately when first result arrives, regardless of success/failure
// Note: Still waits for all goroutines to complete before returning to avoid resource leaks
func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
results := make([]*types.Result, len(reqs))
resultChan := make(chan struct {
@ -203,14 +211,24 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]
wg.Add(1)
go func(idx int, r *types.Request) {
defer wg.Done()
result, _ := s.Search(ctx, r)
// Check if done before starting
select {
case <-done:
return
default:
}
result, _ := s.Search(ctx, r)
// Try to send result
select {
case <-done:
// Already got first result
case resultChan <- struct {
idx int
result *types.Result
}{idx, result}:
case <-done:
// Already got first result, discard this one
}
}(i, req)
}
@ -221,14 +239,17 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]
close(resultChan)
}()
// Return immediately when first result arrives
if res, ok := <-resultChan; ok {
// Get first result and signal others to stop
var gotFirst bool
for res := range resultChan {
results[res.idx] = res.result
close(done) // Signal other goroutines to stop sending
return results, nil
if !gotFirst {
gotFirst = true
close(done) // Signal other goroutines to stop
}
}
// No results (shouldn't happen with valid requests)
// All goroutines have completed (resultChan is closed)
return results, nil
}

1
agent/search/types/types.go
View file

@ -51,6 +51,7 @@ type Request struct {
Collections []string `json:"collections,omitempty"` // KB collection IDs
Threshold float64 `json:"threshold,omitempty"` // Similarity threshold (0-1)
Graph bool `json:"graph,omitempty"` // Enable graph association
Metadata map[string]interface{} `json:"metadata,omitempty"` // Metadata filter for KB search
// Database search specific
Models []string `json:"models,omitempty"` // Model IDs (e.g., "user", "agents.mybot.product")

10
kb/api/README.md
View file

@ -185,7 +185,7 @@ result, err := kb.API.RemoveDocuments(ctx, params)
The Search API supports batch queries with three search modes:
| Mode | Description |
|------|-------------|
| -------- | ------------------------------------------------------ |
| `vector` | Pure vector similarity search |
| `graph` | Graph traversal to find related segments via entities |
| `expand` | Graph-based entity expansion + vector search (default) |
@ -274,7 +274,7 @@ queries := []api.Query{
CollectionID: "my_collection",
Input: "physics",
DocumentID: "specific_doc_id", // filter to specific document
MinScore: 0.5, // minimum similarity score
Threshold: 0.5, // similarity threshold
Metadata: map[string]interface{}{
"category": "science",
},
@ -289,13 +289,13 @@ result, err := kb.API.Search(ctx, queries)
## Query Parameters
| Field | Type | Description |
|-------|------|-------------|
| -------------- | ------------- | ------------------------------------------------------ |
| `CollectionID` | string | Collection to search (required) |
| `Input` | string | Direct query text |
| `Messages` | []ChatMessage | Conversation history (last user message used as query) |
| `Mode` | SearchMode | `vector`, `graph`, or `expand` (default: `expand`) |
| `DocumentID` | string | Filter to specific document |
| `MinScore` | float64 | Minimum similarity threshold |
| `Threshold` | float64 | Similarity threshold (0-1) |
| `Metadata` | map | Filter by metadata fields |
| `MaxDepth` | int | Graph traversal depth (default: 2) |
| `Page` | int | Page number (1-based) |
@ -328,7 +328,7 @@ type ProviderConfigParams struct {
```
Common providers:
- **Chunking**: `__yao.structured` - text splitting
- **Embedding**: `__yao.openai` - vector embeddings
- **Extraction**: `__yao.openai` - entity/relationship extraction for graph

6
kb/api/search.go
View file

@ -15,7 +15,7 @@ import (
const (
DefaultSearchK = 10
DefaultMaxDepth = 2
DefaultMinScore = 0.0
DefaultThreshold = 0.0
MaxSearchK = 100
DefaultSearchPageSize = 20
)
@ -257,7 +257,7 @@ func (kb *KBInstance) searchVector(ctx context.Context, collectionID string, que
DocumentID: query.DocumentID,
Query: queryText,
K: k,
MinScore: query.MinScore,
MinScore: query.Threshold,
Embedding: embedding,
}
@ -356,7 +356,7 @@ func (kb *KBInstance) searchExpand(ctx context.Context, collectionID string, que
DocumentID: query.DocumentID,
Query: queryText,
K: k,
MinScore: query.MinScore,
MinScore: query.Threshold,
Embedding: embedding,
}

14
kb/api/search_test.go
View file

@ -314,30 +314,30 @@ func TestSearchQuery(t *testing.T) {
}
})
t.Run("Search_WithMinScore", func(t *testing.T) {
// Test: Filter by minimum score
t.Run("Search_WithThreshold", func(t *testing.T) {
// Test: Filter by similarity threshold
queries := []api.Query{
{
CollectionID: SearchTestScienceCollection,
Input: "Einstein relativity",
Mode: api.SearchModeVector,
MinScore: 0.5,
Threshold: 0.5,
PageSize: 10,
},
}
result, err := kb.API.Search(ctx, queries)
if err != nil {
t.Logf("MinScore search error: %v", err)
t.Logf("Threshold search error: %v", err)
return
}
assert.NotNil(t, result)
t.Logf("MinScore search returned %d segments", len(result.Segments))
t.Logf("Threshold search returned %d segments", len(result.Segments))
// Verify all results meet minimum score
// Verify all results meet threshold
for _, seg := range result.Segments {
assert.GreaterOrEqual(t, seg.Score, 0.5, "All segments should meet minimum score")
assert.GreaterOrEqual(t, seg.Score, 0.5, "All segments should meet threshold")
}
})

4
kb/api/types.go
View file

@ -233,8 +233,8 @@ type Query struct {
// DocumentID filters results to a specific document (optional)
DocumentID string `json:"document_id,omitempty" yaml:"document_id,omitempty"`
// MinScore filters results below this similarity threshold (optional)
MinScore float64 `json:"min_score,omitempty" yaml:"min_score,omitempty"`
// Threshold filters results below this similarity threshold (optional)
Threshold float64 `json:"threshold,omitempty" yaml:"threshold,omitempty"`
// Metadata filters segments by metadata fields (optional)
Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"`