wizardchen
8ca9c117bb
- Updated the listFAQEntries function to support multiple tag IDs for filtering. - Modified the frontend components to accommodate the new tag_ids parameter. - Enhanced the KbTagManageDrawer and FAQEntryManager components to reflect changes in tag handling. - Improved API documentation to clarify the new tag filtering capabilities. This commit enhances the flexibility of FAQ entry retrieval, allowing users to filter by multiple tags for more precise results.
115 lines
3.2 KiB
Go
115 lines
3.2 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
|
"github.com/Tencent/WeKnora/internal/utils"
|
|
)
|
|
|
|
var dataSchemaTool = BaseTool{
|
|
name: ToolDataSchema,
|
|
description: "Use this tool to get the schema information of a CSV or Excel file loaded into DuckDB. It returns the table name, columns, and row count.",
|
|
schema: utils.GenerateSchema[DataSchemaInput](),
|
|
}
|
|
|
|
type DataSchemaInput struct {
|
|
KnowledgeID string `json:"knowledge_id" jsonschema:"short dN document ID to query"`
|
|
}
|
|
|
|
type DataSchemaTool struct {
|
|
BaseTool
|
|
knowledgeService interfaces.KnowledgeService
|
|
chunkRepo interfaces.ChunkRepository
|
|
targetChunkTypes []types.ChunkType
|
|
}
|
|
|
|
func NewDataSchemaTool(knowledgeService interfaces.KnowledgeService, chunkRepo interfaces.ChunkRepository, targetChunkTypes ...types.ChunkType) *DataSchemaTool {
|
|
if len(targetChunkTypes) == 0 {
|
|
targetChunkTypes = []types.ChunkType{types.ChunkTypeTableSummary, types.ChunkTypeTableColumn}
|
|
}
|
|
return &DataSchemaTool{
|
|
BaseTool: dataSchemaTool,
|
|
knowledgeService: knowledgeService,
|
|
chunkRepo: chunkRepo,
|
|
targetChunkTypes: targetChunkTypes,
|
|
}
|
|
}
|
|
|
|
// Execute executes the tool logic
|
|
func (t *DataSchemaTool) Execute(ctx context.Context, args json.RawMessage) (*types.ToolResult, error) {
|
|
var input DataSchemaInput
|
|
if err := json.Unmarshal(args, &input); err != nil {
|
|
return &types.ToolResult{
|
|
Success: false,
|
|
Error: fmt.Sprintf("Failed to parse input args: %v", err),
|
|
}, err
|
|
}
|
|
|
|
// Get knowledge to get TenantID (use IDOnly to support cross-tenant shared KB)
|
|
knowledge, err := t.knowledgeService.GetKnowledgeByIDOnly(ctx, input.KnowledgeID)
|
|
if err != nil {
|
|
return &types.ToolResult{
|
|
Success: false,
|
|
Error: fmt.Sprintf("Failed to get knowledge '%s': %v", input.KnowledgeID, err),
|
|
}, err
|
|
}
|
|
|
|
// Get chunks for the knowledge ID using ChunkRepository
|
|
// We only need table summary and column chunks
|
|
chunkTypes := t.targetChunkTypes
|
|
page := &types.Pagination{
|
|
Page: 1,
|
|
PageSize: 100, // Should be enough for schema chunks
|
|
}
|
|
|
|
chunks, _, err := t.chunkRepo.ListPagedChunksByKnowledgeID(
|
|
ctx,
|
|
knowledge.TenantID,
|
|
input.KnowledgeID,
|
|
page,
|
|
chunkTypes,
|
|
nil, // tagIDs
|
|
"", // keyword
|
|
"", // searchField
|
|
"", // sortOrder
|
|
"", // knowledgeType
|
|
)
|
|
if err != nil {
|
|
return &types.ToolResult{
|
|
Success: false,
|
|
Error: fmt.Sprintf("Failed to list chunks for knowledge ID '%s': %v", input.KnowledgeID, err),
|
|
}, err
|
|
}
|
|
|
|
var summaryContent, columnContent string
|
|
for _, chunk := range chunks {
|
|
if chunk.ChunkType == types.ChunkTypeTableSummary {
|
|
summaryContent = chunk.Content
|
|
} else if chunk.ChunkType == types.ChunkTypeTableColumn {
|
|
columnContent = chunk.Content
|
|
}
|
|
}
|
|
|
|
if summaryContent == "" || columnContent == "" {
|
|
return &types.ToolResult{
|
|
Success: false,
|
|
Error: fmt.Sprintf("No table schema information found for knowledge ID '%s'", input.KnowledgeID),
|
|
}, fmt.Errorf("no schema info found")
|
|
}
|
|
|
|
output := fmt.Sprintf("%s\n\n%s", summaryContent, columnContent)
|
|
|
|
return &types.ToolResult{
|
|
Success: true,
|
|
Output: output,
|
|
Data: map[string]interface{}{
|
|
"summary": summaryContent,
|
|
"columns": columnContent,
|
|
},
|
|
}, nil
|
|
}
|