CyberStrikeAI/internal/knowledge/types.go

package knowledge

import (
	"encoding/json"
	"time"
)

// formatTime 格式化时间为 RFC3339 格式，零时间返回空字符串
func formatTime(t time.Time) string {
	if t.IsZero() {
		return ""
	}
	return t.Format(time.RFC3339)
}

// KnowledgeItem 知识库项
type KnowledgeItem struct {
	ID        string    `json:"id"`
	Category  string    `json:"category"` // 风险类型（文件夹名）
	Title     string    `json:"title"`    // 标题（文件名）
	FilePath  string    `json:"filePath"` // 文件路径
	Content   string    `json:"content"`  // 文件内容
	CreatedAt time.Time `json:"createdAt"`
	UpdatedAt time.Time `json:"updatedAt"`
}

// KnowledgeItemSummary 知识库项摘要（用于列表，不包含完整内容）
type KnowledgeItemSummary struct {
	ID        string    `json:"id"`
	Category  string    `json:"category"`
	Title     string    `json:"title"`
	FilePath  string    `json:"filePath"`
	Content   string    `json:"content,omitempty"` // 可选：内容预览（如果提供，通常只包含前 150 字符）
	CreatedAt time.Time `json:"createdAt"`
	UpdatedAt time.Time `json:"updatedAt"`
}

// MarshalJSON 自定义 JSON 序列化，确保时间格式正确
func (k *KnowledgeItemSummary) MarshalJSON() ([]byte, error) {
	type Alias KnowledgeItemSummary
	aux := &struct {
		*Alias
		CreatedAt string `json:"createdAt"`
		UpdatedAt string `json:"updatedAt"`
	}{
		Alias: (*Alias)(k),
	}
	aux.CreatedAt = formatTime(k.CreatedAt)
	aux.UpdatedAt = formatTime(k.UpdatedAt)
	return json.Marshal(aux)
}

// MarshalJSON 自定义 JSON 序列化，确保时间格式正确
func (k *KnowledgeItem) MarshalJSON() ([]byte, error) {
	type Alias KnowledgeItem
	aux := &struct {
		*Alias
		CreatedAt string `json:"createdAt"`
		UpdatedAt string `json:"updatedAt"`
	}{
		Alias: (*Alias)(k),
	}
	aux.CreatedAt = formatTime(k.CreatedAt)
	aux.UpdatedAt = formatTime(k.UpdatedAt)
	return json.Marshal(aux)
}

// KnowledgeChunk 知识块（用于向量化）
type KnowledgeChunk struct {
	ID         string    `json:"id"`
	ItemID     string    `json:"itemId"`
	ChunkIndex int       `json:"chunkIndex"`
	ChunkText  string    `json:"chunkText"`
	Embedding  []float32 `json:"-"` // 向量嵌入，不序列化到 JSON
	CreatedAt  time.Time `json:"createdAt"`
}

// RetrievalResult 检索结果
type RetrievalResult struct {
	Chunk      *KnowledgeChunk `json:"chunk"`
	Item       *KnowledgeItem  `json:"item"`
	Similarity float64         `json:"similarity"` // 相似度分数
	Score      float64         `json:"score"`      // 综合分数（混合检索）
}

// RetrievalLog 检索日志
type RetrievalLog struct {
	ID             string    `json:"id"`
	ConversationID string    `json:"conversationId,omitempty"`
	MessageID      string    `json:"messageId,omitempty"`
	Query          string    `json:"query"`
	RiskType       string    `json:"riskType,omitempty"`
	RetrievedItems []string  `json:"retrievedItems"` // 检索到的知识项 ID 列表
	CreatedAt      time.Time `json:"createdAt"`
}

// MarshalJSON 自定义 JSON 序列化，确保时间格式正确
func (r *RetrievalLog) MarshalJSON() ([]byte, error) {
	type Alias RetrievalLog
	return json.Marshal(&struct {
		*Alias
		CreatedAt string `json:"createdAt"`
	}{
		Alias:     (*Alias)(r),
		CreatedAt: formatTime(r.CreatedAt),
	})
}

// CategoryWithItems 分类及其下的知识项（用于按分类分页）
type CategoryWithItems struct {
	Category  string                `json:"category"`           // 分类名称
	ItemCount int                   `json:"itemCount"`          // 该分类下的知识项总数
	Items     []*KnowledgeItemSummary `json:"items"`            // 该分类下的知识项列表
}

// SearchRequest 搜索请求
type SearchRequest struct {
	Query     string  `json:"query"`
	RiskType  string  `json:"riskType,omitempty"`  // 可选：指定风险类型
	TopK      int     `json:"topK,omitempty"`      // 返回 Top-K 结果，默认 5
	Threshold float64 `json:"threshold,omitempty"` // 相似度阈值，默认 0.7
}