Cyrene/backend/ai-core/internal/subsession/knowledge_provider.go

package subsession

import (
	"context"
	"fmt"
	"strings"
	"time"

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

// KnowledgeProvider searches the knowledge base for relevant information.
type KnowledgeProvider struct {
	retriever *rag.Retriever
}

// NewKnowledgeProvider creates a knowledge subsession provider.
func NewKnowledgeProvider(retriever *rag.Retriever) *KnowledgeProvider {
	return &KnowledgeProvider{retriever: retriever}
}

func (p *KnowledgeProvider) Type() model.SubSessionType {
	return model.SubSessionKnowledge
}

// knowledgeKeywords are trigger words from _index.md. Only run expensive embedding search if message matches.
var knowledgeKeywords = []string{
	"翁法罗斯", "泰坦", "城邦", "黑潮", "帝皇权杖", "黄金裔",
	"白厄", "阿格莱雅", "缇宝", "万敌", "那刻夏", "遐蝶", "风堇", "赛飞儿", "海瑟音", "刻律德菈",
	"哀丽秘榭", "昔涟", "星神", "浮黎", "轮回", "始源命途", "无漏净子",
	"剧情", "结局", "逐火", "盗火", "火种", "奥赫玛", "来古士",
	"世界观", "设定", "哲学", "浪漫", "哀怜", "有瑕",
}

func (p *KnowledgeProvider) CanHandle(_ context.Context, intent *model.IntentResult, userMessage string) bool {
	if intent == nil {
		return true
	}
	switch intent.Primary {
	case "knowledge", "technical", "how_to", "factual", "research":
		return true
	case "chat":
		// 仅当消息包含知识库相关关键词时才触发检索，避免每次聊天都跑 embedding
		msg := strings.ToLower(userMessage)
		for _, kw := range knowledgeKeywords {
			if strings.Contains(msg, strings.ToLower(kw)) {
				return true
			}
		}
		return false
	}
	return true
}

func (p *KnowledgeProvider) Priority() int {
	return 3
}

func (p *KnowledgeProvider) Timeout() time.Duration {
	return 15 * time.Second
}

func (p *KnowledgeProvider) CreateContext(ctx context.Context, params CreateContextParams) ([]model.LLMMessage, error) {
	return []model.LLMMessage{
		{Role: model.RoleSystem, Content: "知识库检索子会话"},
		{Role: model.RoleUser, Content: params.UserMessage},
	}, nil
}

func (p *KnowledgeProvider) Execute(ctx context.Context, subCtx []model.LLMMessage) (*model.SubSessionResult, error) {
	userMessage := ""
	for i := len(subCtx) - 1; i >= 0; i-- {
		if subCtx[i].Role == model.RoleUser {
			userMessage = subCtx[i].Content
			break
		}
	}
	if userMessage == "" {
		return nil, fmt.Errorf("无法提取用户消息")
	}

	result := &model.SubSessionResult{
		Type:       model.SubSessionKnowledge,
		Confidence: 0,
	}

	if p.retriever == nil {
		result.Summary = "（知识库未就绪）"
		return result, nil
	}

	retrieval, err := p.retriever.Retrieve(ctx, userMessage, 3)
	if err != nil {
		logger.Printf("[knowledge-subsession] 知识检索失败: %v", err)
		result.Error = fmt.Sprintf("检索失败: %v", err)
		result.Summary = "（知识库检索失败）"
		return result, nil
	}

	if len(retrieval.Results) == 0 {
		result.Summary = "（未找到相关知识）"
		return result, nil
	}

	result.Summary = retrieval.Summary
	result.Confidence = 0.6
	logger.Printf("[knowledge-subsession] 完成: 找到 %d 条知识", len(retrieval.Results))
	return result, nil
}