c
145 lines
4.1 KiB
Go
145 lines
4.1 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"testing"
|
|
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
|
"github.com/hibiken/asynq"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type stubKnowledgeBaseService struct {
|
|
kb *types.KnowledgeBase
|
|
results []*types.SearchResult
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) CreateKnowledgeBase(context.Context, *types.KnowledgeBase) (*types.KnowledgeBase, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) {
|
|
return s.kb, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) GetKnowledgeBaseByIDOnly(context.Context, string) (*types.KnowledgeBase, error) {
|
|
return s.kb, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) GetKnowledgeBasesByIDsOnly(context.Context, []string) ([]*types.KnowledgeBase, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) FillKnowledgeBaseCounts(context.Context, *types.KnowledgeBase) error {
|
|
return nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) ListKnowledgeBases(context.Context) ([]*types.KnowledgeBase, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) ListKnowledgeBasesByTenantID(context.Context, uint64) ([]*types.KnowledgeBase, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) UpdateKnowledgeBase(
|
|
context.Context,
|
|
string,
|
|
string,
|
|
string,
|
|
*types.KnowledgeBaseConfig,
|
|
) (*types.KnowledgeBase, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) DeleteKnowledgeBase(context.Context, string) error {
|
|
return nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) TogglePinKnowledgeBase(context.Context, string) (*types.KnowledgeBase, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) HybridSearch(context.Context, string, types.SearchParams) ([]*types.SearchResult, error) {
|
|
return s.results, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) GetQueryEmbedding(context.Context, string, string) ([]float32, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) ResolveEmbeddingModelKeys(context.Context, []*types.KnowledgeBase) map[string]string {
|
|
return nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) CopyKnowledgeBase(
|
|
context.Context,
|
|
string,
|
|
string,
|
|
) (*types.KnowledgeBase, *types.KnowledgeBase, error) {
|
|
return nil, nil, nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) GetRepository() interfaces.KnowledgeBaseRepository {
|
|
return nil
|
|
}
|
|
|
|
func (s *stubKnowledgeBaseService) ProcessKBDelete(context.Context, *asynq.Task) error {
|
|
return nil
|
|
}
|
|
|
|
func TestQueryKnowledgeGraph_ReportsConfiguredEntityAndRelationTypes(t *testing.T) {
|
|
tool := NewQueryKnowledgeGraphTool(&stubKnowledgeBaseService{
|
|
kb: &types.KnowledgeBase{
|
|
ID: "kb-1",
|
|
ExtractConfig: &types.ExtractConfig{
|
|
Enabled: true,
|
|
Nodes: []*types.GraphNode{
|
|
{Name: "合同"},
|
|
{Name: "部门"},
|
|
},
|
|
Relations: []*types.GraphRelation{
|
|
{Type: "属于"},
|
|
{Type: "管理"},
|
|
},
|
|
},
|
|
},
|
|
results: []*types.SearchResult{
|
|
{
|
|
ID: "chunk-1",
|
|
Content: "合同管理流程由法务部门负责。",
|
|
KnowledgeID: "doc-1",
|
|
KnowledgeTitle: "合同管理制度",
|
|
Score: 0.91,
|
|
},
|
|
},
|
|
})
|
|
|
|
args, err := json.Marshal(QueryKnowledgeGraphInput{
|
|
KnowledgeBaseIDs: []string{"kb-1"},
|
|
Query: "合同管理",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
result, err := tool.Execute(context.Background(), args)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result)
|
|
require.True(t, result.Success)
|
|
t.Logf("tool output:\n%s", result.Output)
|
|
|
|
assert.Contains(t, result.Output, "Entity Types (2)")
|
|
assert.Contains(t, result.Output, "Relationship Types (2)")
|
|
assert.NotContains(t, result.Output, "No entity types configured")
|
|
assert.NotContains(t, result.Output, "No relationship types configured")
|
|
assert.Contains(t, result.Output, "合同")
|
|
assert.Contains(t, result.Output, "管理")
|
|
|
|
graphConfig, ok := result.Data["graph_config"].(map[string]interface{})
|
|
require.True(t, ok)
|
|
assert.ElementsMatch(t, []string{"合同", "部门"}, graphConfig["nodes"])
|
|
assert.ElementsMatch(t, []string{"属于", "管理"}, graphConfig["relations"])
|
|
}
|