Cyrene/backend/ai-core/internal/tools/knowledge_tool.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"

	"git.yeij.top/AskaEth/Cyrene/ai-core/internal/rag"
)

// KnowledgeSearchTool searches the knowledge base.
type KnowledgeSearchTool struct {
	retriever *rag.Retriever
}

// NewKnowledgeSearchTool creates a knowledge search tool.
func NewKnowledgeSearchTool(retriever *rag.Retriever) *KnowledgeSearchTool {
	return &KnowledgeSearchTool{retriever: retriever}
}

func (t *KnowledgeSearchTool) Definition() ToolDefinition {
	return ToolDefinition{
		Name:        "knowledge_search",
		Description: "搜索本地知识库。从文档、代码、笔记等中检索相关信息，支持语义搜索和关键词匹配。",
		Parameters: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"query": map[string]interface{}{
					"type":        "string",
					"description": "搜索查询",
				},
				"top_k": map[string]interface{}{
					"type":        "integer",
					"description": "返回结果数量，默认5条，最大10条",
				},
			},
			"required": []string{"query"},
		},
	}
}

func (t *KnowledgeSearchTool) Execute(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
	query, _ := args["query"].(string)
	if query == "" {
		return &ToolResult{
			ToolName: "knowledge_search",
			Success:  false,
			Error:    "query 参数不能为空",
		}, nil
	}

	topK := 5
	if v, ok := args["top_k"].(float64); ok {
		topK = int(v)
		if topK > 10 {
			topK = 10
		}
	}

	result, err := t.retriever.Retrieve(ctx, query, topK)
	if err != nil {
		return &ToolResult{
			ToolName: "knowledge_search",
			Success:  false,
			Error:    fmt.Sprintf("知识库搜索失败: %v", err),
		}, nil
	}

	count := 0
	if result.Results != nil {
		count = len(result.Results)
	}

	data, _ := json.Marshal(map[string]interface{}{
		"query":   result.Query,
		"summary": result.Summary,
		"count":   count,
	})

	return &ToolResult{
		ToolName: "knowledge_search",
		Success:  true,
		Data:     string(data),
	}, nil
}

// KnowledgeIngestTool allows ingesting documents into the knowledge base.
type KnowledgeIngestTool struct {
	store *rag.KnowledgeStore
}

// NewKnowledgeIngestTool creates a knowledge ingestion tool.
func NewKnowledgeIngestTool(store *rag.KnowledgeStore) *KnowledgeIngestTool {
	return &KnowledgeIngestTool{store: store}
}

func (t *KnowledgeIngestTool) Definition() ToolDefinition {
	return ToolDefinition{
		Name:        "knowledge_ingest",
		Description: "将文件导入知识库。支持 .md .txt .go .py .js .ts .json 等常见文件格式。",
		Parameters: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"path": map[string]interface{}{
					"type":        "string",
					"description": "文件路径或目录路径",
				},
			},
			"required": []string{"path"},
		},
	}
}

func (t *KnowledgeIngestTool) Execute(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
	path, _ := args["path"].(string)
	if path == "" {
		return &ToolResult{
			ToolName: "knowledge_ingest",
			Success:  false,
			Error:    "path 参数不能为空",
		}, nil
	}

	count, err := t.store.IngestFile(ctx, path)
	if err != nil {
		return &ToolResult{
			ToolName: "knowledge_ingest",
			Success:  false,
			Error:    fmt.Sprintf("知识导入失败: %v", err),
		}, nil
	}

	if count == 0 {
		// Try directory
		count, err = t.store.IngestDirectory(ctx)
		if err != nil {
			return &ToolResult{
				ToolName: "knowledge_ingest",
				Success:  false,
				Error:    fmt.Sprintf("目录导入失败: %v", err),
			}, nil
		}
	}

	data, _ := json.Marshal(map[string]interface{}{
		"path":           path,
		"chunks_indexed": count,
		"status":         "ok",
	})

	return &ToolResult{
		ToolName: "knowledge_ingest",
		Success:  true,
		Data:     string(data),
	}, nil
}