MyAiDesk/GovAI/server/internal/handler/chat_llm_test.go

package handler

import (
	"strings"
	"testing"
)

// TestEnhanceCitations 测试来源标注增强功能
func TestEnhanceCitations(t *testing.T) {
	h := &LLMChatHandler{}

	tests := []struct {
		name             string
		response         string
		hasKnowledge     bool
		knowledgeSources []string
		expectedContains []string
	}{
		{
			name: "清单格式 - 应该为每个问题添加标注",
			response: `劳动用工合规审查清单
一、劳动合同管理
劳动合同签订

是否所有员工均签订了书面劳动合同？
合同内容是否符合《劳动合同法》的相关规定？
是否存在劳动合同未明确约定工作地点、工作内容、劳动报酬等情况？`,
			hasKnowledge:     true,
			knowledgeSources: []string{"劳动合同法"},
			expectedContains: []string{
				"是否所有员工均签订了书面劳动合同？ [[",
				"合同内容是否符合《劳动合同法》的相关规定？ [[",
				"是否存在劳动合同未明确约定工作地点、工作内容、劳动报酬等情况？ [[",
			},
		},
		{
			name: "列表项 - 应该为每个列表项添加标注",
			response: `办理条件：
- 有合法稳定就业
- 有合法稳定住所
- 连续就读`,
			hasKnowledge:     true,
			knowledgeSources: []string{"居住证管理规定"},
			expectedContains: []string{
				"- 有合法稳定就业 [[",
				"- 有合法稳定住所 [[",
				"- 连续就读 [[",
			},
		},
		{
			name: "数字列表 - 应该为每个列表项添加标注",
			response: `办理流程：
1. 到派出所提交申请材料
2. 派出所审核材料
3. 制作并发放居住证`,
			hasKnowledge:     true,
			knowledgeSources: []string{"居住证管理规定"},
			expectedContains: []string{
				"1. 到派出所提交申请材料 [[",
				"2. 派出所审核材料 [[",
				"3. 制作并发放居住证 [[",
			},
		},
		{
			name: "陈述句 - 应该添加标注",
			response: `居住证办理需要身份证、居住证明和近期照片。
办理时限为15个工作日。`,
			hasKnowledge:     true,
			knowledgeSources: []string{"居住证管理规定"},
			expectedContains: []string{
				"居住证办理需要身份证、居住证明和近期照片。 [[",
				"办理时限为15个工作日。 [[",
			},
		},
		{
			name: "无知识库 - 应该添加AI建议标注",
			response: `建议您提前准备齐全材料。
如有疑问可先电话咨询。`,
			hasKnowledge:     false,
			knowledgeSources: []string{},
			expectedContains: []string{
				"建议您提前准备齐全材料。 [[AI建议]]",
				"如有疑问可先电话咨询。 [[AI建议]]",
			},
		},
		{
			name: "已有标注 - 不应该重复添加",
			response: `居住证办理需要身份证、居住证明和近期照片 [[知识库:居住证管理规定]]。
建议您提前准备齐全材料 [[AI建议]]。`,
			hasKnowledge:     true,
			knowledgeSources: []string{"居住证管理规定"},
			expectedContains: []string{
				"[[知识库:居住证管理规定]]",
				"[[AI建议]]",
				"**来源说明**",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := h.enhanceCitations(tt.response, tt.hasKnowledge, tt.knowledgeSources)

			for _, expected := range tt.expectedContains {
				if !strings.Contains(result, expected) {
					t.Errorf("期望结果包含: %q\n实际结果: %s", expected, result)
				}
			}

			// 验证来源说明块存在
			if !strings.Contains(result, "**来源说明**") {
				t.Errorf("期望结果包含来源说明块")
			}
		})
	}
}

// TestAddCitationsToResponse 测试自动添加标注的逻辑
func TestAddCitationsToResponse(t *testing.T) {
	h := &LLMChatHandler{}

	response := `劳动合同管理
是否所有员工均签订了书面劳动合同？
合同内容是否符合《劳动合同法》的相关规定？`

	result := h.addCitationsToResponse(response, true, []string{"劳动合同法"})

	// 验证每个问题都被标注了
	if !strings.Contains(result, "是否所有员工均签订了书面劳动合同？ [[") {
		t.Errorf("第一个问题没有被标注: %s", result)
	}

	if !strings.Contains(result, "合同内容是否符合《劳动合同法》的相关规定？ [[") {
		t.Errorf("第二个问题没有被标注: %s", result)
	}
}