Add KB Tests and Update Makefile for Test Management
- Introduced a new unit test target for KB tests in the Makefile, allowing for dedicated testing of the KB module. - Updated the test folder selection logic to exclude additional AI-related components, ensuring focused testing. - Enhanced the GitHub Actions workflows to include KB tests, setting up necessary services like Qdrant, Neo4j, and MongoDB for a comprehensive testing environment. - Refactored search test functions to improve data existence checks and streamline test setup processes, enhancing test reliability and maintainability.
This commit is contained in:
parent
53db8be522
commit
09d368dae7
7 changed files with 696 additions and 23 deletions
186
.github/workflows/pr-test.yml
vendored
186
.github/workflows/pr-test.yml
vendored
|
|
@ -164,6 +164,192 @@ env:
|
|||
TWILIO_TEST_PHONE: ${{ secrets.TWILIO_TEST_PHONE }}
|
||||
|
||||
jobs:
|
||||
# =============================================================================
|
||||
# KB Tests (kb) - Run once with SQLite (requires Qdrant, Neo4j, FastEmbed)
|
||||
# =============================================================================
|
||||
KBTest:
|
||||
runs-on: ubuntu-latest
|
||||
services:
|
||||
qdrant:
|
||||
image: qdrant/qdrant:latest
|
||||
ports:
|
||||
- 6333:6333
|
||||
- 6334:6334
|
||||
|
||||
fastembed:
|
||||
image: yaoapp/fastembed:latest-amd64
|
||||
env:
|
||||
FASTEMBED_PASSWORD: Yao@2026
|
||||
ports:
|
||||
- 6001:8000
|
||||
|
||||
neo4j:
|
||||
image: neo4j:latest
|
||||
ports:
|
||||
- "7687:7687"
|
||||
env:
|
||||
NEO4J_AUTH: neo4j/Yao2026Neo4j
|
||||
|
||||
mongodb:
|
||||
image: mongo:6.0
|
||||
ports:
|
||||
- 27017:27017
|
||||
env:
|
||||
MONGO_INITDB_ROOT_USERNAME: root
|
||||
MONGO_INITDB_ROOT_PASSWORD: 123456
|
||||
MONGO_INITDB_DATABASE: test
|
||||
|
||||
strategy:
|
||||
matrix:
|
||||
go: [1.24]
|
||||
if: >
|
||||
${{ github.event.workflow_run.event == 'pull_request' &&
|
||||
github.event.workflow_run.conclusion == 'success' }}
|
||||
steps:
|
||||
- name: "Download artifact"
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
var artifacts = await github.rest.actions.listWorkflowRunArtifacts({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
run_id: ${{github.event.workflow_run.id }},
|
||||
});
|
||||
var matchArtifact = artifacts.data.artifacts.filter((artifact) => {
|
||||
return artifact.name == "pr"
|
||||
})[0];
|
||||
var download = await github.rest.actions.downloadArtifact({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
artifact_id: matchArtifact.id,
|
||||
archive_format: 'zip',
|
||||
});
|
||||
var fs = require('fs');
|
||||
fs.writeFileSync('${{github.workspace}}/pr.zip', Buffer.from(download.data));
|
||||
|
||||
- name: "Read NR & SHA"
|
||||
run: |
|
||||
unzip pr.zip
|
||||
cat NR
|
||||
cat SHA
|
||||
echo HEAD=$(cat SHA) >> $GITHUB_ENV
|
||||
echo NR=$(cat NR) >> $GITHUB_ENV
|
||||
|
||||
- name: "Comment on PR"
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const { NR } = process.env
|
||||
var issue_number = NR;
|
||||
await github.rest.issues.createComment({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
issue_number: issue_number,
|
||||
body: '🤖 KB Tests (kb) running with SQLite...'
|
||||
});
|
||||
|
||||
- name: Checkout Kun
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/kun
|
||||
path: kun
|
||||
|
||||
- name: Checkout Xun
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/xun
|
||||
path: xun
|
||||
|
||||
- name: Checkout Gou
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/gou
|
||||
path: gou
|
||||
|
||||
- name: Checkout V8Go
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/v8go
|
||||
path: v8go
|
||||
|
||||
- name: Unzip libv8
|
||||
run: |
|
||||
files=$(find ./v8go -name "libv8*.zip")
|
||||
for file in $files; do
|
||||
dir=$(dirname "$file")
|
||||
echo "Extracting $file to directory $dir"
|
||||
unzip -o -d $dir $file
|
||||
rm -rf $dir/__MACOSX
|
||||
done
|
||||
|
||||
- name: Checkout Demo App
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/yao-dev-app
|
||||
path: app
|
||||
|
||||
- name: Checkout Extension
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/yao-extensions-dev
|
||||
path: extension
|
||||
|
||||
- name: Move Dependencies
|
||||
run: |
|
||||
mv kun ../
|
||||
mv xun ../
|
||||
mv gou ../
|
||||
mv v8go ../
|
||||
mv app ../
|
||||
mv extension ../
|
||||
|
||||
- name: Checkout pull request HEAD commit
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ env.HEAD }}
|
||||
|
||||
- name: Setup Apple Private Key
|
||||
run: |
|
||||
mkdir -p ../app/openapi/certs/apple
|
||||
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
|
||||
|
||||
- name: Setup Go ${{ matrix.go }}
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: ${{ matrix.go }}
|
||||
|
||||
- name: Start Redis
|
||||
uses: supercharge/redis-github-action@1.4.0
|
||||
with:
|
||||
redis-version: 6
|
||||
|
||||
- name: Setup Go Tools
|
||||
run: make tools
|
||||
|
||||
- name: Setup ENV (SQLite)
|
||||
run: |
|
||||
mkdir -p ${{ github.WORKSPACE }}/../app/db
|
||||
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
|
||||
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
|
||||
|
||||
- name: Run KB Tests (kb)
|
||||
run: make unit-test-kb
|
||||
|
||||
- name: "Comment on PR - KB Tests Done"
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const { NR } = process.env
|
||||
var issue_number = NR;
|
||||
await github.rest.issues.createComment({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
issue_number: issue_number,
|
||||
body: '✅ KB Tests (kb) passed!'
|
||||
});
|
||||
|
||||
# =============================================================================
|
||||
# AI Tests (agent, aigc) - Run once with SQLite
|
||||
# =============================================================================
|
||||
|
|
|
|||
124
.github/workflows/unit-test.yml
vendored
124
.github/workflows/unit-test.yml
vendored
|
|
@ -171,6 +171,130 @@ env:
|
|||
TWILIO_TEST_PHONE: ${{ secrets.TWILIO_TEST_PHONE }}
|
||||
|
||||
jobs:
|
||||
# =============================================================================
|
||||
# KB Tests (kb) - Run once with SQLite (requires Qdrant, Neo4j, FastEmbed)
|
||||
# =============================================================================
|
||||
kb-test:
|
||||
runs-on: ubuntu-latest
|
||||
services:
|
||||
qdrant:
|
||||
image: qdrant/qdrant:latest
|
||||
ports:
|
||||
- 6333:6333
|
||||
- 6334:6334
|
||||
|
||||
fastembed:
|
||||
image: yaoapp/fastembed:latest-amd64
|
||||
env:
|
||||
FASTEMBED_PASSWORD: Yao@2026
|
||||
ports:
|
||||
- 6001:8000
|
||||
|
||||
neo4j:
|
||||
image: neo4j:latest
|
||||
ports:
|
||||
- "7687:7687"
|
||||
env:
|
||||
NEO4J_AUTH: neo4j/Yao2026Neo4j
|
||||
|
||||
mongodb:
|
||||
image: mongo:6.0
|
||||
ports:
|
||||
- 27017:27017
|
||||
env:
|
||||
MONGO_INITDB_ROOT_USERNAME: root
|
||||
MONGO_INITDB_ROOT_PASSWORD: 123456
|
||||
MONGO_INITDB_DATABASE: test
|
||||
|
||||
strategy:
|
||||
matrix:
|
||||
go: [1.24]
|
||||
steps:
|
||||
- name: Checkout Kun
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: ${{ env.REPO_KUN }}
|
||||
path: kun
|
||||
|
||||
- name: Checkout Xun
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: ${{ env.REPO_XUN }}
|
||||
path: xun
|
||||
|
||||
- name: Checkout Gou
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: ${{ env.REPO_GOU }}
|
||||
path: gou
|
||||
|
||||
- name: Checkout V8Go
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/v8go
|
||||
path: v8go
|
||||
|
||||
- name: Unzip libv8
|
||||
run: |
|
||||
files=$(find ./v8go -name "libv8*.zip")
|
||||
for file in $files; do
|
||||
dir=$(dirname "$file")
|
||||
echo "Extracting $file to directory $dir"
|
||||
unzip -o -d $dir $file
|
||||
rm -rf $dir/__MACOSX
|
||||
done
|
||||
|
||||
- name: Checkout Demo App
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/yao-dev-app
|
||||
path: app
|
||||
|
||||
- name: Checkout Extension
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: yaoapp/yao-extensions-dev
|
||||
path: extension
|
||||
|
||||
- name: Move Dependencies
|
||||
run: |
|
||||
mv kun ../
|
||||
mv xun ../
|
||||
mv gou ../
|
||||
mv v8go ../
|
||||
mv app ../
|
||||
mv extension ../
|
||||
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Setup Apple Private Key
|
||||
run: |
|
||||
mkdir -p ../app/openapi/certs/apple
|
||||
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
|
||||
|
||||
- name: Setup Go ${{ matrix.go }}
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: ${{ matrix.go }}
|
||||
|
||||
- name: Start Redis
|
||||
uses: supercharge/redis-github-action@1.4.0
|
||||
with:
|
||||
redis-version: 6
|
||||
|
||||
- name: Setup Go Tools
|
||||
run: make tools
|
||||
|
||||
- name: Setup ENV (SQLite)
|
||||
run: |
|
||||
mkdir -p ${{ github.WORKSPACE }}/../app/db
|
||||
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
|
||||
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
|
||||
|
||||
- name: Run KB Tests (kb)
|
||||
run: make unit-test-kb
|
||||
|
||||
# =============================================================================
|
||||
# AI Tests (agent, aigc) - Run once with SQLite
|
||||
# =============================================================================
|
||||
|
|
|
|||
38
Makefile
38
Makefile
|
|
@ -11,10 +11,12 @@ OS := $(shell uname)
|
|||
|
||||
# ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
|
||||
TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*' | awk '!/\/tests\// || /openapi\/tests/')
|
||||
# Core tests (exclude AI-related: agent, aigc, openai)
|
||||
TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent' | awk '!/\/tests\// || /openapi\/tests/')
|
||||
# Core tests (exclude AI-related: agent, aigc, openai, and KB)
|
||||
TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb' | awk '!/\/tests\// || /openapi\/tests/')
|
||||
# AI tests (agent, aigc)
|
||||
TESTFOLDER_AI := $(shell $(GO) list ./agent/... ./aigc/...)
|
||||
# KB tests (kb)
|
||||
TESTFOLDER_KB := $(shell $(GO) list ./kb/...)
|
||||
TESTTAGS ?= ""
|
||||
|
||||
# TESTWIDGETS := $(shell $(GO) list ./widgets/...)
|
||||
|
|
@ -103,6 +105,38 @@ unit-test-ai:
|
|||
fi; \
|
||||
done
|
||||
|
||||
# KB Unit Test (kb)
|
||||
.PHONY: unit-test-kb
|
||||
unit-test-kb:
|
||||
echo "mode: count" > coverage-kb.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; \
|
||||
if grep -q "^--- FAIL" tmp.out; then \
|
||||
rm tmp.out; \
|
||||
exit 1; \
|
||||
elif grep -q "^FAIL" tmp.out; then \
|
||||
rm tmp.out; \
|
||||
exit 1; \
|
||||
elif grep -q "^panic:" tmp.out; then \
|
||||
rm tmp.out; \
|
||||
exit 1; \
|
||||
elif grep -q "build failed" tmp.out; then \
|
||||
rm tmp.out; \
|
||||
exit 1; \
|
||||
elif grep -q "setup failed" tmp.out; then \
|
||||
rm tmp.out; \
|
||||
exit 1; \
|
||||
elif grep -q "runtime error" tmp.out; then \
|
||||
rm tmp.out; \
|
||||
exit 1; \
|
||||
fi; \
|
||||
if [ -f profile.out ]; then \
|
||||
cat profile.out | grep -v "mode:" >> coverage-kb.out; \
|
||||
rm profile.out; \
|
||||
fi; \
|
||||
done
|
||||
|
||||
# Benchmark Test
|
||||
.PHONY: benchmark
|
||||
benchmark:
|
||||
|
|
|
|||
334
kb/api/README.md
Normal file
334
kb/api/README.md
Normal file
|
|
@ -0,0 +1,334 @@
|
|||
# KB API
|
||||
|
||||
The `kb/api` package provides a unified Go API for Knowledge Base operations including collection management, document ingestion, and semantic search.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```go
|
||||
import (
|
||||
"context"
|
||||
"github.com/yaoapp/yao/kb"
|
||||
"github.com/yaoapp/yao/kb/api"
|
||||
)
|
||||
|
||||
// After kb.Load(), use kb.API to access all operations
|
||||
ctx := context.Background()
|
||||
```
|
||||
|
||||
## API Interface
|
||||
|
||||
```go
|
||||
type API interface {
|
||||
// Collection operations
|
||||
CreateCollection(ctx, params) (*CreateCollectionResult, error)
|
||||
RemoveCollection(ctx, collectionID) (*RemoveCollectionResult, error)
|
||||
GetCollection(ctx, collectionID) (map[string]interface{}, error)
|
||||
CollectionExists(ctx, collectionID) (*CollectionExistsResult, error)
|
||||
ListCollections(ctx, filter) (*ListCollectionsResult, error)
|
||||
UpdateCollectionMetadata(ctx, collectionID, params) (*UpdateMetadataResult, error)
|
||||
|
||||
// Document operations
|
||||
ListDocuments(ctx, filter) (*ListDocumentsResult, error)
|
||||
GetDocument(ctx, docID, params) (map[string]interface{}, error)
|
||||
RemoveDocuments(ctx, params) (*RemoveDocumentsResult, error)
|
||||
|
||||
// Document add operations (sync)
|
||||
AddFile(ctx, params) (*AddDocumentResult, error)
|
||||
AddText(ctx, params) (*AddDocumentResult, error)
|
||||
AddURL(ctx, params) (*AddDocumentResult, error)
|
||||
|
||||
// Document add operations (async)
|
||||
AddFileAsync(ctx, params) (*AddDocumentAsyncResult, error)
|
||||
AddTextAsync(ctx, params) (*AddDocumentAsyncResult, error)
|
||||
AddURLAsync(ctx, params) (*AddDocumentAsyncResult, error)
|
||||
|
||||
// Search operations
|
||||
Search(ctx, queries) (*SearchResult, error)
|
||||
}
|
||||
```
|
||||
|
||||
## Collection Operations
|
||||
|
||||
### Create Collection
|
||||
|
||||
```go
|
||||
params := &api.CreateCollectionParams{
|
||||
ID: "my_collection",
|
||||
Metadata: map[string]interface{}{
|
||||
"name": "My Knowledge Base",
|
||||
"description": "Collection description",
|
||||
},
|
||||
EmbeddingProviderID: "__yao.openai",
|
||||
EmbeddingOptionID: "text-embedding-3-small",
|
||||
Locale: "en",
|
||||
Config: &types.CreateCollectionOptions{
|
||||
Distance: "cosine",
|
||||
IndexType: "hnsw",
|
||||
},
|
||||
}
|
||||
|
||||
result, err := kb.API.CreateCollection(ctx, params)
|
||||
// result.CollectionID = "my_collection"
|
||||
```
|
||||
|
||||
### Get Collection
|
||||
|
||||
```go
|
||||
collection, err := kb.API.GetCollection(ctx, "my_collection")
|
||||
// collection["id"], collection["name"], collection["config"], etc.
|
||||
```
|
||||
|
||||
### List Collections
|
||||
|
||||
```go
|
||||
filter := &api.ListCollectionsFilter{
|
||||
Page: 1,
|
||||
PageSize: 20,
|
||||
Keywords: "knowledge",
|
||||
Status: []string{"active"},
|
||||
}
|
||||
|
||||
result, err := kb.API.ListCollections(ctx, filter)
|
||||
// result.Data, result.Total, result.PageCnt
|
||||
```
|
||||
|
||||
### Remove Collection
|
||||
|
||||
```go
|
||||
result, err := kb.API.RemoveCollection(ctx, "my_collection")
|
||||
// result.Removed = true
|
||||
```
|
||||
|
||||
## Document Operations
|
||||
|
||||
### Add Text
|
||||
|
||||
```go
|
||||
params := &api.AddTextParams{
|
||||
CollectionID: "my_collection",
|
||||
Text: "Einstein developed the theory of relativity...",
|
||||
DocID: "einstein_bio", // optional, auto-generated if empty
|
||||
Metadata: map[string]interface{}{
|
||||
"title": "Einstein Biography",
|
||||
"author": "John Doe",
|
||||
},
|
||||
Chunking: &api.ProviderConfigParams{
|
||||
ProviderID: "__yao.structured",
|
||||
OptionID: "standard",
|
||||
},
|
||||
Embedding: &api.ProviderConfigParams{
|
||||
ProviderID: "__yao.openai",
|
||||
OptionID: "text-embedding-3-small",
|
||||
},
|
||||
Extraction: &api.ProviderConfigParams{ // optional, for graph extraction
|
||||
ProviderID: "__yao.openai",
|
||||
OptionID: "gpt-4o-mini",
|
||||
},
|
||||
}
|
||||
|
||||
result, err := kb.API.AddText(ctx, params)
|
||||
// result.DocID = "einstein_bio"
|
||||
```
|
||||
|
||||
### Add File
|
||||
|
||||
```go
|
||||
params := &api.AddFileParams{
|
||||
CollectionID: "my_collection",
|
||||
FileID: "uploaded_file_id",
|
||||
Uploader: "local", // or "s3", etc.
|
||||
Chunking: &api.ProviderConfigParams{...},
|
||||
Embedding: &api.ProviderConfigParams{...},
|
||||
}
|
||||
|
||||
result, err := kb.API.AddFile(ctx, params)
|
||||
```
|
||||
|
||||
### Add URL
|
||||
|
||||
```go
|
||||
params := &api.AddURLParams{
|
||||
CollectionID: "my_collection",
|
||||
URL: "https://example.com/article",
|
||||
Chunking: &api.ProviderConfigParams{...},
|
||||
Embedding: &api.ProviderConfigParams{...},
|
||||
}
|
||||
|
||||
result, err := kb.API.AddURL(ctx, params)
|
||||
```
|
||||
|
||||
### List Documents
|
||||
|
||||
```go
|
||||
filter := &api.ListDocumentsFilter{
|
||||
Page: 1,
|
||||
PageSize: 20,
|
||||
CollectionID: "my_collection",
|
||||
Status: []string{"active"},
|
||||
}
|
||||
|
||||
result, err := kb.API.ListDocuments(ctx, filter)
|
||||
```
|
||||
|
||||
### Remove Documents
|
||||
|
||||
```go
|
||||
params := &api.RemoveDocumentsParams{
|
||||
DocumentIDs: []string{"doc1", "doc2"},
|
||||
}
|
||||
|
||||
result, err := kb.API.RemoveDocuments(ctx, params)
|
||||
```
|
||||
|
||||
## Search Operations
|
||||
|
||||
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) |
|
||||
|
||||
### Basic Vector Search
|
||||
|
||||
```go
|
||||
queries := []api.Query{
|
||||
{
|
||||
CollectionID: "my_collection",
|
||||
Input: "What is the theory of relativity?",
|
||||
Mode: api.SearchModeVector,
|
||||
PageSize: 10,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := kb.API.Search(ctx, queries)
|
||||
// result.Segments - matched text segments with scores
|
||||
// result.Total - total count
|
||||
```
|
||||
|
||||
### Graph-Enhanced Search (Expand Mode)
|
||||
|
||||
```go
|
||||
queries := []api.Query{
|
||||
{
|
||||
CollectionID: "my_collection",
|
||||
Input: "Einstein's contributions to physics",
|
||||
Mode: api.SearchModeExpand, // default
|
||||
MaxDepth: 2, // graph traversal depth
|
||||
PageSize: 10,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := kb.API.Search(ctx, queries)
|
||||
// result.Segments - segments from vector + graph expansion
|
||||
// result.Graph.Nodes - related entities
|
||||
// result.Graph.Relationships - entity relationships
|
||||
```
|
||||
|
||||
### Multi-Query Search
|
||||
|
||||
Queries can span multiple collections; results are merged and deduplicated:
|
||||
|
||||
```go
|
||||
queries := []api.Query{
|
||||
{
|
||||
CollectionID: "science_kb",
|
||||
Input: "quantum mechanics",
|
||||
Mode: api.SearchModeVector,
|
||||
},
|
||||
{
|
||||
CollectionID: "tech_kb",
|
||||
Input: "machine learning",
|
||||
Mode: api.SearchModeVector,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := kb.API.Search(ctx, queries)
|
||||
// Merged results from both collections
|
||||
```
|
||||
|
||||
### Search with Messages (Conversation Context)
|
||||
|
||||
```go
|
||||
queries := []api.Query{
|
||||
{
|
||||
CollectionID: "my_collection",
|
||||
Messages: []types.ChatMessage{
|
||||
{Role: "user", Content: "Tell me about Einstein"},
|
||||
{Role: "assistant", Content: "Einstein was a physicist..."},
|
||||
{Role: "user", Content: "What about his discoveries?"}, // used as query
|
||||
},
|
||||
Mode: api.SearchModeExpand,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := kb.API.Search(ctx, queries)
|
||||
```
|
||||
|
||||
### Search with Filters
|
||||
|
||||
```go
|
||||
queries := []api.Query{
|
||||
{
|
||||
CollectionID: "my_collection",
|
||||
Input: "physics",
|
||||
DocumentID: "specific_doc_id", // filter to specific document
|
||||
MinScore: 0.5, // minimum similarity score
|
||||
Metadata: map[string]interface{}{
|
||||
"category": "science",
|
||||
},
|
||||
Page: 1,
|
||||
PageSize: 20,
|
||||
},
|
||||
}
|
||||
|
||||
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 |
|
||||
| `Metadata` | map | Filter by metadata fields |
|
||||
| `MaxDepth` | int | Graph traversal depth (default: 2) |
|
||||
| `Page` | int | Page number (1-based) |
|
||||
| `PageSize` | int | Results per page |
|
||||
|
||||
## Search Result
|
||||
|
||||
```go
|
||||
type SearchResult struct {
|
||||
Segments []types.Segment // Matched segments with scores
|
||||
Graph *GraphData // Nodes and relationships (graph/expand mode)
|
||||
Total int // Total results count
|
||||
Page int // Current page
|
||||
PageSize int // Results per page
|
||||
TotalPages int // Total pages
|
||||
Next int // Next page number
|
||||
Prev int // Previous page number
|
||||
}
|
||||
```
|
||||
|
||||
## Provider Configuration
|
||||
|
||||
Providers handle text processing:
|
||||
|
||||
```go
|
||||
type ProviderConfigParams struct {
|
||||
ProviderID string // e.g., "__yao.openai", "__yao.structured"
|
||||
OptionID string // e.g., "text-embedding-3-small", "gpt-4o-mini"
|
||||
}
|
||||
```
|
||||
|
||||
Common providers:
|
||||
- **Chunking**: `__yao.structured` - text splitting
|
||||
- **Embedding**: `__yao.openai` - vector embeddings
|
||||
- **Extraction**: `__yao.openai` - entity/relationship extraction for graph
|
||||
|
||||
|
|
@ -73,12 +73,17 @@ func (kb *KBInstance) Search(ctx context.Context, queries []Query) (*SearchResul
|
|||
close(errChan)
|
||||
|
||||
// Collect errors
|
||||
var errors []error
|
||||
var errs []error
|
||||
for err := range errChan {
|
||||
errors = append(errors, err)
|
||||
errs = append(errs, err)
|
||||
}
|
||||
if len(errors) > 0 {
|
||||
log.Warn("Search completed with errors: %v", errors)
|
||||
if len(errs) > 0 {
|
||||
// Combine all error messages
|
||||
errMsgs := make([]string, len(errs))
|
||||
for i, e := range errs {
|
||||
errMsgs[i] = e.Error()
|
||||
}
|
||||
return nil, fmt.Errorf("search failed: %v", errMsgs)
|
||||
}
|
||||
|
||||
// 4. Merge and deduplicate results
|
||||
|
|
|
|||
|
|
@ -14,22 +14,12 @@ import (
|
|||
// Note: Test data setup is in search_setup_test.go
|
||||
|
||||
// ========== Search Query Tests ==========
|
||||
// These tests use fixed collection IDs from search_setup_test.go
|
||||
// Run TestSearchSetup first to create test data, then run these tests.
|
||||
//
|
||||
// Usage:
|
||||
// 1. Setup (run once): go test -v -run "TestSearchSetup" ./kb/api/...
|
||||
// 2. Run tests: go test -v -run "TestSearchQuery" ./kb/api/...
|
||||
// 3. Cleanup (optional): go test -v -run "TestSearchCleanup" ./kb/api/...
|
||||
|
||||
// verifyTestDataExists checks if test collections exist, skips if not
|
||||
func verifyTestDataExists(t *testing.T, ctx context.Context) {
|
||||
scienceExists, _ := kb.API.CollectionExists(ctx, SearchTestScienceCollection)
|
||||
techExists, _ := kb.API.CollectionExists(ctx, SearchTestTechCollection)
|
||||
|
||||
if scienceExists == nil || !scienceExists.Exists || techExists == nil || !techExists.Exists {
|
||||
t.Skip("Test data not found. Run 'go test -v -run TestSearchSetup ./kb/api/...' first")
|
||||
}
|
||||
// ensureTestDataExists ensures test collections exist by running setup if needed
|
||||
// Setup will skip creation if data already exists
|
||||
func ensureTestDataExists(t *testing.T, ctx context.Context) {
|
||||
// Run setup - it checks if data exists and skips if already complete
|
||||
TestSearchSetup(t)
|
||||
}
|
||||
|
||||
func TestSearchQuery(t *testing.T) {
|
||||
|
|
@ -38,7 +28,7 @@ func TestSearchQuery(t *testing.T) {
|
|||
}
|
||||
|
||||
ctx := context.Background()
|
||||
verifyTestDataExists(t, ctx)
|
||||
ensureTestDataExists(t, ctx)
|
||||
|
||||
t.Run("VectorSearch_SingleCollection", func(t *testing.T) {
|
||||
// Test: Simple vector search in science collection
|
||||
|
|
|
|||
|
|
@ -239,7 +239,7 @@ type Query struct {
|
|||
// Metadata filters segments by metadata fields (optional)
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"`
|
||||
|
||||
// Graph search options (used when Mode is graph or hybrid)
|
||||
// Graph search options (used when Mode is graph or expand)
|
||||
MaxDepth int `json:"max_depth,omitempty" yaml:"max_depth,omitempty"` // Max traversal depth (default: 2)
|
||||
|
||||
// Pagination options
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue