chore: import clean Agent Desk snapshot (2/7)
This commit is contained in:
1250
internal/ai/application/runtime/agent_loop_engine.go
Normal file
1250
internal/ai/application/runtime/agent_loop_engine.go
Normal file
File diff suppressed because it is too large
Load Diff
864
internal/ai/application/runtime/agent_loop_engine_test.go
Normal file
864
internal/ai/application/runtime/agent_loop_engine_test.go
Normal file
@@ -0,0 +1,864 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
ai "agent-desk/internal/ai"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/dto/request"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
|
||||||
|
"github.com/glebarez/sqlite"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
"gorm.io/gorm"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAgentLoopActivatesSkillInsideSameToolLoop(t *testing.T) {
|
||||||
|
skill := models.SkillDefinition{
|
||||||
|
ID: 7, Name: "退款说明", Instruction: "只根据退款政策回答。",
|
||||||
|
ToolWhitelist: `["builtin/knowledge_retrieve"]`, Status: enums.StatusOk,
|
||||||
|
}
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"skill/7"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(""),
|
||||||
|
Skills: map[int64]models.SkillDefinition{skill.ID: skill},
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
execute := NewAgentLoopEngine().toolSearchExecutor(RunInput{}, turn, &state, &calls)
|
||||||
|
|
||||||
|
result, err := execute(context.Background(), ai.ToolCall{
|
||||||
|
Name: "tool_search", Arguments: `{"toolCode":"skill/7","arguments":{}}`,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("activate Skill: %v", err)
|
||||||
|
}
|
||||||
|
if state.SkillContext.SkillID() != skill.ID || !strings.Contains(result, skill.Instruction) {
|
||||||
|
t.Fatalf("Skill was not activated in the Agent Loop: state=%#v result=%q", state, result)
|
||||||
|
}
|
||||||
|
if len(calls) != 1 || calls[0].ToolCode != "skill/7" || calls[0].Status != "completed" {
|
||||||
|
t.Fatalf("unexpected Skill audit: %#v", calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopRegistersDirectCapabilityAliases(t *testing.T) {
|
||||||
|
turn := agentLoopTurn{AllowedTools: []string{
|
||||||
|
"builtin/conversation_context",
|
||||||
|
"graph/triage_service_request",
|
||||||
|
"graph/triage_service_request",
|
||||||
|
"workflow/47",
|
||||||
|
}}
|
||||||
|
definitions := agentLoopToolDefinitions(turn)
|
||||||
|
names := make(map[string]bool, len(definitions))
|
||||||
|
for _, definition := range definitions {
|
||||||
|
names[definition.Name] = true
|
||||||
|
}
|
||||||
|
for _, expected := range []string{
|
||||||
|
"tool_search",
|
||||||
|
"builtin_conversation_context",
|
||||||
|
"graph_triage_service_request",
|
||||||
|
"workflow_47",
|
||||||
|
} {
|
||||||
|
if !names[expected] {
|
||||||
|
t.Fatalf("missing registered function alias %q: %#v", expected, definitions)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConversationDecisionIsStructuredAndValidated(t *testing.T) {
|
||||||
|
decision, err := parseConversationDecision(`{"action":"handoff","reason":"customer requested a human","reply":"","handoffInitiator":"customer","handoffConfirmed":true}`)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("parse handoff decision: %v", err)
|
||||||
|
}
|
||||||
|
if decision.Action != ConversationActionHandoff || decision.Reason != "customer requested a human" {
|
||||||
|
t.Fatalf("unexpected handoff decision: %#v", decision)
|
||||||
|
}
|
||||||
|
for _, raw := range []string{
|
||||||
|
`{"action":"unknown","reason":"x","reply":"x","handoffInitiator":"none","handoffConfirmed":false}`,
|
||||||
|
`{"action":"reply","reason":"x","reply":"","handoffInitiator":"none","handoffConfirmed":false}`,
|
||||||
|
`{"action":"ask_handoff_confirmation","reason":"x","reply":"confirm?","handoffInitiator":"customer","handoffConfirmed":false}`,
|
||||||
|
} {
|
||||||
|
if _, err := parseConversationDecision(raw); err == nil {
|
||||||
|
t.Fatalf("expected invalid decision to fail: %s", raw)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopRecordsConversationDecision(t *testing.T) {
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
execute := NewAgentLoopEngine().toolSearchExecutor(RunInput{}, agentLoopTurn{}, &state, &calls)
|
||||||
|
if _, err := execute(context.Background(), ai.ToolCall{
|
||||||
|
Name: "conversation_decision", Arguments: `{"action":"handoff","reason":"customer requested a human","reply":"","handoffInitiator":"customer","handoffConfirmed":true}`,
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("record decision: %v", err)
|
||||||
|
}
|
||||||
|
if state.Decision == nil || state.Decision.Action != ConversationActionHandoff {
|
||||||
|
t.Fatalf("decision was not stored: %#v", state.Decision)
|
||||||
|
}
|
||||||
|
if len(calls) != 1 || calls[0].ToolCode != "conversation_decision" || calls[0].Status != "completed" {
|
||||||
|
t.Fatalf("decision audit missing: %#v", calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveAgentLoopReplyKeepsNormalModelReplyWithoutDecision(t *testing.T) {
|
||||||
|
reply, handoff, reason, err := resolveAgentLoopReply("你好,有什么可以帮你?", nil)
|
||||||
|
if err != nil || handoff || reason != "" || reply != "你好,有什么可以帮你?" {
|
||||||
|
t.Fatalf("unexpected normal reply resolution: reply=%q handoff=%t reason=%q err=%v", reply, handoff, reason, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveAgentLoopReplyDetectsModelHandoffClaimWithoutDecision(t *testing.T) {
|
||||||
|
reply, handoff, reason, err := resolveAgentLoopReply("Перевожу вас на оператора, пожалуйста подождите.", nil)
|
||||||
|
if err != nil || !handoff || reply != "" || reason != "model indicated handoff in reply text" {
|
||||||
|
t.Fatalf("expected handoff from model claim text: reply=%q handoff=%t reason=%q err=%v", reply, handoff, reason, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
replyEn, handoffEn, reasonEn, errEn := resolveAgentLoopReply("Connecting you to a human representative right now.", nil)
|
||||||
|
if errEn != nil || !handoffEn || replyEn != "" || reasonEn != "model indicated handoff in reply text" {
|
||||||
|
t.Fatalf("expected handoff from English model claim text: reply=%q handoff=%t reason=%q err=%v", replyEn, handoffEn, reasonEn, errEn)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveAgentLoopReplyFallbackOnUserHandoffRequestWithoutDecision(t *testing.T) {
|
||||||
|
reply, handoff, reason, err := resolveAgentLoopReply("Чем я могу помочь?", nil, "соедините со специалистом")
|
||||||
|
if err != nil || !handoff || reply != "" || reason != "customer requested human support" {
|
||||||
|
t.Fatalf("expected handoff fallback from user handoff request: reply=%q handoff=%t reason=%q err=%v", reply, handoff, reason, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveAgentLoopReplyUsesStructuredHandoffDecision(t *testing.T) {
|
||||||
|
reply, handoff, reason, err := resolveAgentLoopReply("模型自由文本不应生效", &ConversationDecision{
|
||||||
|
Action: ConversationActionHandoff, Reason: "customer requested human support", HandoffInitiator: HandoffInitiatorCustomer, HandoffConfirmed: true,
|
||||||
|
})
|
||||||
|
if err != nil || !handoff || reply != "" || reason != "customer requested human support" {
|
||||||
|
t.Fatalf("unexpected handoff resolution: reply=%q handoff=%t reason=%q err=%v", reply, handoff, reason, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNormalizeAgentLoopReplyAllowsEmptyInternalHandoff(t *testing.T) {
|
||||||
|
reply, err := normalizeAgentLoopReply("", true)
|
||||||
|
if err != nil || reply != "" {
|
||||||
|
t.Fatalf("handoff must bypass customer reply normalization: reply=%q err=%v", reply, err)
|
||||||
|
}
|
||||||
|
if _, err := normalizeAgentLoopReply("", false); err == nil {
|
||||||
|
t.Fatal("ordinary empty model replies must still be rejected")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopDirectCapabilityAliasUsesSamePolicyBoundary(t *testing.T) {
|
||||||
|
skill := models.SkillDefinition{
|
||||||
|
ID: 7, Name: "售后升级处理", Instruction: "先确认升级诉求。", Status: enums.StatusOk,
|
||||||
|
}
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"skill/7"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(""),
|
||||||
|
Skills: map[int64]models.SkillDefinition{skill.ID: skill},
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
execute := NewAgentLoopEngine().toolSearchExecutor(RunInput{}, turn, &state, &calls)
|
||||||
|
|
||||||
|
result, err := execute(context.Background(), ai.ToolCall{Name: "skill/7", Arguments: `{}`})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("execute direct capability alias: %v", err)
|
||||||
|
}
|
||||||
|
if state.SkillContext.SkillID() != skill.ID || !strings.Contains(result, skill.Instruction) {
|
||||||
|
t.Fatalf("direct capability was not routed through Skill activation: state=%#v result=%q", state, result)
|
||||||
|
}
|
||||||
|
if len(calls) != 1 || calls[0].ToolCode != "skill/7" || calls[0].Status != "completed" {
|
||||||
|
t.Fatalf("unexpected direct capability audit: %#v", calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopDirectSanitizedCapabilityAliasRoutesCorrectly(t *testing.T) {
|
||||||
|
skill := models.SkillDefinition{
|
||||||
|
ID: 7, Name: "售后升级处理", Instruction: "先确认升级诉求。", Status: enums.StatusOk,
|
||||||
|
}
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"skill/7"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(""),
|
||||||
|
Skills: map[int64]models.SkillDefinition{skill.ID: skill},
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
execute := NewAgentLoopEngine().toolSearchExecutor(RunInput{}, turn, &state, &calls)
|
||||||
|
|
||||||
|
result, err := execute(context.Background(), ai.ToolCall{Name: "skill_7", Arguments: `{}`})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("execute sanitized direct capability alias: %v", err)
|
||||||
|
}
|
||||||
|
if state.SkillContext.SkillID() != skill.ID || !strings.Contains(result, skill.Instruction) {
|
||||||
|
t.Fatalf("sanitized direct capability was not routed through Skill activation: state=%#v result=%q", state, result)
|
||||||
|
}
|
||||||
|
if len(calls) != 1 || calls[0].ToolCode != "skill/7" || calls[0].Status != "completed" {
|
||||||
|
t.Fatalf("unexpected direct capability audit: %#v", calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopInterruptsBeforeWriteMCPTool(t *testing.T) {
|
||||||
|
configured, err := json.Marshal([]request.AIAgentMCPToolRequest{{
|
||||||
|
ToolCode: "crm/update_customer", ServerCode: "crm", ToolName: "update_customer",
|
||||||
|
Title: "更新客户", RiskLevel: "write", RequireConfirmation: true,
|
||||||
|
}})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal MCP configuration: %v", err)
|
||||||
|
}
|
||||||
|
runInput := RunInput{
|
||||||
|
Conversation: models.Conversation{ID: 9},
|
||||||
|
AIAgent: models.AIAgent{AllowedMCPTools: string(configured)},
|
||||||
|
}
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"crm/update_customer"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(`{"allowedRiskLevels":["read","write"]}`),
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
execute := NewAgentLoopEngine().toolSearchExecutor(runInput, turn, &state, &calls)
|
||||||
|
|
||||||
|
_, err = execute(context.Background(), ai.ToolCall{
|
||||||
|
Name: "tool_search", Arguments: `{"toolCode":"crm/update_customer","arguments":{"name":"Ada"}}`,
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected write MCP Tool to interrupt")
|
||||||
|
}
|
||||||
|
if state.Interrupted == nil || !state.Interrupted.Interrupted || !strings.HasPrefix(state.Interrupted.CheckPointID, "tool:9:") {
|
||||||
|
t.Fatalf("missing MCP confirmation checkpoint: %#v", state.Interrupted)
|
||||||
|
}
|
||||||
|
if state.Interrupted.ReplyText != "即将执行“更新客户”,是否确认继续?" ||
|
||||||
|
len(state.Interrupted.Interrupts) != 1 ||
|
||||||
|
state.Interrupted.Interrupts[0].PromptText != state.Interrupted.ReplyText {
|
||||||
|
t.Fatalf("unexpected customer confirmation prompt: %#v", state.Interrupted)
|
||||||
|
}
|
||||||
|
if len(calls) != 1 || calls[0].RiskLevel != "write" || !calls[0].RequireConfirm || calls[0].Status != "interrupted" {
|
||||||
|
t.Fatalf("unexpected MCP safety audit: %#v", calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopRejectsWriteMCPBeforeConfirmationWhenRiskIsNotAllowed(t *testing.T) {
|
||||||
|
configured, _ := json.Marshal([]request.AIAgentMCPToolRequest{{
|
||||||
|
ToolCode: "crm/update_customer", ServerCode: "crm", ToolName: "update_customer",
|
||||||
|
Title: "更新客户", RiskLevel: "write", RequireConfirmation: true,
|
||||||
|
}})
|
||||||
|
runInput := RunInput{
|
||||||
|
Conversation: models.Conversation{ID: 9},
|
||||||
|
AIAgent: models.AIAgent{AllowedMCPTools: string(configured)},
|
||||||
|
}
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"crm/update_customer"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(`{"allowedRiskLevels":["read"]}`),
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
execute := NewAgentLoopEngine().toolSearchExecutor(runInput, turn, &state, &calls)
|
||||||
|
|
||||||
|
_, err := execute(context.Background(), ai.ToolCall{
|
||||||
|
Name: "tool_search", Arguments: `{"toolCode":"crm/update_customer","arguments":{"name":"Ada"}}`,
|
||||||
|
})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "risk level") {
|
||||||
|
t.Fatalf("expected MCP risk policy rejection, got %v", err)
|
||||||
|
}
|
||||||
|
if state.Interrupted != nil || len(calls) != 1 || calls[0].Status != "failed" {
|
||||||
|
t.Fatalf("disallowed MCP call should fail without a checkpoint: state=%#v calls=%#v", state, calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopRejectsWorkflowWhenWriteRiskIsNotAllowed(t *testing.T) {
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"workflow/23"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(`{"allowedRiskLevels":["read"]}`),
|
||||||
|
Workflows: map[int64]svc.AgentRevisionWorkflowBinding{
|
||||||
|
23: {WorkflowVersionID: 23, ToolName: "创建工单"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
execute := NewAgentLoopEngine().toolSearchExecutor(RunInput{}, turn, &state, &calls)
|
||||||
|
|
||||||
|
_, err := execute(context.Background(), ai.ToolCall{
|
||||||
|
Name: "tool_search", Arguments: `{"toolCode":"workflow/23","arguments":{}}`,
|
||||||
|
})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "risk level") {
|
||||||
|
t.Fatalf("expected Workflow risk policy rejection, got %v", err)
|
||||||
|
}
|
||||||
|
if len(calls) != 1 || calls[0].Status != "failed" || calls[0].RiskLevel != "write" {
|
||||||
|
t.Fatalf("unexpected Workflow policy audit: %#v", calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopKnowledgeFallbackCanRequestHandoff(t *testing.T) {
|
||||||
|
agent := models.AIAgent{
|
||||||
|
KnowledgeIDs: "1", FallbackMode: enums.AIAgentFallbackModeHandoff,
|
||||||
|
FallbackMessage: "我暂时无法核实,马上为你转人工。",
|
||||||
|
}
|
||||||
|
policy := evaluateAgentLoopResponsePolicy(agent, "", nil)
|
||||||
|
prompt := buildAgentLoopSystemPrompt(agent, true, "", nil)
|
||||||
|
if !policy.RequestHandoff || !strings.Contains(prompt, agent.FallbackMessage) {
|
||||||
|
t.Fatalf("knowledge fallback was not applied: policy=%#v prompt=%q", policy, prompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopConfirmationNormalizesHTMLAndKeepsUnknownPending(t *testing.T) {
|
||||||
|
data := normalizeAgentLoopResumeData(enums.IMMessageTypeHTML, map[string]string{
|
||||||
|
"message": "<p>确认。</p>",
|
||||||
|
})
|
||||||
|
if got := parseAgentLoopConfirmation(firstAgentLoopResumeText(data)); got != agentLoopConfirmationConfirmed {
|
||||||
|
t.Fatalf("expected HTML confirmation, got %v from %#v", got, data)
|
||||||
|
}
|
||||||
|
if got := parseAgentLoopConfirmation("取消!"); got != agentLoopConfirmationCancelled {
|
||||||
|
t.Fatalf("expected cancellation, got %v", got)
|
||||||
|
}
|
||||||
|
if got := parseAgentLoopConfirmation("稍后再说"); got != agentLoopConfirmationUnknown {
|
||||||
|
t.Fatalf("ambiguous input must stay pending, got %v", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfiguredMCPToolAppliesTrustedSystemPolicy(t *testing.T) {
|
||||||
|
configured, _ := json.Marshal([]request.AIAgentMCPToolRequest{{
|
||||||
|
ToolCode: "system/server_time",
|
||||||
|
ServerCode: "system",
|
||||||
|
ToolName: "server_time",
|
||||||
|
Title: "server_time",
|
||||||
|
RiskLevel: "write",
|
||||||
|
RequireConfirmation: true,
|
||||||
|
}})
|
||||||
|
tool, err := configuredMCPTool(string(configured), "system/server_time")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("resolve configured system tool: %v", err)
|
||||||
|
}
|
||||||
|
if tool.Title != "获取当前时间" || tool.RiskLevel != "read" || tool.RequireConfirmation {
|
||||||
|
t.Fatalf("trusted policy was not applied at runtime: %#v", tool)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopPromptAvoidsRepeatingWelcomeMessage(t *testing.T) {
|
||||||
|
prompt := buildAgentLoopSystemPrompt(models.AIAgent{}, false, "", nil)
|
||||||
|
if !strings.Contains(prompt, "without repeating the welcome wording") {
|
||||||
|
t.Fatalf("conversation continuity instruction missing: %q", prompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopPromptEnforcesCustomerLanguage(t *testing.T) {
|
||||||
|
prompt := buildAgentLoopSystemPrompt(models.AIAgent{}, false, "", nil)
|
||||||
|
if !strings.Contains(prompt, "CRITICAL LANGUAGE RULE") || !strings.Contains(prompt, "same language as the customer's current message") {
|
||||||
|
t.Fatalf("language rule missing from system prompt: %q", prompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCompleteConfirmedMCPReplyGeneratesCustomerFacingAnswerWithoutTools(t *testing.T) {
|
||||||
|
engine := NewAgentLoopEngine()
|
||||||
|
var systemPrompt string
|
||||||
|
var userPrompt string
|
||||||
|
engine.complete = func(_ context.Context, _ models.AIConfig, system, user string) (*ai.ChatCompletionResult, error) {
|
||||||
|
systemPrompt = system
|
||||||
|
userPrompt = user
|
||||||
|
return &ai.ChatCompletionResult{
|
||||||
|
Content: "当前服务端时间是 2026-07-28 11:51:52。",
|
||||||
|
ModelName: "test-model",
|
||||||
|
PromptTokens: 20,
|
||||||
|
CompletionTokens: 10,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
result, err := engine.completeConfirmedMCPReply(
|
||||||
|
context.Background(),
|
||||||
|
models.AIAgent{},
|
||||||
|
models.AIConfig{ModelName: "test-model"},
|
||||||
|
"获取当前时间",
|
||||||
|
"现在几点钟?",
|
||||||
|
`{"timestamp":"2026-07-28 11:51:52","timezone":"Local"}`,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("complete confirmed MCP reply: %v", err)
|
||||||
|
}
|
||||||
|
if result.Content != "当前服务端时间是 2026-07-28 11:51:52。" {
|
||||||
|
t.Fatalf("unexpected customer reply: %#v", result)
|
||||||
|
}
|
||||||
|
for _, expected := range []string{
|
||||||
|
"Do not request or invoke another tool",
|
||||||
|
"现在几点钟?",
|
||||||
|
"获取当前时间",
|
||||||
|
`"timestamp":"2026-07-28 11:51:52"`,
|
||||||
|
} {
|
||||||
|
if !strings.Contains(systemPrompt+"\n"+userPrompt, expected) {
|
||||||
|
t.Fatalf("post-tool completion context missing %q: system=%q user=%q", expected, systemPrompt, userPrompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConfirmedMCPReplyFallbackDoesNotExposeRawResult(t *testing.T) {
|
||||||
|
got := buildAgentLoopConfirmedMCPFallback("获取当前时间")
|
||||||
|
if got != "“获取当前时间”已成功执行。" || strings.Contains(got, "{") {
|
||||||
|
t.Fatalf("unexpected confirmed MCP fallback: %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentTurnPublishesAllConfiguredCapabilityKinds(t *testing.T) {
|
||||||
|
db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open sqlite: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.AutoMigrate(&models.SkillDefinition{}); err != nil {
|
||||||
|
t.Fatalf("migrate Skill: %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
skill := models.SkillDefinition{Name: "订单查询", Description: "查询订单状态", Status: enums.StatusOk}
|
||||||
|
if err := db.Create(&skill).Error; err != nil {
|
||||||
|
t.Fatalf("create Skill: %v", err)
|
||||||
|
}
|
||||||
|
mcp, _ := json.Marshal([]request.AIAgentMCPToolRequest{{
|
||||||
|
ToolCode: "crm/get_customer", ServerCode: "crm", ToolName: "get_customer",
|
||||||
|
RiskLevel: "read",
|
||||||
|
}})
|
||||||
|
agent := models.AIAgent{SkillIDs: jsonInt64List(skill.ID), AllowedMCPTools: string(mcp)}
|
||||||
|
snapshot := &svc.AgentRevisionSnapshot{
|
||||||
|
Agent: agent,
|
||||||
|
WorkflowBindings: []svc.AgentRevisionWorkflowBinding{{
|
||||||
|
WorkflowVersionID: 23, ToolName: "创建工单", TriggerInstruction: "用户要求创建工单",
|
||||||
|
}},
|
||||||
|
}
|
||||||
|
engine := NewAgentLoopEngine()
|
||||||
|
engine.retrieve = nil
|
||||||
|
engine.history = nil
|
||||||
|
turn := engine.prepareTurn(context.Background(), RunInput{AIAgent: agent}, snapshot)
|
||||||
|
|
||||||
|
for _, code := range []string{"skill/" + jsonInt64List(skill.ID), "workflow/23", "crm/get_customer"} {
|
||||||
|
if !strings.Contains(turn.SystemPrompt, code) {
|
||||||
|
t.Fatalf("capability %q missing from prompt:\n%s", code, turn.SystemPrompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func jsonInt64List(id int64) string {
|
||||||
|
data, _ := json.Marshal([]int64{id})
|
||||||
|
return strings.Trim(string(data), "[]")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentTurnCatalogFormatAndInstruction(t *testing.T) {
|
||||||
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open in-memory db: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.AutoMigrate(&models.SkillDefinition{}); err != nil {
|
||||||
|
t.Fatalf("migrate Skill: %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
skill := models.SkillDefinition{Name: "订单查询", Description: "查询订单状态", Status: enums.StatusOk}
|
||||||
|
if err := db.Create(&skill).Error; err != nil {
|
||||||
|
t.Fatalf("create Skill: %v", err)
|
||||||
|
}
|
||||||
|
mcp, _ := json.Marshal([]request.AIAgentMCPToolRequest{{
|
||||||
|
ToolCode: "crm/get_customer", ServerCode: "crm", ToolName: "get_customer",
|
||||||
|
RiskLevel: "read",
|
||||||
|
}})
|
||||||
|
agent := models.AIAgent{SkillIDs: jsonInt64List(skill.ID), AllowedMCPTools: string(mcp)}
|
||||||
|
snapshot := &svc.AgentRevisionSnapshot{
|
||||||
|
Agent: agent,
|
||||||
|
WorkflowBindings: []svc.AgentRevisionWorkflowBinding{{
|
||||||
|
WorkflowVersionID: 23, ToolName: "创建工单", TriggerInstruction: "用户要求创建工单",
|
||||||
|
}},
|
||||||
|
}
|
||||||
|
engine := NewAgentLoopEngine()
|
||||||
|
engine.retrieve = nil
|
||||||
|
engine.history = nil
|
||||||
|
turn := engine.prepareTurn(context.Background(), RunInput{AIAgent: agent}, snapshot)
|
||||||
|
|
||||||
|
expectedDirectFunctions := []string{
|
||||||
|
"direct function: builtin_conversation_context",
|
||||||
|
"direct function: builtin_knowledge_retrieve",
|
||||||
|
"direct function: graph_triage_service_request",
|
||||||
|
"direct function: skill_" + jsonInt64List(skill.ID),
|
||||||
|
"direct function: workflow_23",
|
||||||
|
"direct function: crm_get_customer",
|
||||||
|
}
|
||||||
|
for _, expected := range expectedDirectFunctions {
|
||||||
|
if !strings.Contains(turn.SystemPrompt, expected) {
|
||||||
|
t.Fatalf("expected direct function %q in system prompt:\n%s", expected, turn.SystemPrompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
prohibition := "Never pass a capability code containing slashes as a tool function name; use the exact direct function name provided in the catalog or call tool_search with toolCode."
|
||||||
|
if !strings.Contains(turn.SystemPrompt, prohibition) {
|
||||||
|
t.Fatalf("expected prohibition instruction in system prompt:\n%s", turn.SystemPrompt)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopEngineEmitsActivitySafelyWithConversationID(t *testing.T) {
|
||||||
|
db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open sqlite: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.AutoMigrate(
|
||||||
|
&models.AIConfig{}, &models.AIAgent{}, &models.AgentRevision{},
|
||||||
|
&models.Conversation{}, &models.Message{},
|
||||||
|
&models.AgentRun{}, &models.AgentStep{}, &models.AgentToolCall{},
|
||||||
|
); err != nil {
|
||||||
|
t.Fatalf("auto migrate: %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
|
||||||
|
config := &models.AIConfig{Status: enums.StatusOk, ModelName: "test-model"}
|
||||||
|
if err := db.Create(config).Error; err != nil {
|
||||||
|
t.Fatalf("create config: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
definitionJSON, _ := json.Marshal(map[string]any{
|
||||||
|
"agent": map[string]any{"name": "test-agent"},
|
||||||
|
})
|
||||||
|
agent := &models.AIAgent{Name: "test-agent", Status: enums.StatusOk, AIConfigID: config.ID}
|
||||||
|
if err := db.Create(agent).Error; err != nil {
|
||||||
|
t.Fatalf("create agent: %v", err)
|
||||||
|
}
|
||||||
|
revision := &models.AgentRevision{
|
||||||
|
AgentID: agent.ID, Status: enums.StatusOk, Revision: 1,
|
||||||
|
Definition: string(definitionJSON),
|
||||||
|
}
|
||||||
|
if err := db.Create(revision).Error; err != nil {
|
||||||
|
t.Fatalf("create revision: %v", err)
|
||||||
|
}
|
||||||
|
agent.PublishedRevisionID = revision.ID
|
||||||
|
if err := db.Save(agent).Error; err != nil {
|
||||||
|
t.Fatalf("save agent: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
conv := models.Conversation{ID: 100}
|
||||||
|
userMsg := models.Message{ID: 200, Content: "test message", SenderType: enums.IMSenderTypeCustomer, MessageType: enums.IMMessageTypeText}
|
||||||
|
|
||||||
|
// 1. Test Run with conversation ID > 0 (success)
|
||||||
|
engine := newAgentLoopEngineWithLoop(func(ctx context.Context, config models.AIConfig, sysPrompt, userPrompt string, tools []ai.ToolDefinition, maxSteps int, executor ai.ToolCallExecutor) (*ai.ToolLoopResult, error) {
|
||||||
|
return &ai.ToolLoopResult{
|
||||||
|
ChatCompletionResult: ai.ChatCompletionResult{
|
||||||
|
Content: "Hello from loop",
|
||||||
|
ModelName: "test-model",
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
})
|
||||||
|
engine.retrieve = nil
|
||||||
|
engine.history = nil
|
||||||
|
|
||||||
|
res, err := engine.Run(context.Background(), RunInput{
|
||||||
|
Conversation: conv,
|
||||||
|
UserMessage: userMsg,
|
||||||
|
AIAgent: *agent,
|
||||||
|
AIConfig: *config,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected Run to succeed, got %v", err)
|
||||||
|
}
|
||||||
|
if res == nil || res.ReplyText != "Hello from loop" {
|
||||||
|
t.Fatalf("unexpected result: %#v", res)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Test Run with conversation ID > 0 (loop error)
|
||||||
|
errEngine := newAgentLoopEngineWithLoop(func(ctx context.Context, config models.AIConfig, sysPrompt, userPrompt string, tools []ai.ToolDefinition, maxSteps int, executor ai.ToolCallExecutor) (*ai.ToolLoopResult, error) {
|
||||||
|
return nil, errors.New("loop execution failed")
|
||||||
|
})
|
||||||
|
errEngine.retrieve = nil
|
||||||
|
errEngine.history = nil
|
||||||
|
|
||||||
|
_, loopErr := errEngine.Run(context.Background(), RunInput{
|
||||||
|
Conversation: conv,
|
||||||
|
UserMessage: userMsg,
|
||||||
|
AIAgent: *agent,
|
||||||
|
AIConfig: *config,
|
||||||
|
})
|
||||||
|
if loopErr == nil || !strings.Contains(loopErr.Error(), "loop execution failed") {
|
||||||
|
t.Fatalf("expected loop error, got: %v", loopErr)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Test toolSearchExecutor with conversation ID > 0
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"builtin/conversation_context"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(""),
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
executor := engine.toolSearchExecutor(RunInput{Conversation: conv}, turn, &state, &calls)
|
||||||
|
|
||||||
|
toolResult, err := executor(context.Background(), ai.ToolCall{
|
||||||
|
Name: "builtin_conversation_context",
|
||||||
|
Arguments: `{}`,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("executor tool call failed: %v", err)
|
||||||
|
}
|
||||||
|
if toolResult == "" {
|
||||||
|
t.Fatal("expected non-empty tool result")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. Test offline / zero conversation ID
|
||||||
|
offlineRes, err := engine.Run(context.Background(), RunInput{
|
||||||
|
Conversation: models.Conversation{ID: 0},
|
||||||
|
UserMessage: userMsg,
|
||||||
|
AIAgent: *agent,
|
||||||
|
AIConfig: *config,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected offline Run to succeed, got %v", err)
|
||||||
|
}
|
||||||
|
if offlineRes == nil || offlineRes.ReplyText != "Hello from loop" {
|
||||||
|
t.Fatalf("unexpected offline result: %#v", offlineRes)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopEngineRetrieveKnowledgeEmitsActivity(t *testing.T) {
|
||||||
|
engine := NewAgentLoopEngine()
|
||||||
|
retrievedCalled := false
|
||||||
|
engine.retrieve = func(ctx context.Context, agent models.AIAgent, query string) (string, int, error) {
|
||||||
|
retrievedCalled = true
|
||||||
|
return "evidence", 1, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Case 1: KnowledgeIDs present and conversation in context
|
||||||
|
conv := models.Conversation{ID: 100}
|
||||||
|
ctx := withAgentLoopConversation(context.Background(), conv)
|
||||||
|
text, count, err := engine.retrieveKnowledge(ctx, models.AIAgent{KnowledgeIDs: "1,2"}, "query")
|
||||||
|
if err != nil || !retrievedCalled || text != "evidence" || count != 1 {
|
||||||
|
t.Fatalf("retrieveKnowledge failed: text=%q count=%d err=%v", text, count, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Case 2: No KnowledgeIDs -> retrieve not called, returns empty
|
||||||
|
retrievedCalled = false
|
||||||
|
text, count, err = engine.retrieveKnowledge(ctx, models.AIAgent{KnowledgeIDs: ""}, "query")
|
||||||
|
if err != nil || retrievedCalled || text != "" || count != 0 {
|
||||||
|
t.Fatalf("expected no retrieve call when KnowledgeIDs is empty")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Case 3: Conversation ID is 0 -> retrieve called safely without panic
|
||||||
|
retrievedCalled = false
|
||||||
|
zeroCtx := withAgentLoopConversation(context.Background(), models.Conversation{ID: 0})
|
||||||
|
text, count, err = engine.retrieveKnowledge(zeroCtx, models.AIAgent{KnowledgeIDs: "1"}, "query")
|
||||||
|
if err != nil || !retrievedCalled || text != "evidence" {
|
||||||
|
t.Fatalf("retrieveKnowledge with zero conversation failed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopToolSearchExecutorEmitsWorkflowActivity(t *testing.T) {
|
||||||
|
turn := agentLoopTurn{
|
||||||
|
AllowedTools: []string{"workflow/42"},
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(""),
|
||||||
|
Workflows: map[int64]svc.AgentRevisionWorkflowBinding{
|
||||||
|
42: {WorkflowVersionID: 42, ToolName: "Cancel Subscription"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
state := agentLoopExecutionState{}
|
||||||
|
var calls []svc.AgentLoopToolCallInput
|
||||||
|
conv := models.Conversation{ID: 100}
|
||||||
|
executor := NewAgentLoopEngine().toolSearchExecutor(RunInput{Conversation: conv}, turn, &state, &calls)
|
||||||
|
|
||||||
|
// Since workflow 42 is not in db, it will fail to resolve workflow version, but activity is published beforehand
|
||||||
|
_, _ = executor(context.Background(), ai.ToolCall{
|
||||||
|
Name: "workflow_42",
|
||||||
|
Arguments: `{}`,
|
||||||
|
})
|
||||||
|
// Check that records captured the attempt
|
||||||
|
if len(calls) != 1 || calls[0].ToolCode != "workflow/42" {
|
||||||
|
t.Fatalf("unexpected calls: %#v", calls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopEngineBuildUserPromptContextWindowLimits(t *testing.T) {
|
||||||
|
engine := NewAgentLoopEngine()
|
||||||
|
// Create 50 dummy messages in history
|
||||||
|
engine.history = func(conversationID int64, limit int) []models.Message {
|
||||||
|
messages := make([]models.Message, 50)
|
||||||
|
for i := 0; i < 50; i++ {
|
||||||
|
messages[i] = models.Message{
|
||||||
|
ID: int64(i + 1),
|
||||||
|
SenderType: enums.IMSenderTypeCustomer,
|
||||||
|
MessageType: enums.IMMessageTypeText,
|
||||||
|
Content: fmt.Sprintf("message %d", i+1),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if limit < len(messages) {
|
||||||
|
return messages[:limit]
|
||||||
|
}
|
||||||
|
return messages
|
||||||
|
}
|
||||||
|
|
||||||
|
testCases := []struct {
|
||||||
|
name string
|
||||||
|
contextWindow int
|
||||||
|
expectedCount int
|
||||||
|
}{
|
||||||
|
{"default when zero", 0, 10},
|
||||||
|
{"default when negative", -3, 10},
|
||||||
|
{"clamped to min 2", 1, 2},
|
||||||
|
{"within range", 15, 15},
|
||||||
|
{"clamped to max 30", 45, 30},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range testCases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
req := RunInput{
|
||||||
|
Conversation: models.Conversation{ID: 1},
|
||||||
|
UserMessage: models.Message{ID: 999, Content: "current query", SenderType: enums.IMSenderTypeCustomer, MessageType: enums.IMMessageTypeText},
|
||||||
|
AIAgent: models.AIAgent{ContextWindow: tc.contextWindow},
|
||||||
|
}
|
||||||
|
prompt, count := engine.buildUserPrompt(req)
|
||||||
|
if count != tc.expectedCount {
|
||||||
|
t.Fatalf("expected history count %d, got %d (prompt=%q)", tc.expectedCount, count, prompt)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopEnginePrepareTurnMaxStepsBounds(t *testing.T) {
|
||||||
|
engine := NewAgentLoopEngine()
|
||||||
|
engine.retrieve = nil
|
||||||
|
engine.history = nil
|
||||||
|
snapshot := &svc.AgentRevisionSnapshot{}
|
||||||
|
|
||||||
|
testCases := []struct {
|
||||||
|
name string
|
||||||
|
inputMaxSteps int
|
||||||
|
expectedSteps int
|
||||||
|
}{
|
||||||
|
{"default when zero", 0, 4},
|
||||||
|
{"default when negative", -2, 4},
|
||||||
|
{"clamped to max 8", 12, 8},
|
||||||
|
{"within range 6", 6, 6},
|
||||||
|
{"within range 2", 2, 2},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range testCases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
turn := engine.prepareTurn(context.Background(), RunInput{
|
||||||
|
AIAgent: models.AIAgent{MaxSteps: tc.inputMaxSteps},
|
||||||
|
}, snapshot)
|
||||||
|
if turn.MaxSteps != tc.expectedSteps {
|
||||||
|
t.Fatalf("expected turn.MaxSteps %d, got %d", tc.expectedSteps, turn.MaxSteps)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopEngineRunPassesBoundedMaxSteps(t *testing.T) {
|
||||||
|
db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open sqlite: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.AutoMigrate(
|
||||||
|
&models.AIConfig{}, &models.AIAgent{}, &models.AgentRevision{},
|
||||||
|
&models.Conversation{}, &models.Message{},
|
||||||
|
&models.AgentRun{}, &models.AgentStep{}, &models.AgentToolCall{},
|
||||||
|
); err != nil {
|
||||||
|
t.Fatalf("auto migrate: %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
|
||||||
|
config := &models.AIConfig{Status: enums.StatusOk, ModelName: "test-model"}
|
||||||
|
if err := db.Create(config).Error; err != nil {
|
||||||
|
t.Fatalf("create config: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
definitionJSON, _ := json.Marshal(map[string]any{
|
||||||
|
"agent": map[string]any{"name": "test-agent", "maxSteps": 15},
|
||||||
|
})
|
||||||
|
agent := &models.AIAgent{Name: "test-agent", Status: enums.StatusOk, AIConfigID: config.ID, MaxSteps: 15}
|
||||||
|
if err := db.Create(agent).Error; err != nil {
|
||||||
|
t.Fatalf("create agent: %v", err)
|
||||||
|
}
|
||||||
|
revision := &models.AgentRevision{
|
||||||
|
AgentID: agent.ID, Status: enums.StatusOk, Revision: 1,
|
||||||
|
Definition: string(definitionJSON),
|
||||||
|
}
|
||||||
|
if err := db.Create(revision).Error; err != nil {
|
||||||
|
t.Fatalf("create revision: %v", err)
|
||||||
|
}
|
||||||
|
agent.PublishedRevisionID = revision.ID
|
||||||
|
if err := db.Save(agent).Error; err != nil {
|
||||||
|
t.Fatalf("save agent: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var capturedMaxSteps int
|
||||||
|
engine := newAgentLoopEngineWithLoop(func(ctx context.Context, config models.AIConfig, sysPrompt, userPrompt string, tools []ai.ToolDefinition, maxSteps int, executor ai.ToolCallExecutor) (*ai.ToolLoopResult, error) {
|
||||||
|
capturedMaxSteps = maxSteps
|
||||||
|
return &ai.ToolLoopResult{
|
||||||
|
ChatCompletionResult: ai.ChatCompletionResult{
|
||||||
|
Content: "Hello",
|
||||||
|
ModelName: "test-model",
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
})
|
||||||
|
engine.retrieve = nil
|
||||||
|
engine.history = nil
|
||||||
|
|
||||||
|
_, err = engine.Run(context.Background(), RunInput{
|
||||||
|
Conversation: models.Conversation{ID: 100},
|
||||||
|
UserMessage: models.Message{ID: 200, Content: "test", SenderType: enums.IMSenderTypeCustomer, MessageType: enums.IMMessageTypeText},
|
||||||
|
AIAgent: *agent,
|
||||||
|
AIConfig: *config,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Run failed: %v", err)
|
||||||
|
}
|
||||||
|
if capturedMaxSteps != 8 {
|
||||||
|
t.Fatalf("expected capturedMaxSteps to be capped at 8, got %d", capturedMaxSteps)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentLoopEngineMultilingualHandling(t *testing.T) {
|
||||||
|
engine := NewAgentLoopEngine()
|
||||||
|
|
||||||
|
// 1. Language directive appended to Russian prompt
|
||||||
|
ruPrompt, _ := engine.buildUserPrompt(RunInput{
|
||||||
|
UserMessage: models.Message{Content: "Здравствуйте, как вернуть товар?"},
|
||||||
|
})
|
||||||
|
if !strings.Contains(ruPrompt, "Language directive: The customer is writing in Russian") {
|
||||||
|
t.Fatalf("expected Russian language directive in prompt, got %q", ruPrompt)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Language directive appended to English/Spanish/global prompt
|
||||||
|
enPrompt, _ := engine.buildUserPrompt(RunInput{
|
||||||
|
UserMessage: models.Message{Content: "Hello, can I get a refund?"},
|
||||||
|
})
|
||||||
|
if !strings.Contains(enPrompt, "Language directive: You MUST inspect the language used by the customer") {
|
||||||
|
t.Fatalf("expected universal language directive in English prompt, got %q", enPrompt)
|
||||||
|
}
|
||||||
|
|
||||||
|
esPrompt, _ := engine.buildUserPrompt(RunInput{
|
||||||
|
UserMessage: models.Message{Content: "Hola, ¿cómo puedo rastrear mi paquete?"},
|
||||||
|
})
|
||||||
|
if !strings.Contains(esPrompt, "Language directive: You MUST inspect the language used by the customer") {
|
||||||
|
t.Fatalf("expected universal language directive in Spanish prompt, got %q", esPrompt)
|
||||||
|
}
|
||||||
|
|
||||||
|
zhPrompt, _ := engine.buildUserPrompt(RunInput{
|
||||||
|
UserMessage: models.Message{Content: "你好,怎么退款?"},
|
||||||
|
})
|
||||||
|
if !strings.Contains(zhPrompt, "Language directive: The customer is writing in Chinese") {
|
||||||
|
t.Fatalf("expected Chinese language directive in Chinese prompt, got %q", zhPrompt)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Russian confirmation and cancellation parsing
|
||||||
|
if parseAgentLoopConfirmation("да") != agentLoopConfirmationConfirmed {
|
||||||
|
t.Fatalf("expected 'да' to be confirmed")
|
||||||
|
}
|
||||||
|
if parseAgentLoopConfirmation("подтверждаю") != agentLoopConfirmationConfirmed {
|
||||||
|
t.Fatalf("expected 'подтверждаю' to be confirmed")
|
||||||
|
}
|
||||||
|
if parseAgentLoopConfirmation("нет") != agentLoopConfirmationCancelled {
|
||||||
|
t.Fatalf("expected 'нет' to be cancelled")
|
||||||
|
}
|
||||||
|
if parseAgentLoopConfirmation("отмена") != agentLoopConfirmationCancelled {
|
||||||
|
t.Fatalf("expected 'отмена' to be cancelled")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. Russian MCP confirmation prompt and fallback
|
||||||
|
ruConfirmationPrompt := buildAgentLoopMCPConfirmationPrompt("Создание заказа", "Создайте заказ, пожалуйста")
|
||||||
|
if !strings.Contains(ruConfirmationPrompt, "Будет выполнено действие «Создание заказа». Подтверждаете продолжение?") {
|
||||||
|
t.Fatalf("unexpected Russian MCP confirmation prompt: %q", ruConfirmationPrompt)
|
||||||
|
}
|
||||||
|
|
||||||
|
ruFallback := buildAgentLoopConfirmedMCPFallback("Создание заказа", "Создайте заказ, пожалуйста")
|
||||||
|
if !strings.Contains(ruFallback, "Действие «Создание заказа» успешно выполнено.") {
|
||||||
|
t.Fatalf("unexpected Russian MCP fallback: %q", ruFallback)
|
||||||
|
}
|
||||||
|
}
|
||||||
102
internal/ai/application/runtime/agent_turn.go
Normal file
102
internal/ai/application/runtime/agent_turn.go
Normal file
@@ -0,0 +1,102 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
"agent-desk/internal/pkg/utils"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
type agentLoopTurn struct {
|
||||||
|
RetrieverCount int
|
||||||
|
RetrieveErr error
|
||||||
|
ResponsePolicy agentLoopResponsePolicy
|
||||||
|
SystemPrompt string
|
||||||
|
UserPrompt string
|
||||||
|
HistoryCount int
|
||||||
|
AllowedTools []string
|
||||||
|
ToolPolicy agentLoopToolPolicy
|
||||||
|
Skills map[int64]models.SkillDefinition
|
||||||
|
Workflows map[int64]svc.AgentRevisionWorkflowBinding
|
||||||
|
MaxSteps int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *AgentLoopEngine) prepareTurn(ctx context.Context, req RunInput, snapshot *svc.AgentRevisionSnapshot) agentLoopTurn {
|
||||||
|
maxSteps := req.AIAgent.MaxSteps
|
||||||
|
if maxSteps <= 0 {
|
||||||
|
maxSteps = 4
|
||||||
|
}
|
||||||
|
if maxSteps > 8 {
|
||||||
|
maxSteps = 8
|
||||||
|
}
|
||||||
|
knowledgeContext, retrieverCount, retrieveErr := e.retrieveKnowledge(ctx, req.AIAgent, req.UserMessage.Content)
|
||||||
|
responsePolicy := evaluateAgentLoopResponsePolicy(req.AIAgent, knowledgeContext, retrieveErr)
|
||||||
|
systemPrompt := buildAgentLoopSystemPrompt(req.AIAgent, len(utils.SplitInt64s(req.AIAgent.KnowledgeIDs)) > 0, knowledgeContext, retrieveErr)
|
||||||
|
userPrompt, historyCount := e.buildUserPrompt(req)
|
||||||
|
if knowledgeContext != "" {
|
||||||
|
userPrompt += "\n\nKnowledge evidence:\n" + knowledgeContext
|
||||||
|
}
|
||||||
|
skills := svc.SkillDefinitionService.GetByIDs(utils.SplitInt64s(req.AIAgent.SkillIDs))
|
||||||
|
workflows := make(map[int64]svc.AgentRevisionWorkflowBinding, len(snapshot.WorkflowBindings))
|
||||||
|
allowedTools := agentLoopSafeBuiltinCodes()
|
||||||
|
// TODO 这么实现我觉得不太好,最好是能够有个统一的能力目录
|
||||||
|
catalog := []string{
|
||||||
|
"- " + toolx.BuiltinConversationContext.Code + " | Builtin | Read current conversation and customer context | direct function: " + toolx.SanitizeToolFunctionName(toolx.BuiltinConversationContext.Code),
|
||||||
|
"- " + toolx.BuiltinKnowledgeRetrieve.Code + " | Builtin | Retrieve linked knowledge base on demand | direct function: " + toolx.SanitizeToolFunctionName(toolx.BuiltinKnowledgeRetrieve.Code),
|
||||||
|
"- " + toolx.GraphTriageServiceRequest.Code + " | Builtin | Analyze service request and generate handling advice | direct function: " + toolx.SanitizeToolFunctionName(toolx.GraphTriageServiceRequest.Code),
|
||||||
|
"- " + toolx.GraphAnalyzeConversation.Code + " | Builtin | Analyze conversation intent and risk signals | direct function: " + toolx.SanitizeToolFunctionName(toolx.GraphAnalyzeConversation.Code),
|
||||||
|
"- " + toolx.GraphPrepareTicketDraft.Code + " | Builtin | Generate ticket draft without writing | direct function: " + toolx.SanitizeToolFunctionName(toolx.GraphPrepareTicketDraft.Code),
|
||||||
|
}
|
||||||
|
for id, skill := range skills {
|
||||||
|
if skill.Status != enums.StatusOk {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
code := agentLoopSkillCode(id)
|
||||||
|
allowedTools = append(allowedTools, code)
|
||||||
|
catalog = append(catalog, fmt.Sprintf("- %s | Skill | %s | %s | direct function: %s", code, strings.TrimSpace(skill.Name), strings.TrimSpace(skill.Description), toolx.SanitizeToolFunctionName(code)))
|
||||||
|
}
|
||||||
|
for _, binding := range snapshot.WorkflowBindings {
|
||||||
|
if binding.WorkflowVersionID <= 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
workflows[binding.WorkflowVersionID] = binding
|
||||||
|
code := agentLoopWorkflowCode(binding.WorkflowVersionID)
|
||||||
|
allowedTools = append(allowedTools, code)
|
||||||
|
catalog = append(catalog, fmt.Sprintf("- %s | Workflow | %s | %s | direct function: %s", code, strings.TrimSpace(binding.ToolName), strings.TrimSpace(binding.TriggerInstruction), toolx.SanitizeToolFunctionName(code)))
|
||||||
|
}
|
||||||
|
mcpTools, _ := toolx.ParseAgentMCPToolsJSON(req.AIAgent.AllowedMCPTools)
|
||||||
|
for _, tool := range mcpTools {
|
||||||
|
if strings.TrimSpace(tool.ToolCode) == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
allowedTools = append(allowedTools, tool.ToolCode)
|
||||||
|
catalog = append(catalog, fmt.Sprintf("- %s | MCP | %s | %s | direct function: %s", tool.ToolCode, tool.Title, tool.Description, toolx.SanitizeToolFunctionName(tool.ToolCode)))
|
||||||
|
}
|
||||||
|
systemPrompt += "\n\nAvailable capabilities:\n" + strings.Join(catalog, "\n")
|
||||||
|
systemPrompt += "\n\nUse tool_search with an exact capability code only when needed. You decide whether to answer directly, activate a Skill, execute a Workflow, retrieve knowledge, or call MCP. A Skill activation returns instructions for this same run. Never invent a capability code. Never pass a capability code containing slashes as a tool function name; use the exact direct function name provided in the catalog or call tool_search with toolCode."
|
||||||
|
systemPrompt += "\n\nHUMAN OPERATOR TRANSFER & SEMANTIC INTENT RECOGNITION:\n" +
|
||||||
|
"- Customers may request to speak to a human operator, specialist, agent, manager, representative, or live person in many different ways across ANY language (e.g., 'переведи на оператора', 'соедини со специалистом', 'позови сотрудника', 'хочу живого человека', 'мне нужен оператор', 'хватит бота', 'с кем я могу поговорить кроме вас', 'can I speak with someone?', 'connect me to an agent', 'I want a real person', 'representative please', 'let me talk to support', 'quiero hablar con un humano', 'ich möchte mit einem Mitarbeiter sprechen', etc.).\n" +
|
||||||
|
"- Do NOT require a specific hardcoded phrase. You MUST understand this intent semantically from the user's message regardless of their wording, slang, or language.\n" +
|
||||||
|
"- When the customer indicates or asks that they want to be connected to a human operator or specialist, you MUST IMMEDIATELY call the conversation_decision tool with action='handoff', handoffInitiator='customer', handoffConfirmed=true, and a short reason. Do NOT ask for confirmation again, do NOT refuse, and do NOT claim in text that the handoff succeeded without calling conversation_decision.\n" +
|
||||||
|
"- Use action='ask_handoff_confirmation' only when YOU (the AI agent), not the customer, proactively suggest transferring to a human, with handoffInitiator='agent' and handoffConfirmed=false.\n" +
|
||||||
|
"- For any regular response to the customer, call conversation_decision with action='reply', handoffInitiator='none', handoffConfirmed=false, and your response in reply."
|
||||||
|
systemPrompt += "\n\nCRITICAL MULTILINGUAL RESPONSE DIRECTIVE: You MUST ALWAYS reply in the EXACT SAME LANGUAGE that the customer uses in their inquiry (e.g. Russian, English, Spanish, German, French, Arabic, Turkish, Kazakh, Italian, Portuguese, Japanese, etc.). Under NO circumstances should you default or drift into Chinese or any language other than the customer's language. Even if capabilities, tool names, or tool outputs contain Chinese or English, translate everything seamlessly into the customer's language."
|
||||||
|
return agentLoopTurn{
|
||||||
|
RetrieverCount: retrieverCount,
|
||||||
|
RetrieveErr: retrieveErr,
|
||||||
|
ResponsePolicy: responsePolicy,
|
||||||
|
SystemPrompt: systemPrompt,
|
||||||
|
UserPrompt: userPrompt,
|
||||||
|
HistoryCount: historyCount,
|
||||||
|
AllowedTools: allowedTools,
|
||||||
|
ToolPolicy: parseAgentLoopToolPolicy(req.AIAgent.ToolPolicy),
|
||||||
|
Skills: skills,
|
||||||
|
Workflows: workflows,
|
||||||
|
MaxSteps: maxSteps,
|
||||||
|
}
|
||||||
|
}
|
||||||
106
internal/ai/application/runtime/application_service.go
Normal file
106
internal/ai/application/runtime/application_service.go
Normal file
@@ -0,0 +1,106 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ApplicationRunInput identifies the persisted inputs for an Agent reply.
|
||||||
|
// Loading these records here keeps channels and debug adapters independent of
|
||||||
|
// individual engine requirements.
|
||||||
|
type ApplicationRunInput struct {
|
||||||
|
ConversationID int64
|
||||||
|
MessageID int64
|
||||||
|
AIAgentID int64
|
||||||
|
}
|
||||||
|
|
||||||
|
type ApplicationResumeInput struct {
|
||||||
|
ApplicationRunInput
|
||||||
|
CheckPointID string
|
||||||
|
ResumeData map[string]string
|
||||||
|
}
|
||||||
|
|
||||||
|
// AgentApplicationService is the single application boundary before engine
|
||||||
|
// dispatch. It owns persisted input loading and relationship validation; the
|
||||||
|
// Agent Loop remains responsible only for runtime execution.
|
||||||
|
type AgentApplicationService struct {
|
||||||
|
runtime *Service
|
||||||
|
}
|
||||||
|
|
||||||
|
var DefaultAgentApplicationService = NewAgentApplicationService()
|
||||||
|
|
||||||
|
func NewAgentApplicationService() *AgentApplicationService {
|
||||||
|
return &AgentApplicationService{runtime: NewService()}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AgentApplicationService) Run(ctx context.Context, input ApplicationRunInput) (*RunResult, error) {
|
||||||
|
req, err := s.loadRequest(input)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return s.RunPrepared(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RunPrepared is for isolated adapters such as the dashboard debug session.
|
||||||
|
// Callers are responsible for constructing an ephemeral or already-validated
|
||||||
|
// request; no persistence side effects are introduced by this boundary.
|
||||||
|
func (s *AgentApplicationService) RunPrepared(ctx context.Context, req RunInput) (*RunResult, error) {
|
||||||
|
return s.runtime.Run(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AgentApplicationService) Resume(ctx context.Context, input ApplicationResumeInput) (*RunResult, error) {
|
||||||
|
req, err := s.loadRequest(input.ApplicationRunInput)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
checkPointID := strings.TrimSpace(input.CheckPointID)
|
||||||
|
interrupt := svc.ConversationInterruptService.GetByCheckPointID(checkPointID)
|
||||||
|
if interrupt == nil || interrupt.ConversationID != req.Conversation.ID {
|
||||||
|
return nil, errorsx.InvalidParam("pending conversation interrupt does not exist")
|
||||||
|
}
|
||||||
|
if interrupt.AIAgentID > 0 && interrupt.AIAgentID != req.AIAgent.ID {
|
||||||
|
return nil, errorsx.InvalidParam("interrupt does not belong to agent")
|
||||||
|
}
|
||||||
|
return s.ResumePrepared(ctx, ResumeInput{
|
||||||
|
Conversation: req.Conversation,
|
||||||
|
UserMessage: req.UserMessage,
|
||||||
|
AIAgent: req.AIAgent,
|
||||||
|
AIConfig: req.AIConfig,
|
||||||
|
CheckPointID: checkPointID,
|
||||||
|
ResumeData: input.ResumeData,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AgentApplicationService) ResumePrepared(ctx context.Context, req ResumeInput) (*RunResult, error) {
|
||||||
|
return s.runtime.Resume(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *AgentApplicationService) loadRequest(input ApplicationRunInput) (RunInput, error) {
|
||||||
|
if input.ConversationID <= 0 || input.MessageID <= 0 || input.AIAgentID <= 0 {
|
||||||
|
return RunInput{}, errorsx.InvalidParam("conversation, message and agent are required")
|
||||||
|
}
|
||||||
|
conversation := svc.ConversationService.Get(input.ConversationID)
|
||||||
|
if conversation == nil {
|
||||||
|
return RunInput{}, errorsx.InvalidParam("conversation does not exist")
|
||||||
|
}
|
||||||
|
message := svc.MessageService.Get(input.MessageID)
|
||||||
|
if message == nil || message.ConversationID != conversation.ID {
|
||||||
|
return RunInput{}, errorsx.InvalidParam("message does not belong to conversation")
|
||||||
|
}
|
||||||
|
agent := svc.AIAgentService.Get(input.AIAgentID)
|
||||||
|
if agent == nil || agent.Status != enums.StatusOk {
|
||||||
|
return RunInput{}, errorsx.InvalidParam("ai agent is unavailable")
|
||||||
|
}
|
||||||
|
if conversation.AIAgentID > 0 && conversation.AIAgentID != agent.ID {
|
||||||
|
return RunInput{}, errorsx.InvalidParam("agent does not belong to conversation")
|
||||||
|
}
|
||||||
|
config := svc.AIConfigService.Get(agent.AIConfigID)
|
||||||
|
if config == nil || config.Status != enums.StatusOk {
|
||||||
|
return RunInput{}, errorsx.InvalidParam("ai config is unavailable")
|
||||||
|
}
|
||||||
|
return RunInput{Conversation: *conversation, UserMessage: *message, AIAgent: *agent, AIConfig: *config}, nil
|
||||||
|
}
|
||||||
54
internal/ai/application/runtime/application_service_test.go
Normal file
54
internal/ai/application/runtime/application_service_test.go
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
|
||||||
|
"github.com/glebarez/sqlite"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
"gorm.io/gorm"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAgentApplicationServiceLoadsConsistentPersistedRequest(t *testing.T) {
|
||||||
|
db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open sqlite: %v", err)
|
||||||
|
}
|
||||||
|
if err := db.AutoMigrate(&models.AIConfig{}, &models.AIAgent{}, &models.Conversation{}, &models.Message{}); err != nil {
|
||||||
|
t.Fatalf("auto migrate: %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
config := &models.AIConfig{Status: enums.StatusOk, ModelName: "test-model"}
|
||||||
|
if err := db.Create(config).Error; err != nil {
|
||||||
|
t.Fatalf("create config: %v", err)
|
||||||
|
}
|
||||||
|
agent := &models.AIAgent{Name: "agent", Status: enums.StatusOk, AIConfigID: config.ID}
|
||||||
|
if err := db.Create(agent).Error; err != nil {
|
||||||
|
t.Fatalf("create agent: %v", err)
|
||||||
|
}
|
||||||
|
conversation := &models.Conversation{AIAgentID: agent.ID}
|
||||||
|
if err := db.Create(conversation).Error; err != nil {
|
||||||
|
t.Fatalf("create conversation: %v", err)
|
||||||
|
}
|
||||||
|
message := &models.Message{ConversationID: conversation.ID, SenderType: enums.IMSenderTypeCustomer, MessageType: enums.IMMessageTypeText, Content: "hello"}
|
||||||
|
if err := db.Create(message).Error; err != nil {
|
||||||
|
t.Fatalf("create message: %v", err)
|
||||||
|
}
|
||||||
|
req, err := NewAgentApplicationService().loadRequest(ApplicationRunInput{ConversationID: conversation.ID, MessageID: message.ID, AIAgentID: agent.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("loadRequest: %v", err)
|
||||||
|
}
|
||||||
|
if req.Conversation.ID != conversation.ID || req.UserMessage.ID != message.ID || req.AIAgent.ID != agent.ID || req.AIConfig.ID != config.ID {
|
||||||
|
t.Fatalf("unexpected request: %#v", req)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgentApplicationServiceRejectsMismatchedMessage(t *testing.T) {
|
||||||
|
service := NewAgentApplicationService()
|
||||||
|
if _, err := service.loadRequest(ApplicationRunInput{ConversationID: 1, MessageID: 0, AIAgentID: 1}); err == nil {
|
||||||
|
t.Fatal("expected invalid identifiers error")
|
||||||
|
}
|
||||||
|
}
|
||||||
162
internal/ai/application/runtime/eino_agent_loop.go
Normal file
162
internal/ai/application/runtime/eino_agent_loop.go
Normal file
@@ -0,0 +1,162 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
ai "agent-desk/internal/ai"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
|
||||||
|
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
|
||||||
|
einomodel "github.com/cloudwego/eino/components/model"
|
||||||
|
einotool "github.com/cloudwego/eino/components/tool"
|
||||||
|
"github.com/cloudwego/eino/compose"
|
||||||
|
"github.com/cloudwego/eino/flow/agent/react"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
einojsonschema "github.com/eino-contrib/jsonschema"
|
||||||
|
)
|
||||||
|
|
||||||
|
// einoAgentLoop is the production model/tool loop. AgentDesk still owns tool
|
||||||
|
// authorization, business execution, interrupts, idempotency, and auditing.
|
||||||
|
func einoAgentLoop(
|
||||||
|
ctx context.Context,
|
||||||
|
config models.AIConfig,
|
||||||
|
systemPrompt string,
|
||||||
|
userPrompt string,
|
||||||
|
definitions []ai.ToolDefinition,
|
||||||
|
maxSteps int,
|
||||||
|
execute ai.ToolCallExecutor,
|
||||||
|
) (*ai.ToolLoopResult, error) {
|
||||||
|
model, err := newEinoChatModel(ctx, config)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if maxSteps <= 0 {
|
||||||
|
maxSteps = 4
|
||||||
|
}
|
||||||
|
if maxSteps > 8 {
|
||||||
|
maxSteps = 8
|
||||||
|
}
|
||||||
|
tools := make([]einotool.BaseTool, 0, len(definitions))
|
||||||
|
for _, definition := range definitions {
|
||||||
|
tool, buildErr := newEinoFunctionTool(definition, execute)
|
||||||
|
if buildErr != nil {
|
||||||
|
return nil, buildErr
|
||||||
|
}
|
||||||
|
tools = append(tools, tool)
|
||||||
|
}
|
||||||
|
agent, err := react.NewAgent(ctx, &react.AgentConfig{
|
||||||
|
ToolCallingModel: model,
|
||||||
|
ToolsConfig: compose.ToolsNodeConfig{Tools: tools},
|
||||||
|
MaxStep: maxSteps,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("create Eino agent loop: %w", err)
|
||||||
|
}
|
||||||
|
messages := make([]*schema.Message, 0, 2)
|
||||||
|
if value := strings.TrimSpace(systemPrompt); value != "" {
|
||||||
|
messages = append(messages, schema.SystemMessage(value))
|
||||||
|
}
|
||||||
|
messages = append(messages, schema.UserMessage(strings.TrimSpace(userPrompt)))
|
||||||
|
result, err := agent.Generate(ctx, messages)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if result == nil {
|
||||||
|
return nil, fmt.Errorf("Eino agent loop returned no result")
|
||||||
|
}
|
||||||
|
ret := &ai.ToolLoopResult{ChatCompletionResult: ai.ChatCompletionResult{
|
||||||
|
Content: strings.TrimSpace(result.Content),
|
||||||
|
ModelName: config.ModelName,
|
||||||
|
}}
|
||||||
|
if result.ResponseMeta != nil && result.ResponseMeta.Usage != nil {
|
||||||
|
ret.PromptTokens = result.ResponseMeta.Usage.PromptTokens
|
||||||
|
ret.CompletionTokens = result.ResponseMeta.Usage.CompletionTokens
|
||||||
|
}
|
||||||
|
return ret, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func newEinoChatModel(ctx context.Context, config models.AIConfig) (einomodel.ToolCallingChatModel, error) {
|
||||||
|
if strings.TrimSpace(config.APIKey) == "" || strings.TrimSpace(config.BaseURL) == "" || strings.TrimSpace(config.ModelName) == "" {
|
||||||
|
return nil, fmt.Errorf("ai config base URL, API key, and model name are required")
|
||||||
|
}
|
||||||
|
modelConfig := &einoopenai.ChatModelConfig{
|
||||||
|
APIKey: strings.TrimSpace(config.APIKey),
|
||||||
|
BaseURL: strings.TrimSpace(config.BaseURL),
|
||||||
|
Model: strings.TrimSpace(config.ModelName),
|
||||||
|
}
|
||||||
|
if config.TimeoutMS > 0 {
|
||||||
|
modelConfig.Timeout = time.Duration(config.TimeoutMS) * time.Millisecond
|
||||||
|
}
|
||||||
|
if config.MaxOutputTokens > 0 {
|
||||||
|
maxTokens := config.MaxOutputTokens
|
||||||
|
modelConfig.MaxCompletionTokens = &maxTokens
|
||||||
|
}
|
||||||
|
if isDashScopeQwenThinkingModel(config) {
|
||||||
|
modelConfig.ExtraFields = map[string]any{"enable_thinking": false}
|
||||||
|
}
|
||||||
|
model, err := einoopenai.NewChatModel(ctx, modelConfig)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("create Eino OpenAI-compatible model: %w", err)
|
||||||
|
}
|
||||||
|
return model, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func isDashScopeQwenThinkingModel(config models.AIConfig) bool {
|
||||||
|
baseURL := strings.ToLower(strings.TrimSpace(config.BaseURL))
|
||||||
|
modelName := strings.ToLower(strings.TrimSpace(config.ModelName))
|
||||||
|
return strings.Contains(baseURL, "dashscope.aliyuncs.com") && strings.HasPrefix(modelName, "qwen3")
|
||||||
|
}
|
||||||
|
|
||||||
|
type einoFunctionTool struct {
|
||||||
|
info *schema.ToolInfo
|
||||||
|
execute ai.ToolCallExecutor
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ einotool.InvokableTool = (*einoFunctionTool)(nil)
|
||||||
|
|
||||||
|
func newEinoFunctionTool(definition ai.ToolDefinition, execute ai.ToolCallExecutor) (*einoFunctionTool, error) {
|
||||||
|
if strings.TrimSpace(definition.Name) == "" || execute == nil {
|
||||||
|
return nil, fmt.Errorf("Eino tool name and executor are required")
|
||||||
|
}
|
||||||
|
info := &schema.ToolInfo{
|
||||||
|
Name: strings.TrimSpace(definition.Name),
|
||||||
|
Desc: strings.TrimSpace(definition.Description),
|
||||||
|
}
|
||||||
|
if len(definition.Parameters) > 0 {
|
||||||
|
data, err := json.Marshal(definition.Parameters)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("encode Eino tool schema: %w", err)
|
||||||
|
}
|
||||||
|
var params einojsonschema.Schema
|
||||||
|
if err := json.Unmarshal(data, ¶ms); err != nil {
|
||||||
|
return nil, fmt.Errorf("decode Eino tool schema: %w", err)
|
||||||
|
}
|
||||||
|
info.ParamsOneOf = schema.NewParamsOneOfByJSONSchema(¶ms)
|
||||||
|
}
|
||||||
|
return &einoFunctionTool{info: info, execute: execute}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *einoFunctionTool) Info(context.Context) (*schema.ToolInfo, error) {
|
||||||
|
return t.info, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *einoFunctionTool) InvokableRun(ctx context.Context, arguments string, _ ...einotool.Option) (string, error) {
|
||||||
|
result, err := t.execute(ctx, ai.ToolCall{Name: t.info.Name, Arguments: arguments})
|
||||||
|
if err == nil {
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
var interrupt *agentLoopInterruptError
|
||||||
|
if errors.As(err, &interrupt) {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
observation, marshalErr := json.Marshal(map[string]string{"error": err.Error()})
|
||||||
|
if marshalErr != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(observation), nil
|
||||||
|
}
|
||||||
145
internal/ai/application/runtime/evaluation.go
Normal file
145
internal/ai/application/runtime/evaluation.go
Normal file
@@ -0,0 +1,145 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/csv"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
)
|
||||||
|
|
||||||
|
// OfflineEvaluationCase is an isolated customer-service evaluation sample.
|
||||||
|
// Expectations are intentionally declarative so the same baseline can evolve
|
||||||
|
// without changing the runner's request contract.
|
||||||
|
type OfflineEvaluationCase struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Message string `json:"message"`
|
||||||
|
History []string `json:"history,omitempty"`
|
||||||
|
Expect map[string]any `json:"expect,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type OfflineEvaluationResult struct {
|
||||||
|
CaseID string `json:"caseId"`
|
||||||
|
Category string `json:"category"`
|
||||||
|
Passed bool `json:"passed"`
|
||||||
|
ReplyText string `json:"replyText"`
|
||||||
|
Interrupted bool `json:"interrupted"`
|
||||||
|
Error string `json:"error,omitempty"`
|
||||||
|
Finding string `json:"finding,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type OfflineEvaluationReport struct {
|
||||||
|
Total int `json:"total"`
|
||||||
|
Passed int `json:"passed"`
|
||||||
|
Results []OfflineEvaluationResult `json:"results"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// OfflineEvaluationRunner executes only isolated Debug requests. The supplied
|
||||||
|
// runner makes it testable without a real model and lets callers choose an
|
||||||
|
// explicit Engine implementation for mode comparison.
|
||||||
|
type OfflineEvaluationRunner struct {
|
||||||
|
run func(context.Context, RunInput) (*RunResult, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewOfflineEvaluationRunner(run func(context.Context, RunInput) (*RunResult, error)) *OfflineEvaluationRunner {
|
||||||
|
return &OfflineEvaluationRunner{run: run}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *OfflineEvaluationRunner) Run(ctx context.Context, agent models.AIAgent, config models.AIConfig, cases []OfflineEvaluationCase) OfflineEvaluationReport {
|
||||||
|
report := OfflineEvaluationReport{Results: make([]OfflineEvaluationResult, 0, len(cases))}
|
||||||
|
for _, item := range cases {
|
||||||
|
result := OfflineEvaluationResult{CaseID: strings.TrimSpace(item.ID), Category: strings.TrimSpace(item.Category)}
|
||||||
|
if r == nil || r.run == nil {
|
||||||
|
result.Error, result.Finding = "evaluation runner is not configured", "runner_missing"
|
||||||
|
report.Results = append(report.Results, result)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
summary, err := r.run(ctx, RunInput{
|
||||||
|
Conversation: models.Conversation{AIAgentID: agent.ID, LastMessageSummary: strings.Join(item.History, "\n")},
|
||||||
|
UserMessage: models.Message{SenderType: enums.IMSenderTypeCustomer, MessageType: enums.IMMessageTypeText, Content: strings.TrimSpace(item.Message), RequestID: "offline-eval:" + strings.TrimSpace(item.ID)},
|
||||||
|
AIAgent: agent,
|
||||||
|
AIConfig: config,
|
||||||
|
Debug: true,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
result.Error, result.Finding = err.Error(), "engine_error"
|
||||||
|
report.Results = append(report.Results, result)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if summary != nil {
|
||||||
|
result.ReplyText = strings.TrimSpace(summary.ReplyText)
|
||||||
|
result.Interrupted = summary.Interrupted
|
||||||
|
}
|
||||||
|
result.Passed, result.Finding = evaluateOfflineCase(item.Expect, summary)
|
||||||
|
if result.Passed {
|
||||||
|
report.Passed++
|
||||||
|
}
|
||||||
|
report.Results = append(report.Results, result)
|
||||||
|
}
|
||||||
|
report.Total = len(report.Results)
|
||||||
|
return report
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r OfflineEvaluationReport) CSV() (string, error) {
|
||||||
|
var output strings.Builder
|
||||||
|
writer := csv.NewWriter(&output)
|
||||||
|
if err := writer.Write([]string{"caseId", "category", "passed", "interrupted", "finding", "error", "replyText"}); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
for _, item := range r.Results {
|
||||||
|
if err := writer.Write([]string{item.CaseID, item.Category, strconv.FormatBool(item.Passed), strconv.FormatBool(item.Interrupted), item.Finding, item.Error, item.ReplyText}); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
writer.Flush()
|
||||||
|
return output.String(), writer.Error()
|
||||||
|
}
|
||||||
|
|
||||||
|
func evaluateOfflineCase(expect map[string]any, summary *RunResult) (bool, string) {
|
||||||
|
if summary == nil || strings.TrimSpace(summary.ReplyText) == "" {
|
||||||
|
return false, "empty_reply"
|
||||||
|
}
|
||||||
|
if requiresConfirmation, _ := expect["requiresConfirmation"].(bool); requiresConfirmation && !summary.Interrupted {
|
||||||
|
return false, "confirmation_not_reached"
|
||||||
|
}
|
||||||
|
if maxWrites, ok := evaluationExpectationInt(expect["maxWriteToolCalls"]); ok {
|
||||||
|
if maxWrites < 0 {
|
||||||
|
return false, "invalid_expectation"
|
||||||
|
}
|
||||||
|
if writeToolCalls(summary) > maxWrites {
|
||||||
|
return false, "write_tool_limit_exceeded"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true, ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func evaluationExpectationInt(value any) (int, bool) {
|
||||||
|
switch item := value.(type) {
|
||||||
|
case int:
|
||||||
|
return item, true
|
||||||
|
case int64:
|
||||||
|
return int(item), true
|
||||||
|
case float64:
|
||||||
|
return int(item), item == float64(int(item))
|
||||||
|
default:
|
||||||
|
return 0, false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeToolCalls(summary *RunResult) int {
|
||||||
|
if summary == nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
count := 0
|
||||||
|
for _, code := range summary.InvokedToolCodes {
|
||||||
|
switch toolx.NormalizeToolCodeAlias(code) {
|
||||||
|
case toolx.GraphCreateTicketConfirm.Code, toolx.GraphHandoffConversation.Code:
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return count
|
||||||
|
}
|
||||||
220
internal/ai/application/runtime/service.go
Normal file
220
internal/ai/application/runtime/service.go
Normal file
@@ -0,0 +1,220 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
workflowexecutor "agent-desk/internal/ai/runtime/workflow"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Service struct {
|
||||||
|
engine *AgentLoopEngine
|
||||||
|
}
|
||||||
|
|
||||||
|
const (
|
||||||
|
workflowRunStatusCompleted = 1
|
||||||
|
workflowRunStatusInterrupted = 2
|
||||||
|
workflowRunStatusFailed = 3
|
||||||
|
)
|
||||||
|
|
||||||
|
func NewService() *Service {
|
||||||
|
return NewServiceWithEngine(NewAgentLoopEngine())
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewServiceWithEngine(engine *AgentLoopEngine) *Service {
|
||||||
|
return &Service{engine: engine}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Service) Run(ctx context.Context, req RunInput) (*RunResult, error) {
|
||||||
|
return s.engine.Run(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Service) Resume(ctx context.Context, req ResumeInput) (*RunResult, error) {
|
||||||
|
return s.engine.Resume(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Service) RunOfflineEvaluation(ctx context.Context, agent models.AIAgent, config models.AIConfig, cases []OfflineEvaluationCase) (OfflineEvaluationReport, error) {
|
||||||
|
runner := NewOfflineEvaluationRunner(s.engine.Run)
|
||||||
|
return runner.Run(ctx, agent, config, cases), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func toWorkflowResult(result *workflowexecutor.Result, modelName string, workflow resolvedWorkflow, workflowRunID int64) *RunResult {
|
||||||
|
if result == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
trace := map[string]any{
|
||||||
|
"status": result.Status,
|
||||||
|
"workflowId": workflow.WorkflowID,
|
||||||
|
"workflowVersionId": workflow.VersionID,
|
||||||
|
"workflowRunId": workflowRunID,
|
||||||
|
"nodePath": result.NodePath,
|
||||||
|
}
|
||||||
|
traceData, _ := json.Marshal(trace)
|
||||||
|
return &RunResult{
|
||||||
|
Status: result.Status,
|
||||||
|
ReplyText: result.ReplyText,
|
||||||
|
ModelName: modelName,
|
||||||
|
PromptTokens: result.PromptTokens,
|
||||||
|
CompletionTokens: result.CompletionTokens,
|
||||||
|
RetrieverCount: result.RetrieverCount,
|
||||||
|
WorkflowID: workflow.WorkflowID,
|
||||||
|
WorkflowVersionID: workflow.VersionID,
|
||||||
|
WorkflowRunID: workflowRunID,
|
||||||
|
WorkflowNodePath: append([]string(nil), result.NodePath...),
|
||||||
|
TraceData: string(traceData),
|
||||||
|
CheckPointID: result.CheckPointID,
|
||||||
|
CheckPointData: result.CheckPointData,
|
||||||
|
Interrupted: result.Interrupted,
|
||||||
|
Interrupts: toWorkflowInterruptSummaries(result.Interrupts),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func toWorkflowInterruptSummaries(items []workflowexecutor.InterruptSummary) []InterruptContextSummary {
|
||||||
|
if len(items) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
ret := make([]InterruptContextSummary, 0, len(items))
|
||||||
|
for _, item := range items {
|
||||||
|
ret = append(ret, InterruptContextSummary{
|
||||||
|
Type: item.Type,
|
||||||
|
ID: item.ID,
|
||||||
|
InfoPreview: item.InfoPreview,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeWorkflowRun(req RunInput, workflow resolvedWorkflow, result *workflowexecutor.Result, errorMessage string) (int64, error) {
|
||||||
|
return writeWorkflowRunWithExistingID(req, workflow, result, errorMessage, 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeWorkflowRunWithExistingID(req RunInput, workflow resolvedWorkflow, result *workflowexecutor.Result, errorMessage string, existingRunID int64) (int64, error) {
|
||||||
|
if result == nil {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
endedAt := now
|
||||||
|
nodeTypes := make(map[string]string, len(workflow.Definition.Nodes))
|
||||||
|
for _, node := range workflow.Definition.Nodes {
|
||||||
|
nodeTypes[node.ID] = node.Type
|
||||||
|
}
|
||||||
|
runStatus := workflowRunStatus(result.Status, errorMessage)
|
||||||
|
var runID int64
|
||||||
|
err := sqls.WithTransaction(func(ctx *sqls.TxContext) error {
|
||||||
|
run := repositories.AIWorkflowRunRepository.Get(ctx.Tx, existingRunID)
|
||||||
|
if run == nil {
|
||||||
|
run = &models.AIWorkflowRun{
|
||||||
|
WorkflowID: workflow.WorkflowID,
|
||||||
|
WorkflowVersionID: workflow.VersionID,
|
||||||
|
ConversationID: req.Conversation.ID,
|
||||||
|
AIAgentID: req.AIAgent.ID,
|
||||||
|
MessageID: req.UserMessage.ID,
|
||||||
|
Status: runStatus,
|
||||||
|
StartedAt: now,
|
||||||
|
EndedAt: &endedAt,
|
||||||
|
InterruptType: firstWorkflowInterruptType(result),
|
||||||
|
InterruptNodeID: firstWorkflowInterruptNodeID(result),
|
||||||
|
ErrorMessage: errorMessage,
|
||||||
|
}
|
||||||
|
if err := repositories.AIWorkflowRunRepository.Create(ctx.Tx, run); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
} else if err := repositories.AIWorkflowRunRepository.Updates(ctx.Tx, run.ID, map[string]any{
|
||||||
|
"status": runStatus,
|
||||||
|
"ended_at": &endedAt,
|
||||||
|
"interrupt_type": firstWorkflowInterruptType(result),
|
||||||
|
"interrupt_node_id": firstWorkflowInterruptNodeID(result),
|
||||||
|
"error_message": errorMessage,
|
||||||
|
"updated_at": now,
|
||||||
|
}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
runID = run.ID
|
||||||
|
nodeTraces := result.NodeTraces
|
||||||
|
if len(nodeTraces) == 0 {
|
||||||
|
nodeTraces = fallbackWorkflowNodeTraces(result.NodePath, nodeTypes, result.Status)
|
||||||
|
}
|
||||||
|
for _, nodeTrace := range nodeTraces {
|
||||||
|
nodeRun := &models.AIWorkflowNodeRun{
|
||||||
|
WorkflowRunID: run.ID,
|
||||||
|
NodeID: nodeTrace.NodeID,
|
||||||
|
NodeType: firstNonEmpty(nodeTrace.NodeType, nodeTypes[nodeTrace.NodeID]),
|
||||||
|
Status: workflowRunStatus(nodeTrace.Status, nodeTrace.ErrorMessage),
|
||||||
|
InputPreview: nodeTrace.InputPreview,
|
||||||
|
OutputPreview: nodeTrace.OutputPreview,
|
||||||
|
ErrorMessage: nodeTrace.ErrorMessage,
|
||||||
|
StartedAt: now,
|
||||||
|
EndedAt: &endedAt,
|
||||||
|
DurationMS: nodeTrace.DurationMS,
|
||||||
|
}
|
||||||
|
if err := repositories.AIWorkflowNodeRunRepository.Create(ctx.Tx, nodeRun); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
return runID, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func workflowAgentRunStatus(status string, errorMessage string) string {
|
||||||
|
if strings.TrimSpace(errorMessage) != "" || strings.TrimSpace(status) == "error" {
|
||||||
|
return "failed"
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(status) == "interrupted" {
|
||||||
|
return "interrupted"
|
||||||
|
}
|
||||||
|
return "completed"
|
||||||
|
}
|
||||||
|
|
||||||
|
func workflowRunStatus(status string, errorMessage string) int {
|
||||||
|
if strings.TrimSpace(errorMessage) != "" || strings.TrimSpace(status) == "error" {
|
||||||
|
return workflowRunStatusFailed
|
||||||
|
}
|
||||||
|
switch strings.TrimSpace(status) {
|
||||||
|
case "interrupted":
|
||||||
|
return workflowRunStatusInterrupted
|
||||||
|
default:
|
||||||
|
return workflowRunStatusCompleted
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func fallbackWorkflowNodeTraces(nodePath []string, nodeTypes map[string]string, status string) []workflowexecutor.NodeTrace {
|
||||||
|
ret := make([]workflowexecutor.NodeTrace, 0, len(nodePath))
|
||||||
|
for _, nodeID := range nodePath {
|
||||||
|
ret = append(ret, workflowexecutor.NodeTrace{
|
||||||
|
NodeID: nodeID,
|
||||||
|
NodeType: nodeTypes[nodeID],
|
||||||
|
Status: status,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstWorkflowInterruptType(result *workflowexecutor.Result) string {
|
||||||
|
if result == nil || len(result.Interrupts) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(result.Interrupts[0].Type)
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstWorkflowInterruptNodeID(result *workflowexecutor.Result) string {
|
||||||
|
if result == nil || len(result.Interrupts) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(result.Interrupts[0].ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstNonEmpty(items ...string) string {
|
||||||
|
for _, item := range items {
|
||||||
|
if strings.TrimSpace(item) != "" {
|
||||||
|
return strings.TrimSpace(item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
113
internal/ai/application/runtime/types.go
Normal file
113
internal/ai/application/runtime/types.go
Normal file
@@ -0,0 +1,113 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// RunInput is the normalized, fully prepared input for the Agent Loop.
|
||||||
|
type RunInput struct {
|
||||||
|
Conversation models.Conversation
|
||||||
|
UserMessage models.Message
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
AIConfig models.AIConfig
|
||||||
|
CheckPointID string
|
||||||
|
Debug bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResumeInput extends the prepared input with an approved interrupt payload.
|
||||||
|
type ResumeInput struct {
|
||||||
|
Conversation models.Conversation
|
||||||
|
UserMessage models.Message
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
AIConfig models.AIConfig
|
||||||
|
CheckPointID string
|
||||||
|
ResumeData map[string]string
|
||||||
|
Debug bool
|
||||||
|
}
|
||||||
|
|
||||||
|
type InterruptContextSummary struct {
|
||||||
|
Type string `json:"type,omitempty"`
|
||||||
|
ID string `json:"id"`
|
||||||
|
DisplayName string `json:"displayName,omitempty"`
|
||||||
|
PromptText string `json:"promptText,omitempty"`
|
||||||
|
InfoPreview string `json:"infoPreview,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// RunResult is the normalized Agent Loop result.
|
||||||
|
type RunResult struct {
|
||||||
|
RunID string
|
||||||
|
Status string
|
||||||
|
ReplyText string
|
||||||
|
PlannedSkillID int64
|
||||||
|
PlannedSkillName string
|
||||||
|
SkillAllowedToolCodes []string
|
||||||
|
ModelName string
|
||||||
|
PromptTokens int
|
||||||
|
CompletionTokens int
|
||||||
|
HistoryMessageCount int
|
||||||
|
RetrieverCount int
|
||||||
|
ToolCallCount int
|
||||||
|
InvokedToolCodes []string
|
||||||
|
WorkflowID int64
|
||||||
|
WorkflowVersionID int64
|
||||||
|
WorkflowRunID int64
|
||||||
|
AgentRunID int64
|
||||||
|
WorkflowNodePath []string
|
||||||
|
CheckPointID string
|
||||||
|
CheckPointData string
|
||||||
|
Interrupted bool
|
||||||
|
HandoffRequested bool
|
||||||
|
HandoffReason string
|
||||||
|
ConversationDecision *ConversationDecision
|
||||||
|
Interrupts []InterruptContextSummary
|
||||||
|
TraceData string
|
||||||
|
ErrorMessage string
|
||||||
|
}
|
||||||
|
|
||||||
|
type ConversationAction string
|
||||||
|
|
||||||
|
const (
|
||||||
|
ConversationActionReply ConversationAction = "reply"
|
||||||
|
ConversationActionHandoff ConversationAction = "handoff"
|
||||||
|
ConversationActionAskHandoffConfirmation ConversationAction = "ask_handoff_confirmation"
|
||||||
|
)
|
||||||
|
|
||||||
|
type HandoffInitiator string
|
||||||
|
|
||||||
|
const (
|
||||||
|
HandoffInitiatorNone HandoffInitiator = "none"
|
||||||
|
HandoffInitiatorCustomer HandoffInitiator = "customer"
|
||||||
|
HandoffInitiatorAgent HandoffInitiator = "agent"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ConversationDecision is the model's structured, non-side-effect decision.
|
||||||
|
// The runtime remains the only component that can execute a handoff.
|
||||||
|
type ConversationDecision struct {
|
||||||
|
Action ConversationAction `json:"action"`
|
||||||
|
Reason string `json:"reason"`
|
||||||
|
Reply string `json:"reply"`
|
||||||
|
HandoffInitiator HandoffInitiator `json:"handoffInitiator"`
|
||||||
|
HandoffConfirmed bool `json:"handoffConfirmed"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type StreamEventType string
|
||||||
|
|
||||||
|
const (
|
||||||
|
StreamEventStarted StreamEventType = "started"
|
||||||
|
StreamEventStep StreamEventType = "step"
|
||||||
|
StreamEventOutput StreamEventType = "output"
|
||||||
|
StreamEventCompleted StreamEventType = "completed"
|
||||||
|
StreamEventFailed StreamEventType = "failed"
|
||||||
|
)
|
||||||
|
|
||||||
|
// StreamEvent is the transport-neutral event contract for future streaming.
|
||||||
|
type StreamEvent struct {
|
||||||
|
Type StreamEventType `json:"type"`
|
||||||
|
RunID string `json:"runId,omitempty"`
|
||||||
|
AgentRunID int64 `json:"agentRunId,omitempty"`
|
||||||
|
StepCode string `json:"stepCode,omitempty"`
|
||||||
|
Content string `json:"content,omitempty"`
|
||||||
|
Error string `json:"error,omitempty"`
|
||||||
|
OccurredAt time.Time `json:"occurredAt"`
|
||||||
|
}
|
||||||
37
internal/ai/application/runtime/workflow_runtime.go
Normal file
37
internal/ai/application/runtime/workflow_runtime.go
Normal file
@@ -0,0 +1,37 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/workflow/dsl"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
type resolvedWorkflow struct {
|
||||||
|
Definition dsl.Definition
|
||||||
|
WorkflowID int64
|
||||||
|
VersionID int64
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveWorkflowVersion(workflowVersionID int64) (resolvedWorkflow, error) {
|
||||||
|
if workflowVersionID <= 0 {
|
||||||
|
return resolvedWorkflow{}, errorsx.InvalidParam("workflow version is required")
|
||||||
|
}
|
||||||
|
version := repositories.AIWorkflowVersionRepository.Get(sqls.DB(), workflowVersionID)
|
||||||
|
if version == nil || version.Status != enums.StatusOk {
|
||||||
|
return resolvedWorkflow{}, errorsx.InvalidParam("workflow version does not exist")
|
||||||
|
}
|
||||||
|
var def dsl.Definition
|
||||||
|
if err := json.Unmarshal([]byte(version.Definition), &def); err != nil {
|
||||||
|
return resolvedWorkflow{}, errorsx.InvalidParam("workflow definition is invalid")
|
||||||
|
}
|
||||||
|
return resolvedWorkflow{
|
||||||
|
Definition: def,
|
||||||
|
WorkflowID: version.WorkflowID,
|
||||||
|
VersionID: version.ID,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
101
internal/ai/embedding.go
Normal file
101
internal/ai/embedding.go
Normal file
@@ -0,0 +1,101 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
openai "github.com/openai/openai-go/v3"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
)
|
||||||
|
|
||||||
|
type EmbeddingResult struct {
|
||||||
|
Vector []float32
|
||||||
|
TokensUsed int
|
||||||
|
ModelName string
|
||||||
|
Dimension int
|
||||||
|
}
|
||||||
|
|
||||||
|
type embedding struct{}
|
||||||
|
|
||||||
|
var Embedding = &embedding{}
|
||||||
|
|
||||||
|
func (s *embedding) GetModel(ctx context.Context) (*models.AIConfig, error) {
|
||||||
|
config, err := GetEnabledAIConfig(enums.AIModelTypeEmbedding)
|
||||||
|
if err != nil {
|
||||||
|
return nil, errorsx.BusinessErrorI18n(2001, "error.embeddingModel.noneEnabled")
|
||||||
|
}
|
||||||
|
return config, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *embedding) GenerateEmbedding(ctx context.Context, text string) (*EmbeddingResult, error) {
|
||||||
|
if text == "" {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0215")
|
||||||
|
}
|
||||||
|
|
||||||
|
result, err := s.callEmbeddingAPI(ctx, text)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *embedding) GenerateBatchEmbeddings(ctx context.Context, texts []string) ([]EmbeddingResult, error) {
|
||||||
|
if len(texts) == 0 {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0216")
|
||||||
|
}
|
||||||
|
|
||||||
|
results := make([]EmbeddingResult, 0, len(texts))
|
||||||
|
for _, text := range texts {
|
||||||
|
result, err := s.callEmbeddingAPI(ctx, text)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to generate embedding for text: %w", err)
|
||||||
|
}
|
||||||
|
results = append(results, *result)
|
||||||
|
}
|
||||||
|
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *embedding) callEmbeddingAPI(ctx context.Context, text string) (*EmbeddingResult, error) {
|
||||||
|
config, err := GetEnabledAIConfig(enums.AIModelTypeEmbedding)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
client := newOpenAIClient(*config)
|
||||||
|
embeddingResp, err := client.Embeddings.New(ctx, openai.EmbeddingNewParams{
|
||||||
|
Input: openai.EmbeddingNewParamsInputUnion{
|
||||||
|
OfString: openai.String(text),
|
||||||
|
},
|
||||||
|
Model: openai.EmbeddingModel(config.ModelName),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to call embedding api: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(embeddingResp.Data) == 0 {
|
||||||
|
return nil, fmt.Errorf("no embedding data in response")
|
||||||
|
}
|
||||||
|
vector := make([]float32, 0, len(embeddingResp.Data[0].Embedding))
|
||||||
|
for _, item := range embeddingResp.Data[0].Embedding {
|
||||||
|
vector = append(vector, float32(item))
|
||||||
|
}
|
||||||
|
|
||||||
|
return &EmbeddingResult{
|
||||||
|
Vector: vector,
|
||||||
|
TokensUsed: int(embeddingResp.Usage.TotalTokens),
|
||||||
|
ModelName: embeddingResp.Model,
|
||||||
|
Dimension: len(vector),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *embedding) GetDimension(ctx context.Context) (int, error) {
|
||||||
|
model, err := s.GetModel(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
return model.Dimension, nil
|
||||||
|
}
|
||||||
312
internal/ai/handoff_intent.go
Normal file
312
internal/ai/handoff_intent.go
Normal file
@@ -0,0 +1,312 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"unicode"
|
||||||
|
)
|
||||||
|
|
||||||
|
// IsCustomerHandoffRequest determines whether a user's message expresses an intent
|
||||||
|
// to speak with a human operator, agent, specialist, manager, or live person, across
|
||||||
|
// varied natural phrasing, colloquialisms, and multiple languages (Russian, English,
|
||||||
|
// Spanish, German, French, Chinese, etc.).
|
||||||
|
func IsCustomerHandoffRequest(text string) bool {
|
||||||
|
raw := strings.TrimSpace(text)
|
||||||
|
if raw == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
lower := strings.ToLower(raw)
|
||||||
|
|
||||||
|
// Direct Chinese keywords
|
||||||
|
if IsHan(lower) {
|
||||||
|
for _, kw := range []string{
|
||||||
|
"人工", "转人工", "真人", "人工客服", "接人工", "转接人工", "转坐席", "人工服务", "人工专员",
|
||||||
|
} {
|
||||||
|
if strings.Contains(lower, kw) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clean punctuation for token / n-gram checks
|
||||||
|
cleaned := strings.Map(func(r rune) rune {
|
||||||
|
if unicode.IsPunct(r) || unicode.IsSymbol(r) {
|
||||||
|
return ' '
|
||||||
|
}
|
||||||
|
return r
|
||||||
|
}, lower)
|
||||||
|
fields := strings.Fields(cleaned)
|
||||||
|
if len(fields) == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
normalized := " " + strings.Join(fields, " ") + " "
|
||||||
|
|
||||||
|
// Russian handoff patterns & stems:
|
||||||
|
// Anti-bot phrases: "хватит бота", "не хочу бота", "отключи бота", "надоел бот"
|
||||||
|
if strings.Contains(normalized, " хватит бота ") ||
|
||||||
|
strings.Contains(normalized, " хватит робота ") ||
|
||||||
|
strings.Contains(normalized, " не хочу бота ") ||
|
||||||
|
strings.Contains(normalized, " не хочу общаться с ботом ") ||
|
||||||
|
strings.Contains(normalized, " надоел бот ") ||
|
||||||
|
strings.Contains(normalized, " отключи бота ") ||
|
||||||
|
strings.Contains(normalized, " без бота ") ||
|
||||||
|
strings.Contains(normalized, " убрать бота ") ||
|
||||||
|
strings.Contains(normalized, " stop bot ") ||
|
||||||
|
strings.Contains(normalized, " no bot ") ||
|
||||||
|
strings.Contains(normalized, " hate bot ") ||
|
||||||
|
strings.Contains(normalized, " not a bot ") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
russianTargetStems := []string{
|
||||||
|
"оператор", "человек", "специалист", "сотрудник", "консультант",
|
||||||
|
"менеджер", "живо", "живой", "поддержк", "саппорт", "люди", "людей",
|
||||||
|
"админ", "администратор",
|
||||||
|
}
|
||||||
|
|
||||||
|
russianActionStems := []string{
|
||||||
|
"переведи", "переключ", "соедин", "позов", "дайт", "дай", "свяж", "связать",
|
||||||
|
"вызов", "вызовите", "хочу", "нужен", "нужна", "нужно", "можно", "позовите",
|
||||||
|
"соедините", "переведите", "переключите", "пригласи", "пригласите",
|
||||||
|
"перенаправ", "перенаправь", "перенаправьте", "передай", "передайте",
|
||||||
|
"направь", "направьте",
|
||||||
|
}
|
||||||
|
|
||||||
|
// Action + Target combinations in Russian (e.g. "позови сотрудника", "соедини со специалистом", "хочу человека")
|
||||||
|
for _, action := range russianActionStems {
|
||||||
|
if strings.Contains(normalized, " "+action) {
|
||||||
|
for _, target := range russianTargetStems {
|
||||||
|
if strings.Contains(normalized, " "+target) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Single-word requests (e.g. user just typed "оператор", "человек", "operator", "human")
|
||||||
|
if normalized == " оператор " ||
|
||||||
|
normalized == " оператора " ||
|
||||||
|
normalized == " человек " ||
|
||||||
|
normalized == " специалист " ||
|
||||||
|
normalized == " support " ||
|
||||||
|
normalized == " operator " ||
|
||||||
|
normalized == " agent " ||
|
||||||
|
normalized == " human " {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Direct targets in Russian with prepositions or adjectives (e.g. "к оператору", "с человеком", "со специалистом", "живой человек")
|
||||||
|
if strings.Contains(normalized, " к оператору") ||
|
||||||
|
strings.Contains(normalized, " на оператора") ||
|
||||||
|
strings.Contains(normalized, " с оператором") ||
|
||||||
|
strings.Contains(normalized, " со специалистом") ||
|
||||||
|
strings.Contains(normalized, " к специалисту") ||
|
||||||
|
strings.Contains(normalized, " с человеком") ||
|
||||||
|
strings.Contains(normalized, " с сотрудником") ||
|
||||||
|
strings.Contains(normalized, " живой человек") ||
|
||||||
|
strings.Contains(normalized, " живого человека") ||
|
||||||
|
strings.Contains(normalized, " живым человеком") ||
|
||||||
|
strings.Contains(normalized, " поговорить с кем") ||
|
||||||
|
strings.Contains(normalized, " поговорить со ") ||
|
||||||
|
strings.Contains(normalized, " реальный человек") ||
|
||||||
|
strings.Contains(normalized, " живой есть") ||
|
||||||
|
strings.Contains(normalized, " есть живой") ||
|
||||||
|
strings.Contains(normalized, " есть кто живой") ||
|
||||||
|
strings.Contains(normalized, " можно человека") ||
|
||||||
|
strings.Contains(normalized, " дай человека") ||
|
||||||
|
strings.Contains(normalized, " дай оператора") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// English handoff patterns
|
||||||
|
englishTargets := []string{
|
||||||
|
"operator", "agent", "human", "representative", "person", "specialist", "someone", "support",
|
||||||
|
}
|
||||||
|
englishActions := []string{
|
||||||
|
"talk to", "speak to", "speak with", "connect", "transfer", "get me", "call", "reach", "need", "want", "let me",
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, action := range englishActions {
|
||||||
|
if strings.Contains(normalized, " "+action) {
|
||||||
|
for _, target := range englishTargets {
|
||||||
|
if strings.Contains(normalized, " "+target) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if strings.Contains(normalized, " real person ") ||
|
||||||
|
strings.Contains(normalized, " live agent ") ||
|
||||||
|
strings.Contains(normalized, " live person ") ||
|
||||||
|
strings.Contains(normalized, " human agent ") ||
|
||||||
|
strings.Contains(normalized, " human please ") ||
|
||||||
|
strings.Contains(normalized, " customer service ") ||
|
||||||
|
strings.Contains(normalized, " human support ") ||
|
||||||
|
strings.Contains(normalized, " talk to someone ") ||
|
||||||
|
strings.Contains(normalized, " speak to someone ") ||
|
||||||
|
strings.Contains(normalized, " representative please ") ||
|
||||||
|
strings.Contains(normalized, " representative ") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Spanish handoff patterns
|
||||||
|
if strings.Contains(normalized, " hablar con ") ||
|
||||||
|
strings.Contains(normalized, " pasar con ") ||
|
||||||
|
strings.Contains(normalized, " contactar con ") ||
|
||||||
|
strings.Contains(normalized, " quiero hablar con ") ||
|
||||||
|
strings.Contains(normalized, " necesito hablar con ") {
|
||||||
|
if strings.Contains(normalized, " persona") ||
|
||||||
|
strings.Contains(normalized, " humano") ||
|
||||||
|
strings.Contains(normalized, " agente") ||
|
||||||
|
strings.Contains(normalized, " operador") ||
|
||||||
|
strings.Contains(normalized, " alguien") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if strings.Contains(normalized, " atencion al cliente") ||
|
||||||
|
strings.Contains(normalized, " persona real") ||
|
||||||
|
strings.Contains(normalized, " agente humano") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// German handoff patterns
|
||||||
|
if strings.Contains(normalized, " mitarbeiter ") ||
|
||||||
|
strings.Contains(normalized, " echter mensch ") ||
|
||||||
|
strings.Contains(normalized, " mit jemandem sprechen ") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// French handoff patterns
|
||||||
|
if strings.Contains(normalized, " parler a quelqu un ") ||
|
||||||
|
strings.Contains(normalized, " conseiller ") ||
|
||||||
|
strings.Contains(normalized, " agent humain ") ||
|
||||||
|
strings.Contains(normalized, " etre mis en relation ") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsModelClaimingHandoff checks whether an AI model's text reply claims that it is
|
||||||
|
// transferring, connecting, or forwarding the conversation to a human operator/agent.
|
||||||
|
func IsModelClaimingHandoff(text string) bool {
|
||||||
|
raw := strings.TrimSpace(text)
|
||||||
|
if raw == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
lower := strings.ToLower(raw)
|
||||||
|
|
||||||
|
// Direct Russian phrases (exact matches)
|
||||||
|
russianModelClaims := []string{
|
||||||
|
"перевожу вас на оператора",
|
||||||
|
"переключаю вас на оператора",
|
||||||
|
"перевожу на оператора",
|
||||||
|
"переключаю на оператора",
|
||||||
|
"соединяю с оператором",
|
||||||
|
"соединяю со специалистом",
|
||||||
|
"передаю диалог оператору",
|
||||||
|
"передаю ваш запрос оператору",
|
||||||
|
"перевожу диалог на оператора",
|
||||||
|
"переключаю на специалиста",
|
||||||
|
"перевожу на специалиста",
|
||||||
|
"соединяю с живым",
|
||||||
|
"связываю вас с оператором",
|
||||||
|
"сейчас переведу на оператора",
|
||||||
|
"сейчас переключу на оператора",
|
||||||
|
"перенаправляю оператору",
|
||||||
|
"перенаправляю вас к оператору",
|
||||||
|
"перенаправляю вас на оператора",
|
||||||
|
"перенаправляю на оператора",
|
||||||
|
"перенаправляю вас на специалиста",
|
||||||
|
"перенаправляю на специалиста",
|
||||||
|
"перенаправляю в службу поддержки",
|
||||||
|
"перенаправляю вас",
|
||||||
|
"передаю вас оператору",
|
||||||
|
"передаю диалог",
|
||||||
|
"направляю диалог",
|
||||||
|
"направляю ваш запрос",
|
||||||
|
}
|
||||||
|
for _, claim := range russianModelClaims {
|
||||||
|
if strings.Contains(lower, claim) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Flexible Russian stem matching (Action + Target)
|
||||||
|
// Catches arbitrary phrasing like "Похоже я не могу решить ваш вопрос, перенаправляю оператору. Ожидайте!"
|
||||||
|
ruActionStems := []string{"перенаправл", "перевод", "переключ", "соедин", "передаю", "передаем", "передам", "направляю", "направляем", "направлю", "связываю", "свяжу"}
|
||||||
|
ruTargetStems := []string{"оператор", "специалист", "сотрудник", "человек", "поддержк", "менеджер", "коллег", "живой"}
|
||||||
|
for _, act := range ruActionStems {
|
||||||
|
if strings.Contains(lower, act) {
|
||||||
|
for _, tgt := range ruTargetStems {
|
||||||
|
if strings.Contains(lower, tgt) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// English model handoff statements
|
||||||
|
englishModelClaims := []string{
|
||||||
|
"transferring you to an operator",
|
||||||
|
"transferring you to an agent",
|
||||||
|
"transferring you to a human",
|
||||||
|
"transferring you to human",
|
||||||
|
"transferring you to our support",
|
||||||
|
"connecting you to an operator",
|
||||||
|
"connecting you to an agent",
|
||||||
|
"connecting you to a human",
|
||||||
|
"connecting you with an agent",
|
||||||
|
"connecting you with an operator",
|
||||||
|
"routing you to an agent",
|
||||||
|
"handing you over to",
|
||||||
|
}
|
||||||
|
for _, claim := range englishModelClaims {
|
||||||
|
if strings.Contains(lower, claim) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Flexible English stem matching (Action + Target)
|
||||||
|
enActionStems := []string{"transfer", "connect", "rout", "hand over", "handing over", "forward", "pass you", "passing you"}
|
||||||
|
enTargetStems := []string{"operator", "agent", "human", "specialist", "representative", "support", "colleague"}
|
||||||
|
for _, act := range enActionStems {
|
||||||
|
if strings.Contains(lower, act) {
|
||||||
|
for _, tgt := range enTargetStems {
|
||||||
|
if strings.Contains(lower, tgt) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Chinese model handoff statements
|
||||||
|
chineseModelClaims := []string{
|
||||||
|
"为您转接人工",
|
||||||
|
"正在为您转接",
|
||||||
|
"为您接入人工",
|
||||||
|
"转接人工客服",
|
||||||
|
"转接人工",
|
||||||
|
"已为您转接",
|
||||||
|
}
|
||||||
|
for _, claim := range chineseModelClaims {
|
||||||
|
if strings.Contains(lower, claim) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Flexible Chinese stem matching
|
||||||
|
zhActionStems := []string{"转接", "转交", "接入", "转给", "转到", "交由"}
|
||||||
|
zhTargetStems := []string{"人工", "客服", "专员", "同事"}
|
||||||
|
for _, act := range zhActionStems {
|
||||||
|
if strings.Contains(lower, act) {
|
||||||
|
for _, tgt := range zhTargetStems {
|
||||||
|
if strings.Contains(lower, tgt) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return false
|
||||||
|
}
|
||||||
126
internal/ai/handoff_intent_test.go
Normal file
126
internal/ai/handoff_intent_test.go
Normal file
@@ -0,0 +1,126 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestIsCustomerHandoffRequest(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
input string
|
||||||
|
expected bool
|
||||||
|
}{
|
||||||
|
// Russian varied phrasing (not just "переведи на оператора")
|
||||||
|
{"переведи на оператора", true},
|
||||||
|
{"переведите на оператора пожалуйста", true},
|
||||||
|
{"соедините со специалистом", true},
|
||||||
|
{"позови сотрудника", true},
|
||||||
|
{"позовите сотрудника поддержки", true},
|
||||||
|
{"мне нужен живой человек", true},
|
||||||
|
{"хочу поговорить с человеком", true},
|
||||||
|
{"хочу поговорить со специалистом", true},
|
||||||
|
{"хватит бота", true},
|
||||||
|
{"хватит робота, дайте человека", true},
|
||||||
|
{"не хочу общаться с ботом", true},
|
||||||
|
{"можно человека?", true},
|
||||||
|
{"есть кто живой?", true},
|
||||||
|
{"свяжите с оператором", true},
|
||||||
|
{"переключи на оператора", true},
|
||||||
|
{"дай оператора", true},
|
||||||
|
{"дай человека", true},
|
||||||
|
{"перенаправь оператору", true},
|
||||||
|
{"перенаправьте на специалиста", true},
|
||||||
|
{"тупой бот, дай человека", true},
|
||||||
|
{"нужен оператор", true},
|
||||||
|
{"нужен специалист", true},
|
||||||
|
{"оператор", true},
|
||||||
|
{"человек", true},
|
||||||
|
{"к оператору", true},
|
||||||
|
|
||||||
|
// English varied phrasing
|
||||||
|
{"can I speak to someone?", true},
|
||||||
|
{"connect me to an agent", true},
|
||||||
|
{"I want to talk to a real person", true},
|
||||||
|
{"let me speak with support", true},
|
||||||
|
{"transfer me to a representative", true},
|
||||||
|
{"representative please", true},
|
||||||
|
{"human please", true},
|
||||||
|
{"stop bot and get human", true},
|
||||||
|
{"live agent", true},
|
||||||
|
{"operator", true},
|
||||||
|
{"human", true},
|
||||||
|
|
||||||
|
// Other languages
|
||||||
|
{"necesito hablar con una persona", true},
|
||||||
|
{"quiero hablar con un humano", true},
|
||||||
|
{"mit einem Mitarbeiter sprechen", true},
|
||||||
|
{"转人工", true},
|
||||||
|
{"请帮我转接人工客服", true},
|
||||||
|
|
||||||
|
// Non-handoff questions & greetings (including mentions of operators without transfer intent)
|
||||||
|
{"Привет хочу возрат денег", false},
|
||||||
|
{"Здравствуйте, как вернуть товар?", false},
|
||||||
|
{"Добрый день, сколько стоит доставка?", false},
|
||||||
|
{"Hello, how do I track my order?", false},
|
||||||
|
{"Can I get a refund for my purchase?", false},
|
||||||
|
{"Операторы не создаются", false},
|
||||||
|
{"Операторы не создаються", false},
|
||||||
|
{"Кто такой оператор?", false},
|
||||||
|
{"Оператор работает ночью?", false},
|
||||||
|
{"Спасибо, всё понятно!", false},
|
||||||
|
{"OK, thank you very much", false},
|
||||||
|
{"", false},
|
||||||
|
{" ", false},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.input, func(t *testing.T) {
|
||||||
|
got := IsCustomerHandoffRequest(tt.input)
|
||||||
|
if got != tt.expected {
|
||||||
|
t.Errorf("IsCustomerHandoffRequest(%q) = %v, want %v", tt.input, got, tt.expected)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsModelClaimingHandoff(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
input string
|
||||||
|
expected bool
|
||||||
|
}{
|
||||||
|
// Russian model claims
|
||||||
|
{"Похоже я не могу решить ваш вопрос, перенаправляю оператору. Ожидайте!", true},
|
||||||
|
{"Перевожу вас на оператора, пожалуйста подождите.", true},
|
||||||
|
{"Переключаю вас на оператора службы поддержки.", true},
|
||||||
|
{"Перенаправляю вас на специалиста.", true},
|
||||||
|
{"Соединяю со специалистом, оставайтесь на связи.", true},
|
||||||
|
{"Передаю ваш запрос оператору.", true},
|
||||||
|
{"Передаю диалог коллегам из поддержки.", true},
|
||||||
|
{"Сейчас переведу на оператора.", true},
|
||||||
|
|
||||||
|
// English model claims
|
||||||
|
{"Transferring you to an agent, please wait a moment.", true},
|
||||||
|
{"Connecting you to an operator right away.", true},
|
||||||
|
{"Connecting you to a human representative.", true},
|
||||||
|
{"Routing you to an agent.", true},
|
||||||
|
{"I am passing you to our support team.", true},
|
||||||
|
|
||||||
|
// Chinese model claims
|
||||||
|
{"正在为您转接人工客服,请稍候。", true},
|
||||||
|
{"为您转接人工客服处理。", true},
|
||||||
|
|
||||||
|
// Normal replies
|
||||||
|
{"Здравствуйте! Чем я могу вам помочь?", false},
|
||||||
|
{"Чтобы оформить возврат, перейдите в личный кабинет.", false},
|
||||||
|
{"To return an item, please open your order history.", false},
|
||||||
|
{"", false},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.input, func(t *testing.T) {
|
||||||
|
got := IsModelClaimingHandoff(tt.input)
|
||||||
|
if got != tt.expected {
|
||||||
|
t.Errorf("IsModelClaimingHandoff(%q) = %v, want %v", tt.input, got, tt.expected)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
53
internal/ai/language.go
Normal file
53
internal/ai/language.go
Normal file
@@ -0,0 +1,53 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"unicode"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DetectLanguage inspects a text and returns the predominant language: "ru", "zh", or "en".
|
||||||
|
func DetectLanguage(text string) string {
|
||||||
|
cyrillicCount := 0
|
||||||
|
hanCount := 0
|
||||||
|
latinCount := 0
|
||||||
|
|
||||||
|
for _, r := range text {
|
||||||
|
if unicode.Is(unicode.Cyrillic, r) {
|
||||||
|
cyrillicCount++
|
||||||
|
} else if unicode.Is(unicode.Han, r) {
|
||||||
|
hanCount++
|
||||||
|
} else if unicode.Is(unicode.Latin, r) {
|
||||||
|
latinCount++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if cyrillicCount > 0 && cyrillicCount >= hanCount {
|
||||||
|
return "ru"
|
||||||
|
}
|
||||||
|
if hanCount > 0 && hanCount > cyrillicCount {
|
||||||
|
return "zh"
|
||||||
|
}
|
||||||
|
if latinCount > 0 {
|
||||||
|
return "en"
|
||||||
|
}
|
||||||
|
return "en"
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsCyrillic returns true if the string contains at least one Cyrillic rune.
|
||||||
|
func IsCyrillic(text string) bool {
|
||||||
|
for _, r := range text {
|
||||||
|
if unicode.Is(unicode.Cyrillic, r) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsHan returns true if the string contains at least one Han (Chinese) rune.
|
||||||
|
func IsHan(text string) bool {
|
||||||
|
for _, r := range text {
|
||||||
|
if unicode.Is(unicode.Han, r) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
46
internal/ai/language_test.go
Normal file
46
internal/ai/language_test.go
Normal file
@@ -0,0 +1,46 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDetectLanguage(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
input string
|
||||||
|
expected string
|
||||||
|
}{
|
||||||
|
{"Russian greeting", "Здравствуйте, подскажите статус заказа", "ru"},
|
||||||
|
{"Russian short", "Привет!", "ru"},
|
||||||
|
{"English inquiry", "Hello, can you help me with a refund?", "en"},
|
||||||
|
{"Chinese inquiry", "你好,请问怎么办理退款?", "zh"},
|
||||||
|
{"Mixed Russian and English", "Привет, check this please", "ru"},
|
||||||
|
{"Mixed Chinese and English", "你好 hello world", "zh"},
|
||||||
|
{"Numbers only", "12345 67890", "en"},
|
||||||
|
{"Empty string", "", "en"},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
got := DetectLanguage(tt.input)
|
||||||
|
if got != tt.expected {
|
||||||
|
t.Errorf("DetectLanguage(%q) = %q, want %q", tt.input, got, tt.expected)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsCyrillicAndIsHan(t *testing.T) {
|
||||||
|
if !IsCyrillic("Привет") {
|
||||||
|
t.Errorf("IsCyrillic('Привет') should be true")
|
||||||
|
}
|
||||||
|
if IsCyrillic("Hello world") {
|
||||||
|
t.Errorf("IsCyrillic('Hello world') should be false")
|
||||||
|
}
|
||||||
|
if !IsHan("你好世界") {
|
||||||
|
t.Errorf("IsHan('你好世界') should be true")
|
||||||
|
}
|
||||||
|
if IsHan("Hello world") {
|
||||||
|
t.Errorf("IsHan('Hello world') should be false")
|
||||||
|
}
|
||||||
|
}
|
||||||
98
internal/ai/llm.go
Normal file
98
internal/ai/llm.go
Normal file
@@ -0,0 +1,98 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/common/strs"
|
||||||
|
openai "github.com/openai/openai-go/v3"
|
||||||
|
"github.com/openai/openai-go/v3/shared"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
type ChatCompletionResult struct {
|
||||||
|
Content string
|
||||||
|
ModelName string
|
||||||
|
PromptTokens int
|
||||||
|
CompletionTokens int
|
||||||
|
}
|
||||||
|
|
||||||
|
type llm struct{}
|
||||||
|
|
||||||
|
var LLM = &llm{}
|
||||||
|
|
||||||
|
func (s *llm) Chat(ctx context.Context, systemPrompt string, userPrompt string) (*ChatCompletionResult, error) {
|
||||||
|
config, err := GetEnabledAIConfig(enums.AIModelTypeLLM)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return s.ChatWithConfig(ctx, *config, systemPrompt, userPrompt)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *llm) ChatWithConfig(ctx context.Context, config models.AIConfig, systemPrompt string, userPrompt string) (*ChatCompletionResult, error) {
|
||||||
|
messages := make([]openai.ChatCompletionMessageParamUnion, 0, 2)
|
||||||
|
if strs.IsNotBlank(systemPrompt) {
|
||||||
|
messages = append(messages, openai.ChatCompletionMessageParamUnion{
|
||||||
|
OfSystem: &openai.ChatCompletionSystemMessageParam{
|
||||||
|
Content: openai.ChatCompletionSystemMessageParamContentUnion{
|
||||||
|
OfString: openai.String(systemPrompt),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
messages = append(messages, openai.ChatCompletionMessageParamUnion{
|
||||||
|
OfUser: &openai.ChatCompletionUserMessageParam{
|
||||||
|
Content: openai.ChatCompletionUserMessageParamContentUnion{
|
||||||
|
OfString: openai.String(userPrompt),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
params := openai.ChatCompletionNewParams{
|
||||||
|
Messages: messages,
|
||||||
|
Model: shared.ChatModel(config.ModelName),
|
||||||
|
}
|
||||||
|
if config.MaxOutputTokens > 0 {
|
||||||
|
params.MaxCompletionTokens = openai.Int(int64(config.MaxOutputTokens))
|
||||||
|
}
|
||||||
|
applyProviderSpecificChatParams(¶ms, config)
|
||||||
|
|
||||||
|
client := newOpenAIClient(config)
|
||||||
|
chatResp, err := client.Chat.Completions.New(ctx, params)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to call llm api (model=%s provider=%s system_chars=%d user_chars=%d max_output_tokens=%d): %w",
|
||||||
|
config.ModelName, config.Provider, utf8.RuneCountInString(systemPrompt), utf8.RuneCountInString(userPrompt), config.MaxOutputTokens, err)
|
||||||
|
}
|
||||||
|
if len(chatResp.Choices) == 0 {
|
||||||
|
return nil, fmt.Errorf("no llm choices in response")
|
||||||
|
}
|
||||||
|
|
||||||
|
content := strings.TrimSpace(chatResp.Choices[0].Message.Content)
|
||||||
|
return &ChatCompletionResult{
|
||||||
|
Content: content,
|
||||||
|
ModelName: config.ModelName,
|
||||||
|
PromptTokens: int(chatResp.Usage.PromptTokens),
|
||||||
|
CompletionTokens: int(chatResp.Usage.CompletionTokens),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func applyProviderSpecificChatParams(params *openai.ChatCompletionNewParams, config models.AIConfig) {
|
||||||
|
if params == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if isDashScopeQwenThinkingModel(config) {
|
||||||
|
params.SetExtraFields(map[string]any{
|
||||||
|
"enable_thinking": false,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func isDashScopeQwenThinkingModel(config models.AIConfig) bool {
|
||||||
|
baseURL := strings.ToLower(strings.TrimSpace(config.BaseURL))
|
||||||
|
modelName := strings.ToLower(strings.TrimSpace(config.ModelName))
|
||||||
|
return strings.Contains(baseURL, "dashscope.aliyuncs.com") && strings.HasPrefix(modelName, "qwen3")
|
||||||
|
}
|
||||||
43
internal/ai/llm_test.go
Normal file
43
internal/ai/llm_test.go
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
openai "github.com/openai/openai-go/v3"
|
||||||
|
"github.com/openai/openai-go/v3/shared"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestApplyProviderSpecificChatParamsIncludesDashScopeThinkingFlag(t *testing.T) {
|
||||||
|
params := openai.ChatCompletionNewParams{
|
||||||
|
Messages: []openai.ChatCompletionMessageParamUnion{
|
||||||
|
{
|
||||||
|
OfUser: &openai.ChatCompletionUserMessageParam{
|
||||||
|
Content: openai.ChatCompletionUserMessageParamContentUnion{
|
||||||
|
OfString: openai.String("hello"),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Model: shared.ChatModel("qwen3.5-plus"),
|
||||||
|
}
|
||||||
|
|
||||||
|
applyProviderSpecificChatParams(¶ms, models.AIConfig{
|
||||||
|
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||||
|
ModelName: "qwen3.5-plus",
|
||||||
|
})
|
||||||
|
|
||||||
|
raw, err := json.Marshal(params)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal params: %v", err)
|
||||||
|
}
|
||||||
|
var body map[string]any
|
||||||
|
if err := json.Unmarshal(raw, &body); err != nil {
|
||||||
|
t.Fatalf("unmarshal params: %v", err)
|
||||||
|
}
|
||||||
|
if got, ok := body["enable_thinking"].(bool); !ok || got {
|
||||||
|
t.Fatalf("expected enable_thinking=false in request body, got body=%s", raw)
|
||||||
|
}
|
||||||
|
}
|
||||||
200
internal/ai/mcps/client.go
Normal file
200
internal/ai/mcps/client.go
Normal file
@@ -0,0 +1,200 @@
|
|||||||
|
package mcps
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
"agent-desk/internal/pkg/i18nx"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Client struct{}
|
||||||
|
|
||||||
|
func NewClient() *Client {
|
||||||
|
return &Client{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) TestConnection(ctx context.Context, cfg ServerConfig) (*ConnectionResult, error) {
|
||||||
|
session, closeFn, err := c.connect(ctx, cfg)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer closeFn()
|
||||||
|
|
||||||
|
initResult := session.InitializeResult()
|
||||||
|
serverName := ""
|
||||||
|
version := ""
|
||||||
|
protocol := ""
|
||||||
|
if initResult != nil {
|
||||||
|
serverName = initResult.ServerInfo.Name
|
||||||
|
version = initResult.ServerInfo.Version
|
||||||
|
protocol = initResult.ProtocolVersion
|
||||||
|
}
|
||||||
|
return &ConnectionResult{
|
||||||
|
ServerCode: cfg.Code,
|
||||||
|
Endpoint: cfg.Endpoint,
|
||||||
|
Protocol: protocol,
|
||||||
|
ServerName: serverName,
|
||||||
|
Version: version,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) ListTools(ctx context.Context, cfg ServerConfig) ([]ToolInfo, error) {
|
||||||
|
session, closeFn, err := c.connect(ctx, cfg)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer closeFn()
|
||||||
|
|
||||||
|
result, err := session.ListTools(ctx, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, i18nx.Errorf("error.mcp.listToolsFailed", err)
|
||||||
|
}
|
||||||
|
ret := make([]ToolInfo, 0, len(result.Tools))
|
||||||
|
for _, tool := range result.Tools {
|
||||||
|
readOnlyHint := tool.Annotations != nil && tool.Annotations.ReadOnlyHint
|
||||||
|
ret = append(ret, ToolInfo{
|
||||||
|
Name: tool.Name,
|
||||||
|
Title: tool.Title,
|
||||||
|
Description: tool.Description,
|
||||||
|
InputSchema: tool.InputSchema,
|
||||||
|
OutputSchema: tool.OutputSchema,
|
||||||
|
ReadOnlyHint: readOnlyHint,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return ret, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) CallTool(ctx context.Context, cfg ServerConfig, toolName string, arguments map[string]any) (*ToolCallResult, error) {
|
||||||
|
toolName = strings.TrimSpace(toolName)
|
||||||
|
if toolName == "" {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0076")
|
||||||
|
}
|
||||||
|
|
||||||
|
session, closeFn, err := c.connect(ctx, cfg)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer closeFn()
|
||||||
|
|
||||||
|
result, err := session.CallTool(ctx, &mcp.CallToolParams{
|
||||||
|
Name: toolName,
|
||||||
|
Arguments: arguments,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, i18nx.Errorf("error.mcp.callToolFailed", err)
|
||||||
|
}
|
||||||
|
return &ToolCallResult{
|
||||||
|
ServerCode: cfg.Code,
|
||||||
|
ToolName: toolName,
|
||||||
|
IsError: result.IsError,
|
||||||
|
Content: convertContents(result.Content),
|
||||||
|
StructuredContent: result.StructuredContent,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) connect(ctx context.Context, cfg ServerConfig) (*mcp.ClientSession, func(), error) {
|
||||||
|
if strings.TrimSpace(cfg.Code) == "" {
|
||||||
|
return nil, nil, errorsx.InvalidParamI18n("error.e0070")
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(cfg.Endpoint) == "" {
|
||||||
|
return nil, nil, errorsx.InvalidParamI18n("error.e0032")
|
||||||
|
}
|
||||||
|
|
||||||
|
timeout := time.Duration(cfg.TimeoutMS) * time.Millisecond
|
||||||
|
if timeout <= 0 {
|
||||||
|
timeout = 15 * time.Second
|
||||||
|
}
|
||||||
|
connCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||||
|
|
||||||
|
httpClient := &http.Client{
|
||||||
|
Transport: &headerRoundTripper{
|
||||||
|
next: http.DefaultTransport,
|
||||||
|
headers: cfg.Headers,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
client := mcp.NewClient(&mcp.Implementation{
|
||||||
|
Name: "agent-desk-mcp-client",
|
||||||
|
Version: "v1",
|
||||||
|
}, nil)
|
||||||
|
transport := &mcp.StreamableClientTransport{
|
||||||
|
Endpoint: cfg.Endpoint,
|
||||||
|
HTTPClient: httpClient,
|
||||||
|
MaxRetries: 0,
|
||||||
|
DisableStandaloneSSE: true,
|
||||||
|
}
|
||||||
|
session, err := client.Connect(connCtx, transport, nil)
|
||||||
|
if err != nil {
|
||||||
|
cancel()
|
||||||
|
return nil, nil, i18nx.Errorf("error.mcp.connectServerFailed", err)
|
||||||
|
}
|
||||||
|
return session, func() {
|
||||||
|
_ = session.Close()
|
||||||
|
cancel()
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func convertContents(contents []mcp.Content) []ToolResultContent {
|
||||||
|
ret := make([]ToolResultContent, 0, len(contents))
|
||||||
|
for _, item := range contents {
|
||||||
|
switch v := item.(type) {
|
||||||
|
case *mcp.TextContent:
|
||||||
|
ret = append(ret, ToolResultContent{
|
||||||
|
Type: "text",
|
||||||
|
Text: v.Text,
|
||||||
|
})
|
||||||
|
case *mcp.ImageContent:
|
||||||
|
ret = append(ret, ToolResultContent{
|
||||||
|
Type: "image",
|
||||||
|
Data: map[string]any{
|
||||||
|
"mimeType": v.MIMEType,
|
||||||
|
"data": v.Data,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
case *mcp.AudioContent:
|
||||||
|
ret = append(ret, ToolResultContent{
|
||||||
|
Type: "audio",
|
||||||
|
Data: map[string]any{
|
||||||
|
"mimeType": v.MIMEType,
|
||||||
|
"data": v.Data,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
case *mcp.EmbeddedResource:
|
||||||
|
ret = append(ret, ToolResultContent{
|
||||||
|
Type: "resource",
|
||||||
|
Data: v.Resource,
|
||||||
|
})
|
||||||
|
default:
|
||||||
|
ret = append(ret, ToolResultContent{
|
||||||
|
Type: "unknown",
|
||||||
|
Data: v,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
|
||||||
|
type headerRoundTripper struct {
|
||||||
|
next http.RoundTripper
|
||||||
|
headers map[string]string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
next := r.next
|
||||||
|
if next == nil {
|
||||||
|
next = http.DefaultTransport
|
||||||
|
}
|
||||||
|
clone := req.Clone(req.Context())
|
||||||
|
for key, value := range r.headers {
|
||||||
|
key = strings.TrimSpace(key)
|
||||||
|
if key == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
clone.Header.Set(key, value)
|
||||||
|
}
|
||||||
|
return next.RoundTrip(clone)
|
||||||
|
}
|
||||||
77
internal/ai/mcps/providers/system_tools_provider.go
Normal file
77
internal/ai/mcps/providers/system_tools_provider.go
Normal file
@@ -0,0 +1,77 @@
|
|||||||
|
package providers
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
"context"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
type systemToolProvider struct{}
|
||||||
|
|
||||||
|
func NewSystemToolProvider() ToolProvider {
|
||||||
|
return &systemToolProvider{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *systemToolProvider) Name() string {
|
||||||
|
return "system"
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *systemToolProvider) Register(server *mcp.Server) error {
|
||||||
|
mcp.AddTool(
|
||||||
|
server,
|
||||||
|
&mcp.Tool{
|
||||||
|
Name: "server_time",
|
||||||
|
Title: "获取当前时间",
|
||||||
|
Description: "获取当前服务端时间,可选传入时区。",
|
||||||
|
Annotations: &mcp.ToolAnnotations{
|
||||||
|
ReadOnlyHint: true,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
func(_ context.Context, _ *mcp.CallToolRequest, args serverTimeArgs) (*mcp.CallToolResult, map[string]any, error) {
|
||||||
|
loc := time.Local
|
||||||
|
timezone := args.Timezone
|
||||||
|
if timezone == "" {
|
||||||
|
timezone = "Local"
|
||||||
|
} else if loaded, err := time.LoadLocation(timezone); err == nil {
|
||||||
|
loc = loaded
|
||||||
|
}
|
||||||
|
now := time.Now().In(loc)
|
||||||
|
return nil, map[string]any{
|
||||||
|
"timezone": timezone,
|
||||||
|
"timestamp": now.Format("2006-01-02 15:04:05"),
|
||||||
|
"unix": now.Unix(),
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
mcp.AddTool(
|
||||||
|
server,
|
||||||
|
&mcp.Tool{
|
||||||
|
Name: "service_info",
|
||||||
|
Title: "查看服务信息",
|
||||||
|
Description: "查看当前 agent-desk 服务的基础运行信息。",
|
||||||
|
Annotations: &mcp.ToolAnnotations{
|
||||||
|
ReadOnlyHint: true,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
func(ctx context.Context, req *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, map[string]any, error) {
|
||||||
|
cfg := config.Current()
|
||||||
|
return nil, map[string]any{
|
||||||
|
"name": "agent-desk",
|
||||||
|
"version": "v1",
|
||||||
|
"mcpPath": "/api/mcp",
|
||||||
|
"port": cfg.Server.Port,
|
||||||
|
"mcpEnabled": cfg.MCP.Enabled,
|
||||||
|
"vectorDb": cfg.VectorDB.Type,
|
||||||
|
"storageType": cfg.Storage.Default,
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type serverTimeArgs struct {
|
||||||
|
Timezone string `json:"timezone,omitempty" jsonschema:"可选时区名称,例如 Asia/Shanghai 或 UTC"`
|
||||||
|
}
|
||||||
10
internal/ai/mcps/providers/types.go
Normal file
10
internal/ai/mcps/providers/types.go
Normal file
@@ -0,0 +1,10 @@
|
|||||||
|
package providers
|
||||||
|
|
||||||
|
import (
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
type ToolProvider interface {
|
||||||
|
Name() string
|
||||||
|
Register(server *mcp.Server) error
|
||||||
|
}
|
||||||
23
internal/ai/mcps/registry.go
Normal file
23
internal/ai/mcps/registry.go
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
package mcps
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/ai/mcps/providers"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
func defaultProviders() []providers.ToolProvider {
|
||||||
|
return []providers.ToolProvider{
|
||||||
|
providers.NewSystemToolProvider(),
|
||||||
|
// 在这里注册其他的 ToolProvider
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func registerProviders(server *mcp.Server) error {
|
||||||
|
for _, provider := range defaultProviders() {
|
||||||
|
if err := provider.Register(server); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
72
internal/ai/mcps/runtime.go
Normal file
72
internal/ai/mcps/runtime.go
Normal file
@@ -0,0 +1,72 @@
|
|||||||
|
package mcps
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
)
|
||||||
|
|
||||||
|
type RuntimeService struct {
|
||||||
|
client *Client
|
||||||
|
}
|
||||||
|
|
||||||
|
var Runtime = NewRuntimeService()
|
||||||
|
|
||||||
|
func NewRuntimeService() *RuntimeService {
|
||||||
|
return &RuntimeService{
|
||||||
|
client: NewClient(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RuntimeService) CallTool(ctx context.Context, serverCode string, toolName string, arguments map[string]any) (*ToolCallResult, error) {
|
||||||
|
server, err := s.resolveServer(serverCode)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return s.client.CallTool(ctx, server, toolName, arguments)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RuntimeService) ListTools(ctx context.Context, serverCode string) ([]ToolInfo, error) {
|
||||||
|
server, err := s.resolveServer(serverCode)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return s.client.ListTools(ctx, server)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *RuntimeService) resolveServer(serverCode string) (ServerConfig, error) {
|
||||||
|
cfg := config.Current()
|
||||||
|
if !cfg.MCP.Enabled {
|
||||||
|
return ServerConfig{}, errorsx.InvalidParamI18n("error.e0035")
|
||||||
|
}
|
||||||
|
serverCode = strings.TrimSpace(serverCode)
|
||||||
|
if serverCode == "" {
|
||||||
|
return ServerConfig{}, errorsx.InvalidParamI18n("error.e0070")
|
||||||
|
}
|
||||||
|
server, ok := cfg.MCP.Servers[serverCode]
|
||||||
|
if !ok {
|
||||||
|
return ServerConfig{}, errorsx.InvalidParamI18n("error.e0034")
|
||||||
|
}
|
||||||
|
if !server.Enabled {
|
||||||
|
return ServerConfig{}, errorsx.InvalidParamI18n("error.e0033")
|
||||||
|
}
|
||||||
|
return ServerConfig{
|
||||||
|
Code: serverCode,
|
||||||
|
Endpoint: strings.TrimSpace(server.Endpoint),
|
||||||
|
TimeoutMS: server.TimeoutMS,
|
||||||
|
Headers: cloneRuntimeHeaders(server.Headers),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func cloneRuntimeHeaders(headers map[string]string) map[string]string {
|
||||||
|
if len(headers) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
ret := make(map[string]string, len(headers))
|
||||||
|
for key, value := range headers {
|
||||||
|
ret[key] = value
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
32
internal/ai/mcps/server.go
Normal file
32
internal/ai/mcps/server.go
Normal file
@@ -0,0 +1,32 @@
|
|||||||
|
package mcps
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
func NewHTTPHandler() http.Handler {
|
||||||
|
server := newServer()
|
||||||
|
return mcp.NewStreamableHTTPHandler(func(_ *http.Request) *mcp.Server {
|
||||||
|
return server
|
||||||
|
}, &mcp.StreamableHTTPOptions{
|
||||||
|
JSONResponse: true,
|
||||||
|
SessionTimeout: 2 * time.Minute,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func newServer() *mcp.Server {
|
||||||
|
server := mcp.NewServer(&mcp.Implementation{
|
||||||
|
Name: "agent-desk-mcp-server",
|
||||||
|
Title: "CS Agent MCP Server",
|
||||||
|
Version: "v1",
|
||||||
|
WebsiteURL: "https://github.com/modelcontextprotocol",
|
||||||
|
}, nil)
|
||||||
|
if err := registerProviders(server); err != nil {
|
||||||
|
panic(fmt.Sprintf("register mcp providers failed: %v", err))
|
||||||
|
}
|
||||||
|
return server
|
||||||
|
}
|
||||||
46
internal/ai/mcps/types.go
Normal file
46
internal/ai/mcps/types.go
Normal file
@@ -0,0 +1,46 @@
|
|||||||
|
package mcps
|
||||||
|
|
||||||
|
type ServerConfig struct {
|
||||||
|
Code string
|
||||||
|
Endpoint string
|
||||||
|
TimeoutMS int
|
||||||
|
Headers map[string]string
|
||||||
|
}
|
||||||
|
|
||||||
|
type ServerInfo struct {
|
||||||
|
Code string `json:"code"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
Endpoint string `json:"endpoint"`
|
||||||
|
TimeoutMS int `json:"timeoutMs"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type ConnectionResult struct {
|
||||||
|
ServerCode string `json:"serverCode"`
|
||||||
|
Endpoint string `json:"endpoint"`
|
||||||
|
Protocol string `json:"protocol"`
|
||||||
|
ServerName string `json:"serverName"`
|
||||||
|
Version string `json:"version"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type ToolInfo struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Title string `json:"title"`
|
||||||
|
Description string `json:"description"`
|
||||||
|
InputSchema any `json:"inputSchema"`
|
||||||
|
OutputSchema any `json:"outputSchema,omitempty"`
|
||||||
|
ReadOnlyHint bool `json:"readOnlyHint"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type ToolResultContent struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Text string `json:"text,omitempty"`
|
||||||
|
Data any `json:"data,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type ToolCallResult struct {
|
||||||
|
ServerCode string `json:"serverCode"`
|
||||||
|
ToolName string `json:"toolName"`
|
||||||
|
IsError bool `json:"isError"`
|
||||||
|
Content []ToolResultContent `json:"content"`
|
||||||
|
StructuredContent any `json:"structuredContent,omitempty"`
|
||||||
|
}
|
||||||
37
internal/ai/openai_client.go
Normal file
37
internal/ai/openai_client.go
Normal file
@@ -0,0 +1,37 @@
|
|||||||
|
package ai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
openai "github.com/openai/openai-go/v3"
|
||||||
|
"github.com/openai/openai-go/v3/option"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
)
|
||||||
|
|
||||||
|
func newOpenAIClient(config models.AIConfig) openai.Client {
|
||||||
|
opts := []option.RequestOption{
|
||||||
|
option.WithAPIKey(config.APIKey),
|
||||||
|
option.WithBaseURL(config.BaseURL),
|
||||||
|
}
|
||||||
|
if config.TimeoutMS > 0 {
|
||||||
|
opts = append(opts, option.WithRequestTimeout(time.Duration(config.TimeoutMS)*time.Millisecond))
|
||||||
|
}
|
||||||
|
if config.MaxRetryCount >= 0 {
|
||||||
|
opts = append(opts, option.WithMaxRetries(config.MaxRetryCount))
|
||||||
|
}
|
||||||
|
|
||||||
|
return openai.NewClient(opts...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func GetEnabledAIConfig(modelType enums.AIModelType) (*models.AIConfig, error) {
|
||||||
|
item := repositories.AIConfigRepository.GetEnabled(sqls.DB(), modelType)
|
||||||
|
if item == nil {
|
||||||
|
return nil, errorsx.BusinessErrorI18n(2005, "error.aiConfig.noneEnabled")
|
||||||
|
}
|
||||||
|
return item, nil
|
||||||
|
}
|
||||||
422
internal/ai/rag/answer.go
Normal file
422
internal/ai/rag/answer.go
Normal file
@@ -0,0 +1,422 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/dto"
|
||||||
|
"agent-desk/internal/pkg/dto/request"
|
||||||
|
"agent-desk/internal/pkg/dto/response"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
)
|
||||||
|
|
||||||
|
type answer struct {
|
||||||
|
}
|
||||||
|
|
||||||
|
var Answer = &answer{}
|
||||||
|
|
||||||
|
func (s *answer) DebugSearch(ctx context.Context, req request.KnowledgeSearchRequest) (*response.KnowledgeSearchResponse, error) {
|
||||||
|
if strings.TrimSpace(req.Question) == "" {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0340")
|
||||||
|
}
|
||||||
|
startedAt := time.Now()
|
||||||
|
results, err := s.retrieve(req, ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
respResults := make([]response.KnowledgeSearchResult, 0, len(results))
|
||||||
|
for _, item := range results {
|
||||||
|
respResults = append(respResults, response.KnowledgeSearchResult{
|
||||||
|
KnowledgeBaseID: item.KnowledgeBaseID,
|
||||||
|
ChunkID: item.ChunkID,
|
||||||
|
DocumentID: item.DocumentID,
|
||||||
|
DocumentTitle: item.DocumentTitle,
|
||||||
|
FaqID: item.FaqID,
|
||||||
|
FaqQuestion: item.FaqQuestion,
|
||||||
|
ChunkNo: item.ChunkNo,
|
||||||
|
Title: item.Title,
|
||||||
|
SectionPath: item.SectionPath,
|
||||||
|
Content: item.Content,
|
||||||
|
Score: float64(item.Score),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return &response.KnowledgeSearchResponse{
|
||||||
|
Question: req.Question,
|
||||||
|
Results: respResults,
|
||||||
|
HitCount: len(respResults),
|
||||||
|
LatencyMs: time.Since(startedAt).Milliseconds(),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *answer) DebugAnswer(ctx context.Context, req request.KnowledgeAnswerRequest, operator *dto.AuthPrincipal) (*response.KnowledgeAnswerResponse, error) {
|
||||||
|
if strings.TrimSpace(req.Question) == "" {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0340")
|
||||||
|
}
|
||||||
|
startedAt := time.Now()
|
||||||
|
|
||||||
|
retrieveStartedAt := time.Now()
|
||||||
|
results, err := s.retrieve(request.KnowledgeSearchRequest{
|
||||||
|
KnowledgeBaseIDs: req.KnowledgeBaseIDs,
|
||||||
|
Question: req.Question,
|
||||||
|
TopK: req.TopK,
|
||||||
|
ScoreThreshold: req.ScoreThreshold,
|
||||||
|
RerankLimit: req.RerankLimit,
|
||||||
|
}, ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
retrieveMs := time.Since(retrieveStartedAt).Milliseconds()
|
||||||
|
knowledgeBase := s.resolveAnswerKnowledgeBase(req.KnowledgeBaseIDs, results)
|
||||||
|
contextResults := buildContextHits(Retrieve.SelectContextResults(results, 4000))
|
||||||
|
|
||||||
|
hits := make([]response.KnowledgeSearchResult, 0, len(results))
|
||||||
|
topScore := 0.0
|
||||||
|
for i, item := range results {
|
||||||
|
score := float64(item.Score)
|
||||||
|
if i == 0 {
|
||||||
|
topScore = score
|
||||||
|
}
|
||||||
|
hits = append(hits, response.KnowledgeSearchResult{
|
||||||
|
KnowledgeBaseID: item.KnowledgeBaseID,
|
||||||
|
ChunkID: item.ChunkID,
|
||||||
|
DocumentID: item.DocumentID,
|
||||||
|
DocumentTitle: item.DocumentTitle,
|
||||||
|
FaqID: item.FaqID,
|
||||||
|
FaqQuestion: item.FaqQuestion,
|
||||||
|
ChunkNo: item.ChunkNo,
|
||||||
|
Title: item.Title,
|
||||||
|
SectionPath: item.SectionPath,
|
||||||
|
Content: item.Content,
|
||||||
|
Score: score,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
citations := buildKnowledgeCitations(hits, 3)
|
||||||
|
|
||||||
|
answerMode := enums.KnowledgeAnswerMode(req.AnswerMode)
|
||||||
|
if answerMode == 0 {
|
||||||
|
if knowledgeBase != nil {
|
||||||
|
answerMode = enums.KnowledgeAnswerMode(knowledgeBase.AnswerMode)
|
||||||
|
}
|
||||||
|
if answerMode == 0 {
|
||||||
|
answerMode = enums.KnowledgeAnswerModeStrict
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fallbackMode := enums.AIAgentFallbackModeNoAnswer
|
||||||
|
|
||||||
|
answerStatus := enums.KnowledgeAnswerStatusNormal
|
||||||
|
answer := ""
|
||||||
|
modelName := ""
|
||||||
|
promptTokens := 0
|
||||||
|
completionTokens := 0
|
||||||
|
generateStartedAt := time.Now()
|
||||||
|
|
||||||
|
if len(hits) == 0 {
|
||||||
|
answerStatus = enums.KnowledgeAnswerStatusNoAnswer
|
||||||
|
answer = buildFallbackAnswer(fallbackMode)
|
||||||
|
} else {
|
||||||
|
contextText := Retrieve.BuildContext(ctx, results, 4000)
|
||||||
|
systemPrompt := buildAnswerSystemPrompt(answerMode)
|
||||||
|
userPrompt := fmt.Sprintf("用户问题:%s\n\n参考资料:\n%s", req.Question, contextText)
|
||||||
|
llmResult, llmErr := ai.LLM.Chat(ctx, systemPrompt, userPrompt)
|
||||||
|
if llmErr != nil {
|
||||||
|
answerStatus = enums.KnowledgeAnswerStatusFallback
|
||||||
|
answer = buildFallbackAnswer(fallbackMode)
|
||||||
|
} else {
|
||||||
|
answer = llmResult.Content
|
||||||
|
modelName = llmResult.ModelName
|
||||||
|
promptTokens = llmResult.PromptTokens
|
||||||
|
completionTokens = llmResult.CompletionTokens
|
||||||
|
if strings.TrimSpace(answer) == "" {
|
||||||
|
answerStatus = enums.KnowledgeAnswerStatusFallback
|
||||||
|
answer = buildFallbackAnswer(fallbackMode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
generateMs := time.Since(generateStartedAt).Milliseconds()
|
||||||
|
rerankLimit := 0
|
||||||
|
chunkProvider := ""
|
||||||
|
chunkTargetTokens := 0
|
||||||
|
chunkMaxTokens := 0
|
||||||
|
chunkOverlapTokens := 0
|
||||||
|
if knowledgeBase != nil {
|
||||||
|
rerankLimit = resolveRerankLimit(req.RerankLimit, knowledgeBase.DefaultRerankLimit)
|
||||||
|
chunkProvider = knowledgeBase.ChunkProvider
|
||||||
|
chunkTargetTokens = knowledgeBase.ChunkTargetTokens
|
||||||
|
chunkMaxTokens = knowledgeBase.ChunkMaxTokens
|
||||||
|
chunkOverlapTokens = knowledgeBase.ChunkOverlapTokens
|
||||||
|
}
|
||||||
|
|
||||||
|
logItem, err := RetrieveLog.CreateRetrieveLog(&CreateRetrieveLogRequest{
|
||||||
|
KnowledgeBaseID: firstKnowledgeBaseID(req.KnowledgeBaseIDs),
|
||||||
|
Channel: defaultRetrieveChannel(req.Channel),
|
||||||
|
Scene: defaultRetrieveScene(req.Scene),
|
||||||
|
SessionID: req.SessionID,
|
||||||
|
ConversationID: req.ConversationID,
|
||||||
|
Question: req.Question,
|
||||||
|
RewriteQuestion: "",
|
||||||
|
Answer: answer,
|
||||||
|
AnswerStatus: int(answerStatus),
|
||||||
|
ChunkProvider: chunkProvider,
|
||||||
|
ChunkTargetTokens: chunkTargetTokens,
|
||||||
|
ChunkMaxTokens: chunkMaxTokens,
|
||||||
|
ChunkOverlapTokens: chunkOverlapTokens,
|
||||||
|
RerankEnabled: rerankLimit > 0,
|
||||||
|
RerankLimit: rerankLimit,
|
||||||
|
Hits: hits,
|
||||||
|
UsedHits: contextResults,
|
||||||
|
Citations: citations,
|
||||||
|
LatencyMs: time.Since(startedAt).Milliseconds(),
|
||||||
|
RetrieveMs: retrieveMs,
|
||||||
|
GenerateMs: generateMs,
|
||||||
|
PromptTokens: promptTokens,
|
||||||
|
CompletionTokens: completionTokens,
|
||||||
|
ModelName: modelName,
|
||||||
|
}, operator)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return &response.KnowledgeAnswerResponse{
|
||||||
|
Question: req.Question,
|
||||||
|
Answer: answer,
|
||||||
|
AnswerStatus: int(answerStatus),
|
||||||
|
AnswerStatusName: getAnswerStatusName(answerStatus),
|
||||||
|
Citations: citations,
|
||||||
|
Hits: hits,
|
||||||
|
HitCount: len(hits),
|
||||||
|
TopScore: topScore,
|
||||||
|
LatencyMs: time.Since(startedAt).Milliseconds(),
|
||||||
|
RetrieveMs: retrieveMs,
|
||||||
|
GenerateMs: generateMs,
|
||||||
|
PromptTokens: promptTokens,
|
||||||
|
CompletionTokens: completionTokens,
|
||||||
|
ModelName: modelName,
|
||||||
|
RetrieveLogID: logItem.ID,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildContextHits(results []RetrieveResult) []response.KnowledgeSearchResult {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
hits := make([]response.KnowledgeSearchResult, 0, len(results))
|
||||||
|
for _, item := range results {
|
||||||
|
hits = append(hits, response.KnowledgeSearchResult{
|
||||||
|
KnowledgeBaseID: item.KnowledgeBaseID,
|
||||||
|
ChunkID: item.ChunkID,
|
||||||
|
DocumentID: item.DocumentID,
|
||||||
|
DocumentTitle: item.DocumentTitle,
|
||||||
|
FaqID: item.FaqID,
|
||||||
|
FaqQuestion: item.FaqQuestion,
|
||||||
|
ChunkNo: item.ChunkNo,
|
||||||
|
Title: item.Title,
|
||||||
|
SectionPath: item.SectionPath,
|
||||||
|
Content: item.Content,
|
||||||
|
Score: float64(item.Score),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return hits
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildKnowledgeCitations(hits []response.KnowledgeSearchResult, limit int) []response.KnowledgeCitation {
|
||||||
|
if len(hits) == 0 || limit <= 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
citations := make([]response.KnowledgeCitation, 0, limit)
|
||||||
|
seen := make(map[string]struct{})
|
||||||
|
for _, item := range hits {
|
||||||
|
key := fmt.Sprintf("%d|%d|%s|%d", item.DocumentID, item.FaqID, item.SectionPath, item.ChunkNo)
|
||||||
|
if _, ok := seen[key]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[key] = struct{}{}
|
||||||
|
citations = append(citations, response.KnowledgeCitation{
|
||||||
|
DocumentID: item.DocumentID,
|
||||||
|
DocumentTitle: item.DocumentTitle,
|
||||||
|
FaqID: item.FaqID,
|
||||||
|
FaqQuestion: item.FaqQuestion,
|
||||||
|
ChunkNo: item.ChunkNo,
|
||||||
|
Title: item.Title,
|
||||||
|
SectionPath: item.SectionPath,
|
||||||
|
Snippet: truncateCitationSnippet(item.Content, 120),
|
||||||
|
Score: item.Score,
|
||||||
|
})
|
||||||
|
if len(citations) >= limit {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return citations
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateCitationSnippet(text string, limit int) string {
|
||||||
|
if limit <= 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
runes := []rune(strings.TrimSpace(text))
|
||||||
|
if len(runes) <= limit {
|
||||||
|
return string(runes)
|
||||||
|
}
|
||||||
|
return string(runes[:limit]) + "..."
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *answer) BuildDocumentIndex(ctx context.Context, documentID int64) error {
|
||||||
|
return Index.IndexDocumentByID(ctx, documentID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *answer) retrieve(req request.KnowledgeSearchRequest, ctx context.Context) ([]RetrieveResult, error) {
|
||||||
|
if len(normalizeKnowledgeBaseIDs(req.KnowledgeBaseIDs)) == 0 {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0284")
|
||||||
|
}
|
||||||
|
knowledgeBases := s.loadKnowledgeBases(req.KnowledgeBaseIDs)
|
||||||
|
|
||||||
|
results, err := Retrieve.Retrieve(ctx, RetrieveRequest{
|
||||||
|
KnowledgeBaseIDs: req.KnowledgeBaseIDs,
|
||||||
|
Query: req.Question,
|
||||||
|
TopK: req.TopK,
|
||||||
|
ScoreThreshold: req.ScoreThreshold,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
defaultRerankLimit := resolveDefaultRerankLimit(knowledgeBases)
|
||||||
|
rerankLimit := resolveRerankLimit(req.RerankLimit, defaultRerankLimit)
|
||||||
|
if rerankLimit > 0 && len(results) > rerankLimit {
|
||||||
|
return Retrieve.RetrieveWithRerank(ctx, RetrieveRequest{
|
||||||
|
KnowledgeBaseIDs: req.KnowledgeBaseIDs,
|
||||||
|
Query: req.Question,
|
||||||
|
TopK: req.TopK,
|
||||||
|
ScoreThreshold: req.ScoreThreshold,
|
||||||
|
}, rerankLimit)
|
||||||
|
}
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *answer) loadKnowledgeBases(knowledgeBaseIDs []int64) []models.KnowledgeBase {
|
||||||
|
normalized := normalizeKnowledgeBaseIDs(knowledgeBaseIDs)
|
||||||
|
if len(normalized) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
items := repositories.KnowledgeBaseRepository.Find(sqls.DB(), sqls.NewCnd().In("id", normalized))
|
||||||
|
if len(items) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
itemMap := make(map[int64]models.KnowledgeBase, len(items))
|
||||||
|
for _, item := range items {
|
||||||
|
itemMap[item.ID] = item
|
||||||
|
}
|
||||||
|
results := make([]models.KnowledgeBase, 0, len(normalized))
|
||||||
|
for _, id := range normalized {
|
||||||
|
if item, ok := itemMap[id]; ok {
|
||||||
|
results = append(results, item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return results
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *answer) resolvePrimaryKnowledgeBase(knowledgeBaseIDs []int64) *models.KnowledgeBase {
|
||||||
|
items := s.loadKnowledgeBases(knowledgeBaseIDs)
|
||||||
|
for _, item := range items {
|
||||||
|
return &item
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *answer) resolveAnswerKnowledgeBase(knowledgeBaseIDs []int64, results []RetrieveResult) *models.KnowledgeBase {
|
||||||
|
items := s.loadKnowledgeBases(knowledgeBaseIDs)
|
||||||
|
if len(items) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if len(results) > 0 {
|
||||||
|
for _, item := range items {
|
||||||
|
if item.ID == results[0].KnowledgeBaseID {
|
||||||
|
return &item
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return &items[0]
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstKnowledgeBaseID(ids []int64) int64 {
|
||||||
|
normalized := normalizeKnowledgeBaseIDs(ids)
|
||||||
|
if len(normalized) == 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return normalized[0]
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveRerankLimit(requestLimit int, defaultLimit int) int {
|
||||||
|
if requestLimit > 0 {
|
||||||
|
return requestLimit
|
||||||
|
}
|
||||||
|
if defaultLimit > 0 {
|
||||||
|
return defaultLimit
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveDefaultRerankLimit(items []models.KnowledgeBase) int {
|
||||||
|
limit := 0
|
||||||
|
for _, item := range items {
|
||||||
|
if item.DefaultRerankLimit > limit {
|
||||||
|
limit = item.DefaultRerankLimit
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return limit
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildAnswerSystemPrompt(answerMode enums.KnowledgeAnswerMode) string {
|
||||||
|
if answerMode == enums.KnowledgeAnswerModeAssist {
|
||||||
|
return "你是客服知识库助手。请优先依据提供的知识片段回答,可以做轻度归纳,但不要编造未提供的事实。"
|
||||||
|
}
|
||||||
|
return "你是严格的客服知识库助手。只能依据提供的知识片段回答;如果资料不足,请明确说明知识库暂无明确信息。"
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildFallbackAnswer(fallbackMode enums.AIAgentFallbackMode) string {
|
||||||
|
switch fallbackMode {
|
||||||
|
case enums.AIAgentFallbackModeSuggestRetry:
|
||||||
|
return "当前知识库里没有找到足够明确的信息,你可以换个更具体的问法再试一次。"
|
||||||
|
default:
|
||||||
|
return "当前知识库暂无明确信息。"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func defaultRetrieveChannel(channel string) string {
|
||||||
|
if strings.TrimSpace(channel) == "" {
|
||||||
|
return string(enums.KnowledgeRetrieveChannelDebug)
|
||||||
|
}
|
||||||
|
return channel
|
||||||
|
}
|
||||||
|
|
||||||
|
func defaultRetrieveScene(scene string) string {
|
||||||
|
if strings.TrimSpace(scene) == "" {
|
||||||
|
return string(enums.KnowledgeRetrieveSceneQA)
|
||||||
|
}
|
||||||
|
return scene
|
||||||
|
}
|
||||||
|
|
||||||
|
func getAnswerStatusName(status enums.KnowledgeAnswerStatus) string {
|
||||||
|
switch status {
|
||||||
|
case enums.KnowledgeAnswerStatusNormal:
|
||||||
|
return "正常"
|
||||||
|
case enums.KnowledgeAnswerStatusNoAnswer:
|
||||||
|
return "无答案"
|
||||||
|
case enums.KnowledgeAnswerStatusFallback:
|
||||||
|
return "兜底"
|
||||||
|
case enums.KnowledgeAnswerStatusBlocked:
|
||||||
|
return "风控拦截"
|
||||||
|
default:
|
||||||
|
return "未知"
|
||||||
|
}
|
||||||
|
}
|
||||||
123
internal/ai/rag/answer_test.go
Normal file
123
internal/ai/rag/answer_test.go
Normal file
@@ -0,0 +1,123 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/dto/response"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuildFallbackAnswer(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
mode enums.AIAgentFallbackMode
|
||||||
|
expected string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "no answer",
|
||||||
|
mode: enums.AIAgentFallbackModeNoAnswer,
|
||||||
|
expected: "当前知识库暂无明确信息。",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "suggest retry",
|
||||||
|
mode: enums.AIAgentFallbackModeSuggestRetry,
|
||||||
|
expected: "当前知识库里没有找到足够明确的信息,你可以换个更具体的问法再试一次。",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
if got := buildFallbackAnswer(tt.mode); got != tt.expected {
|
||||||
|
t.Fatalf("%s: expected %q, got %q", tt.name, tt.expected, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetAnswerStatusName(t *testing.T) {
|
||||||
|
if got := getAnswerStatusName(enums.KnowledgeAnswerStatusNoAnswer); got != "无答案" {
|
||||||
|
t.Fatalf("expected no-answer label, got %q", got)
|
||||||
|
}
|
||||||
|
if got := getAnswerStatusName(enums.KnowledgeAnswerStatusFallback); got != "兜底" {
|
||||||
|
t.Fatalf("expected fallback label, got %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveRerankLimit(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
requestLimit int
|
||||||
|
defaultLimit int
|
||||||
|
expected int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "request overrides default",
|
||||||
|
requestLimit: 3,
|
||||||
|
defaultLimit: 5,
|
||||||
|
expected: 3,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "default used when request missing",
|
||||||
|
requestLimit: 0,
|
||||||
|
defaultLimit: 5,
|
||||||
|
expected: 5,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "zero when both missing",
|
||||||
|
requestLimit: 0,
|
||||||
|
defaultLimit: 0,
|
||||||
|
expected: 0,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
if got := resolveRerankLimit(tt.requestLimit, tt.defaultLimit); got != tt.expected {
|
||||||
|
t.Fatalf("%s: expected %d, got %d", tt.name, tt.expected, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveDefaultRerankLimit(t *testing.T) {
|
||||||
|
items := []models.KnowledgeBase{
|
||||||
|
{ID: 11, DefaultRerankLimit: 3},
|
||||||
|
{ID: 22, DefaultRerankLimit: 7},
|
||||||
|
{ID: 33, DefaultRerankLimit: 5},
|
||||||
|
}
|
||||||
|
|
||||||
|
if got := resolveDefaultRerankLimit(items); got != 7 {
|
||||||
|
t.Fatalf("expected max rerank limit 7, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildKnowledgeCitations(t *testing.T) {
|
||||||
|
hits := []response.KnowledgeSearchResult{
|
||||||
|
{
|
||||||
|
DocumentID: 11,
|
||||||
|
DocumentTitle: "退款手册",
|
||||||
|
ChunkNo: 0,
|
||||||
|
Title: "退款说明",
|
||||||
|
SectionPath: "售后 > 退款说明",
|
||||||
|
Content: "退款申请提交后,预计1-3个工作日到账。",
|
||||||
|
Score: 0.91,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
DocumentID: 11,
|
||||||
|
DocumentTitle: "退款手册",
|
||||||
|
ChunkNo: 0,
|
||||||
|
Title: "退款说明",
|
||||||
|
SectionPath: "售后 > 退款说明",
|
||||||
|
Content: "重复内容",
|
||||||
|
Score: 0.89,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
citations := buildKnowledgeCitations(hits, 3)
|
||||||
|
if len(citations) != 1 {
|
||||||
|
t.Fatalf("expected 1 citation, got %d", len(citations))
|
||||||
|
}
|
||||||
|
if citations[0].DocumentID != 11 {
|
||||||
|
t.Fatalf("expected document id 11, got %d", citations[0].DocumentID)
|
||||||
|
}
|
||||||
|
if citations[0].SectionPath != "售后 > 退款说明" {
|
||||||
|
t.Fatalf("unexpected section path: %q", citations[0].SectionPath)
|
||||||
|
}
|
||||||
|
}
|
||||||
44
internal/ai/rag/chunk/fixed_provider.go
Normal file
44
internal/ai/rag/chunk/fixed_provider.go
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
package chunk
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"context"
|
||||||
|
)
|
||||||
|
|
||||||
|
type fixedProvider struct{}
|
||||||
|
|
||||||
|
func NewFixedProvider() Provider {
|
||||||
|
return &fixedProvider{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *fixedProvider) Name() string {
|
||||||
|
return string(enums.KnowledgeChunkProviderFixed)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *fixedProvider) Supports(contentType enums.KnowledgeDocumentContentType) bool {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *fixedProvider) Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error) {
|
||||||
|
text := req.PlainText
|
||||||
|
if text == "" {
|
||||||
|
text = req.Content
|
||||||
|
}
|
||||||
|
parts := splitPlainText(text, req.Options)
|
||||||
|
results := make([]ChunkResult, 0, len(parts))
|
||||||
|
for i, part := range parts {
|
||||||
|
results = append(results, ChunkResult{
|
||||||
|
ChunkNo: i,
|
||||||
|
Title: req.DocumentTitle,
|
||||||
|
Content: part,
|
||||||
|
ChunkType: enums.KnowledgeChunkTypeText,
|
||||||
|
SectionPath: req.DocumentTitle,
|
||||||
|
CharCount: len([]rune(part)),
|
||||||
|
TokenCount: estimateTokenCount(part),
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"provider": enums.KnowledgeChunkProviderFixed,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
12
internal/ai/rag/chunk/provider.go
Normal file
12
internal/ai/rag/chunk/provider.go
Normal file
@@ -0,0 +1,12 @@
|
|||||||
|
package chunk
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"context"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Provider interface {
|
||||||
|
Name() string
|
||||||
|
Supports(contentType enums.KnowledgeDocumentContentType) bool
|
||||||
|
Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error)
|
||||||
|
}
|
||||||
59
internal/ai/rag/chunk/registry.go
Normal file
59
internal/ai/rag/chunk/registry.go
Normal file
@@ -0,0 +1,59 @@
|
|||||||
|
package chunk
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Registry struct {
|
||||||
|
providers map[string]Provider
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewRegistry() *Registry {
|
||||||
|
return &Registry{
|
||||||
|
providers: make(map[string]Provider),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewDefaultRegistry() *Registry {
|
||||||
|
r := NewRegistry()
|
||||||
|
r.Register(NewFixedProvider())
|
||||||
|
r.Register(NewStructuredProvider())
|
||||||
|
return r
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Registry) Register(p Provider) {
|
||||||
|
if p == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
r.providers[p.Name()] = p
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Registry) Get(name string) Provider {
|
||||||
|
if name == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return r.providers[name]
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Registry) Resolve(name string, contentType enums.KnowledgeDocumentContentType) Provider {
|
||||||
|
if p := r.Get(name); p != nil && p.Supports(contentType) {
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
if p := r.Get(string(enums.KnowledgeChunkProviderStructured)); p != nil && p.Supports(contentType) {
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
return r.Get(string(enums.KnowledgeChunkProviderFixed))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Registry) Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error) {
|
||||||
|
if req == nil {
|
||||||
|
return nil, fmt.Errorf("chunk request is nil")
|
||||||
|
}
|
||||||
|
provider := r.Resolve(req.Options.Provider, req.ContentType)
|
||||||
|
if provider == nil {
|
||||||
|
return nil, fmt.Errorf("chunk provider not found")
|
||||||
|
}
|
||||||
|
return provider.Chunk(ctx, req)
|
||||||
|
}
|
||||||
261
internal/ai/rag/chunk/structured_provider.go
Normal file
261
internal/ai/rag/chunk/structured_provider.go
Normal file
@@ -0,0 +1,261 @@
|
|||||||
|
package chunk
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/gomarkdown/markdown"
|
||||||
|
"golang.org/x/net/html"
|
||||||
|
)
|
||||||
|
|
||||||
|
type structuredProvider struct{}
|
||||||
|
|
||||||
|
type contentBlock struct {
|
||||||
|
Type string
|
||||||
|
Level int
|
||||||
|
Text string
|
||||||
|
Title string
|
||||||
|
SectionPath string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewStructuredProvider() Provider {
|
||||||
|
return &structuredProvider{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *structuredProvider) Name() string {
|
||||||
|
return string(enums.KnowledgeChunkProviderStructured)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *structuredProvider) Supports(contentType enums.KnowledgeDocumentContentType) bool {
|
||||||
|
switch contentType {
|
||||||
|
case enums.KnowledgeDocumentContentTypeHTML, enums.KnowledgeDocumentContentTypeMarkdown:
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *structuredProvider) Chunk(ctx context.Context, req *ChunkRequest) ([]ChunkResult, error) {
|
||||||
|
content := req.Content
|
||||||
|
if req.ContentType == enums.KnowledgeDocumentContentTypeMarkdown {
|
||||||
|
content = string(markdown.ToHTML([]byte(content), nil, nil))
|
||||||
|
}
|
||||||
|
|
||||||
|
blocks := parseStructuredBlocks(content, req.DocumentTitle)
|
||||||
|
if len(blocks) == 0 {
|
||||||
|
return NewFixedProvider().Chunk(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
results := make([]ChunkResult, 0)
|
||||||
|
chunkNo := 0
|
||||||
|
for _, block := range blocks {
|
||||||
|
parts := splitPlainText(block.Text, req.Options)
|
||||||
|
for _, part := range parts {
|
||||||
|
if part == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
results = append(results, ChunkResult{
|
||||||
|
ChunkNo: chunkNo,
|
||||||
|
Title: block.Title,
|
||||||
|
Content: part,
|
||||||
|
ChunkType: mapBlockType(block.Type),
|
||||||
|
SectionPath: block.SectionPath,
|
||||||
|
CharCount: len([]rune(part)),
|
||||||
|
TokenCount: estimateTokenCount(part),
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"provider": enums.KnowledgeChunkProviderStructured,
|
||||||
|
"blockType": block.Type,
|
||||||
|
"sectionPath": block.SectionPath,
|
||||||
|
"sectionTitle": block.Title,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
chunkNo++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(results) == 0 {
|
||||||
|
return NewFixedProvider().Chunk(ctx, req)
|
||||||
|
}
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseStructuredBlocks(content string, documentTitle string) []contentBlock {
|
||||||
|
content = strings.TrimSpace(content)
|
||||||
|
if content == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
parent := &html.Node{Type: html.ElementNode, Data: "div"}
|
||||||
|
nodes, err := html.ParseFragment(strings.NewReader(content), parent)
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var blocks []contentBlock
|
||||||
|
headings := make([]string, 0)
|
||||||
|
var walk func(node *html.Node)
|
||||||
|
walk = func(node *html.Node) {
|
||||||
|
if node == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if node.Type == html.ElementNode {
|
||||||
|
switch node.Data {
|
||||||
|
case "h1", "h2", "h3", "h4", "h5", "h6":
|
||||||
|
title := normalizeText(nodeText(node))
|
||||||
|
if title != "" {
|
||||||
|
level := int(node.Data[1] - '0')
|
||||||
|
if level <= 0 {
|
||||||
|
level = 1
|
||||||
|
}
|
||||||
|
headings = updateHeadingPath(headings, level, title)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
case "p":
|
||||||
|
appendBlock(&blocks, "paragraph", normalizeText(nodeText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||||
|
return
|
||||||
|
case "ul", "ol":
|
||||||
|
appendBlock(&blocks, "list", normalizeText(listText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||||
|
return
|
||||||
|
case "table":
|
||||||
|
appendBlock(&blocks, "table", normalizeText(tableText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||||
|
return
|
||||||
|
case "pre", "code":
|
||||||
|
appendBlock(&blocks, "code", normalizeText(nodeText(node)), currentTitle(headings, documentTitle), strings.Join(headings, " > "))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
walk(child)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, node := range nodes {
|
||||||
|
walk(node)
|
||||||
|
}
|
||||||
|
return blocks
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendBlock(blocks *[]contentBlock, blockType string, text string, title string, sectionPath string) {
|
||||||
|
text = normalizeText(text)
|
||||||
|
if text == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if sectionPath == "" {
|
||||||
|
sectionPath = title
|
||||||
|
}
|
||||||
|
*blocks = append(*blocks, contentBlock{
|
||||||
|
Type: blockType,
|
||||||
|
Text: text,
|
||||||
|
Title: title,
|
||||||
|
SectionPath: sectionPath,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func updateHeadingPath(headings []string, level int, title string) []string {
|
||||||
|
if level <= 0 {
|
||||||
|
level = 1
|
||||||
|
}
|
||||||
|
if len(headings) >= level {
|
||||||
|
headings = headings[:level-1]
|
||||||
|
}
|
||||||
|
headings = append(headings, title)
|
||||||
|
return headings
|
||||||
|
}
|
||||||
|
|
||||||
|
func currentTitle(headings []string, documentTitle string) string {
|
||||||
|
if len(headings) == 0 {
|
||||||
|
return documentTitle
|
||||||
|
}
|
||||||
|
return headings[len(headings)-1]
|
||||||
|
}
|
||||||
|
|
||||||
|
func mapBlockType(blockType string) enums.KnowledgeChunkType {
|
||||||
|
switch blockType {
|
||||||
|
case "table":
|
||||||
|
return enums.KnowledgeChunkTypeTable
|
||||||
|
case "code":
|
||||||
|
return enums.KnowledgeChunkTypeCode
|
||||||
|
default:
|
||||||
|
return enums.KnowledgeChunkTypeText
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func nodeText(node *html.Node) string {
|
||||||
|
if node == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
var builder strings.Builder
|
||||||
|
writeNodeText(&builder, node)
|
||||||
|
return builder.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeNodeText(builder *strings.Builder, node *html.Node) {
|
||||||
|
if node == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
switch node.Type {
|
||||||
|
case html.TextNode:
|
||||||
|
builder.WriteString(node.Data)
|
||||||
|
case html.ElementNode:
|
||||||
|
if shouldSeparate(node.Data) {
|
||||||
|
builder.WriteByte(' ')
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
writeNodeText(builder, child)
|
||||||
|
}
|
||||||
|
if node.Type == html.ElementNode && shouldSeparate(node.Data) {
|
||||||
|
builder.WriteByte(' ')
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func shouldSeparate(tag string) bool {
|
||||||
|
switch tag {
|
||||||
|
case "p", "div", "br", "li", "ul", "ol", "blockquote", "pre", "table", "tr", "td", "th", "h1", "h2", "h3", "h4", "h5", "h6":
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func listText(node *html.Node) string {
|
||||||
|
items := make([]string, 0)
|
||||||
|
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
if child.Type == html.ElementNode && child.Data == "li" {
|
||||||
|
item := normalizeText(nodeText(child))
|
||||||
|
if item != "" {
|
||||||
|
items = append(items, item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return strings.Join(items, " ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func tableText(node *html.Node) string {
|
||||||
|
rows := make([]string, 0)
|
||||||
|
var walk func(*html.Node)
|
||||||
|
walk = func(n *html.Node) {
|
||||||
|
if n == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if n.Type == html.ElementNode && n.Data == "tr" {
|
||||||
|
cells := make([]string, 0)
|
||||||
|
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
if child.Type == html.ElementNode && (child.Data == "td" || child.Data == "th") {
|
||||||
|
cell := normalizeText(nodeText(child))
|
||||||
|
if cell != "" {
|
||||||
|
cells = append(cells, cell)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(cells) > 0 {
|
||||||
|
rows = append(rows, strings.Join(cells, " | "))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
for child := n.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
walk(child)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
walk(node)
|
||||||
|
return strings.Join(rows, " ")
|
||||||
|
}
|
||||||
32
internal/ai/rag/chunk/types.go
Normal file
32
internal/ai/rag/chunk/types.go
Normal file
@@ -0,0 +1,32 @@
|
|||||||
|
package chunk
|
||||||
|
|
||||||
|
import "agent-desk/internal/pkg/enums"
|
||||||
|
|
||||||
|
type ChunkRequest struct {
|
||||||
|
KnowledgeBaseID int64
|
||||||
|
DocumentID int64
|
||||||
|
DocumentTitle string
|
||||||
|
ContentType enums.KnowledgeDocumentContentType
|
||||||
|
Content string
|
||||||
|
PlainText string
|
||||||
|
Options ChunkOptions
|
||||||
|
}
|
||||||
|
|
||||||
|
type ChunkOptions struct {
|
||||||
|
Provider string
|
||||||
|
TargetTokens int
|
||||||
|
MaxTokens int
|
||||||
|
OverlapTokens int
|
||||||
|
EnableFallback bool
|
||||||
|
}
|
||||||
|
|
||||||
|
type ChunkResult struct {
|
||||||
|
ChunkNo int
|
||||||
|
Title string
|
||||||
|
Content string
|
||||||
|
ChunkType enums.KnowledgeChunkType
|
||||||
|
SectionPath string
|
||||||
|
CharCount int
|
||||||
|
TokenCount int
|
||||||
|
Metadata map[string]any
|
||||||
|
}
|
||||||
218
internal/ai/rag/chunk/utils.go
Normal file
218
internal/ai/rag/chunk/utils.go
Normal file
@@ -0,0 +1,218 @@
|
|||||||
|
package chunk
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"strings"
|
||||||
|
"unicode"
|
||||||
|
"unicode/utf8"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultTargetTokens = 300
|
||||||
|
defaultMaxTokens = 400
|
||||||
|
defaultOverlapTokens = 40
|
||||||
|
)
|
||||||
|
|
||||||
|
func normalizeOptions(opts ChunkOptions) ChunkOptions {
|
||||||
|
if opts.TargetTokens <= 0 {
|
||||||
|
opts.TargetTokens = defaultTargetTokens
|
||||||
|
}
|
||||||
|
if opts.MaxTokens <= 0 {
|
||||||
|
opts.MaxTokens = defaultMaxTokens
|
||||||
|
}
|
||||||
|
if opts.MaxTokens < opts.TargetTokens {
|
||||||
|
opts.MaxTokens = opts.TargetTokens
|
||||||
|
}
|
||||||
|
if opts.OverlapTokens < 0 {
|
||||||
|
opts.OverlapTokens = 0
|
||||||
|
}
|
||||||
|
if opts.OverlapTokens == 0 {
|
||||||
|
opts.OverlapTokens = defaultOverlapTokens
|
||||||
|
}
|
||||||
|
if opts.Provider == "" {
|
||||||
|
opts.Provider = string(enums.KnowledgeChunkProviderStructured)
|
||||||
|
}
|
||||||
|
return opts
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeText(text string) string {
|
||||||
|
return strings.Join(strings.Fields(strings.TrimSpace(text)), " ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func estimateTokenCount(text string) int {
|
||||||
|
text = strings.TrimSpace(text)
|
||||||
|
if text == "" {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
count := 0
|
||||||
|
inWord := false
|
||||||
|
for _, r := range text {
|
||||||
|
switch {
|
||||||
|
case unicode.IsSpace(r):
|
||||||
|
inWord = false
|
||||||
|
case unicode.Is(unicode.Han, r):
|
||||||
|
count++
|
||||||
|
inWord = false
|
||||||
|
case unicode.IsLetter(r) || unicode.IsDigit(r):
|
||||||
|
if !inWord {
|
||||||
|
count++
|
||||||
|
inWord = true
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
count++
|
||||||
|
inWord = false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if count == 0 {
|
||||||
|
return utf8.RuneCountInString(text)
|
||||||
|
}
|
||||||
|
return count
|
||||||
|
}
|
||||||
|
|
||||||
|
func contentHash(text string) string {
|
||||||
|
sum := sha256.Sum256([]byte(text))
|
||||||
|
return hex.EncodeToString(sum[:])
|
||||||
|
}
|
||||||
|
|
||||||
|
func splitSentences(text string) []string {
|
||||||
|
text = strings.TrimSpace(text)
|
||||||
|
if text == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var sentences []string
|
||||||
|
var builder strings.Builder
|
||||||
|
for _, r := range text {
|
||||||
|
builder.WriteRune(r)
|
||||||
|
switch r {
|
||||||
|
case '\n', '。', '!', '?', '!', '?', ';', ';':
|
||||||
|
sentence := normalizeText(builder.String())
|
||||||
|
if sentence != "" {
|
||||||
|
sentences = append(sentences, sentence)
|
||||||
|
}
|
||||||
|
builder.Reset()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if builder.Len() > 0 {
|
||||||
|
sentence := normalizeText(builder.String())
|
||||||
|
if sentence != "" {
|
||||||
|
sentences = append(sentences, sentence)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(sentences) == 0 {
|
||||||
|
return []string{normalizeText(text)}
|
||||||
|
}
|
||||||
|
return sentences
|
||||||
|
}
|
||||||
|
|
||||||
|
func tailTextByTokens(text string, tokenLimit int) string {
|
||||||
|
if tokenLimit <= 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
sentences := splitSentences(text)
|
||||||
|
if len(sentences) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
var selected []string
|
||||||
|
total := 0
|
||||||
|
for i := len(sentences) - 1; i >= 0; i-- {
|
||||||
|
sentence := sentences[i]
|
||||||
|
tokens := estimateTokenCount(sentence)
|
||||||
|
if total > 0 && total+tokens > tokenLimit {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
selected = append([]string{sentence}, selected...)
|
||||||
|
total += tokens
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(strings.Join(selected, " "))
|
||||||
|
}
|
||||||
|
|
||||||
|
func splitPlainText(text string, opts ChunkOptions) []string {
|
||||||
|
text = normalizeText(text)
|
||||||
|
if text == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
opts = normalizeOptions(opts)
|
||||||
|
sentences := splitSentences(text)
|
||||||
|
if len(sentences) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
chunks := make([]string, 0)
|
||||||
|
current := make([]string, 0)
|
||||||
|
currentTokens := 0
|
||||||
|
|
||||||
|
flush := func() {
|
||||||
|
if len(current) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
chunks = append(chunks, strings.Join(current, " "))
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, sentence := range sentences {
|
||||||
|
sentenceTokens := estimateTokenCount(sentence)
|
||||||
|
if sentenceTokens > opts.MaxTokens {
|
||||||
|
if len(current) > 0 {
|
||||||
|
flush()
|
||||||
|
overlap := tailTextByTokens(strings.Join(current, " "), opts.OverlapTokens)
|
||||||
|
current = nil
|
||||||
|
currentTokens = 0
|
||||||
|
if overlap != "" {
|
||||||
|
current = append(current, overlap)
|
||||||
|
currentTokens = estimateTokenCount(overlap)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, piece := range splitLongSentence(sentence, opts.MaxTokens) {
|
||||||
|
piece = normalizeText(piece)
|
||||||
|
if piece != "" {
|
||||||
|
chunks = append(chunks, piece)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if currentTokens > 0 && currentTokens+sentenceTokens > opts.MaxTokens {
|
||||||
|
flush()
|
||||||
|
overlap := tailTextByTokens(strings.Join(current, " "), opts.OverlapTokens)
|
||||||
|
current = nil
|
||||||
|
currentTokens = 0
|
||||||
|
if overlap != "" {
|
||||||
|
current = append(current, overlap)
|
||||||
|
currentTokens = estimateTokenCount(overlap)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
current = append(current, sentence)
|
||||||
|
currentTokens += sentenceTokens
|
||||||
|
}
|
||||||
|
|
||||||
|
flush()
|
||||||
|
return chunks
|
||||||
|
}
|
||||||
|
|
||||||
|
func splitLongSentence(text string, maxTokens int) []string {
|
||||||
|
runes := []rune(strings.TrimSpace(text))
|
||||||
|
if len(runes) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if maxTokens <= 0 {
|
||||||
|
return []string{text}
|
||||||
|
}
|
||||||
|
window := maxTokens * 2
|
||||||
|
if window < 50 {
|
||||||
|
window = 50
|
||||||
|
}
|
||||||
|
var result []string
|
||||||
|
for start := 0; start < len(runes); start += window {
|
||||||
|
end := start + window
|
||||||
|
if end > len(runes) {
|
||||||
|
end = len(runes)
|
||||||
|
}
|
||||||
|
part := normalizeText(string(runes[start:end]))
|
||||||
|
if part != "" {
|
||||||
|
result = append(result, part)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
331
internal/ai/rag/index.go
Normal file
331
internal/ai/rag/index.go
Normal file
@@ -0,0 +1,331 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai"
|
||||||
|
ragchunk "agent-desk/internal/ai/rag/chunk"
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
type ChunkingConfig struct {
|
||||||
|
Provider string
|
||||||
|
TargetTokens int
|
||||||
|
MaxTokens int
|
||||||
|
OverlapTokens int
|
||||||
|
EnableFallback bool
|
||||||
|
}
|
||||||
|
|
||||||
|
type index struct {
|
||||||
|
chunkConfig ChunkingConfig
|
||||||
|
registry *ragchunk.Registry
|
||||||
|
}
|
||||||
|
|
||||||
|
const knowledgeCollectionName = "knowledge_chunks"
|
||||||
|
|
||||||
|
var Index = &index{
|
||||||
|
chunkConfig: ChunkingConfig{
|
||||||
|
Provider: string(enums.KnowledgeChunkProviderStructured),
|
||||||
|
TargetTokens: 300,
|
||||||
|
MaxTokens: 400,
|
||||||
|
OverlapTokens: 40,
|
||||||
|
EnableFallback: true,
|
||||||
|
},
|
||||||
|
registry: ragchunk.NewDefaultRegistry(),
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) IndexDocumentByID(ctx context.Context, documentID int64) error {
|
||||||
|
document, err := s.loadDocumentByID(documentID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return s.IndexDocument(ctx, *document)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) IndexDocument(ctx context.Context, document models.KnowledgeDocument) error {
|
||||||
|
start := time.Now()
|
||||||
|
if err := s.markDocumentIndexPending(document.ID); err != nil {
|
||||||
|
slog.Error("Failed to mark knowledge document index as pending", "document_id", document.ID, "error", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
fail := func(err error) error {
|
||||||
|
if updateErr := s.markDocumentIndexFailed(document.ID, err); updateErr != nil {
|
||||||
|
slog.Error("Failed to mark knowledge document index as failed", "document_id", document.ID, "error", updateErr)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// TODO 这里每次都查询下知识库不太友好
|
||||||
|
knowledgeBase, err := s.loadDocumentKnowledgeBase(document)
|
||||||
|
if err != nil {
|
||||||
|
return fail(err)
|
||||||
|
}
|
||||||
|
vectors, chunkCount, err := s.runDocumentIndex(ctx, document, *knowledgeBase)
|
||||||
|
if err != nil {
|
||||||
|
return fail(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := s.markDocumentIndexIndexed(document.ID); err != nil {
|
||||||
|
slog.Error("Failed to mark knowledge document index as indexed", "document_id", document.ID, "error", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
slog.Info("Document indexed successfully",
|
||||||
|
slog.Any("document_id", document.ID),
|
||||||
|
slog.Any("chunks_count", chunkCount),
|
||||||
|
slog.Any("vectors_count", len(vectors)),
|
||||||
|
slog.Any("time_taken", time.Since(start).String()),
|
||||||
|
)
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) IndexFAQByID(ctx context.Context, faqID int64) error {
|
||||||
|
faq, err := s.loadFAQByID(faqID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := s.markFAQIndexPending(faq.ID); err != nil {
|
||||||
|
slog.Error("Failed to mark knowledge faq index as pending", "faq_id", faq.ID, "error", err)
|
||||||
|
}
|
||||||
|
fail := func(err error) error {
|
||||||
|
if updateErr := s.markFAQIndexFailed(faq.ID, err); updateErr != nil {
|
||||||
|
slog.Error("Failed to mark knowledge faq index as failed", "faq_id", faq.ID, "error", updateErr)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
knowledgeBase, err := s.loadFAQKnowledgeBase(*faq)
|
||||||
|
if err != nil {
|
||||||
|
return fail(err)
|
||||||
|
}
|
||||||
|
if err := s.runFAQIndex(ctx, *faq, *knowledgeBase); err != nil {
|
||||||
|
return fail(err)
|
||||||
|
}
|
||||||
|
if err := s.markFAQIndexIndexed(faq.ID); err != nil {
|
||||||
|
slog.Error("Failed to mark knowledge faq index as indexed", "faq_id", faq.ID, "error", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) RemoveDocumentIndex(ctx context.Context, documentID int64) error {
|
||||||
|
chunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), documentID)
|
||||||
|
if len(chunks) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := s.deleteChunkVectors(ctx, s.collectChunkVectorIDs(chunks)); err != nil {
|
||||||
|
slog.Error("Failed to delete vectors", "error", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := repositories.KnowledgeChunkRepository.DeleteByDocumentID(sqls.DB(), documentID); err != nil {
|
||||||
|
return fmt.Errorf("failed to delete chunks: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
slog.Info("Document index removed", "document_id", documentID, "chunks_removed", len(chunks))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) RemoveFAQIndex(ctx context.Context, faqID int64) error {
|
||||||
|
chunks := repositories.KnowledgeChunkRepository.FindByFaqID(sqls.DB(), faqID)
|
||||||
|
if len(chunks) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err := s.deleteChunkVectors(ctx, s.collectChunkVectorIDs(chunks)); err != nil {
|
||||||
|
slog.Error("Failed to delete faq vectors", "error", err)
|
||||||
|
}
|
||||||
|
if err := repositories.KnowledgeChunkRepository.DeleteByFaqID(sqls.DB(), faqID); err != nil {
|
||||||
|
return fmt.Errorf("failed to delete faq chunks: %w", err)
|
||||||
|
}
|
||||||
|
slog.Info("FAQ index removed", "faq_id", faqID, "chunks_removed", len(chunks))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) RemoveKnowledgeBaseIndex(ctx context.Context, knowledgeBaseID int64) error {
|
||||||
|
chunks := repositories.KnowledgeChunkRepository.FindByKnowledgeBaseID(sqls.DB(), knowledgeBaseID)
|
||||||
|
if len(chunks) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err := s.deleteChunkVectors(ctx, s.collectChunkVectorIDs(chunks)); err != nil {
|
||||||
|
slog.Error("Failed to delete knowledge base vectors", "knowledge_base_id", knowledgeBaseID, "error", err)
|
||||||
|
}
|
||||||
|
if err := repositories.KnowledgeChunkRepository.DeleteByKnowledgeBaseID(sqls.DB(), knowledgeBaseID); err != nil {
|
||||||
|
return fmt.Errorf("failed to delete chunks for knowledge base %d: %w", knowledgeBaseID, err)
|
||||||
|
}
|
||||||
|
slog.Info("Knowledge base index removed", "knowledge_base_id", knowledgeBaseID, "chunks_removed", len(chunks))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) getCollectionName() string {
|
||||||
|
return knowledgeCollectionName
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildKnowledgeChunkVectorID(knowledgeBaseID int64, documentID int64, chunkNo int) string {
|
||||||
|
raw := fmt.Sprintf("kb:%d:doc:%d:chunk:%d", knowledgeBaseID, documentID, chunkNo)
|
||||||
|
return uuid.NewSHA1(uuid.NameSpaceOID, []byte(raw)).String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildKnowledgeFAQChunkVectorID(knowledgeBaseID int64, faqID int64, chunkNo int) string {
|
||||||
|
raw := fmt.Sprintf("kb:%d:faq:%d:chunk:%d", knowledgeBaseID, faqID, chunkNo)
|
||||||
|
return uuid.NewSHA1(uuid.NameSpaceOID, []byte(raw)).String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildChunkContentHash(content string) string {
|
||||||
|
sum := sha256.Sum256([]byte(content))
|
||||||
|
return hex.EncodeToString(sum[:])
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstPositiveInt(values ...int) int {
|
||||||
|
for _, value := range values {
|
||||||
|
if value > 0 {
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstNonEmptyString(values ...string) string {
|
||||||
|
for _, value := range values {
|
||||||
|
if value != "" {
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) EnsureCollection(ctx context.Context) error {
|
||||||
|
dimension, err := ai.Embedding.GetDimension(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to get embedding dimension: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
collectionName := s.getCollectionName()
|
||||||
|
provider := vectordb.GetProvider()
|
||||||
|
if provider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
|
||||||
|
existing, err := provider.GetCollection(ctx, collectionName)
|
||||||
|
if err == nil && existing != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return provider.CreateCollection(ctx, collectionName, dimension)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) RebuildKnowledgeBaseIndex(ctx context.Context, knowledgeBaseID int64) error {
|
||||||
|
knowledgeBase := repositories.KnowledgeBaseRepository.Get(sqls.DB(), knowledgeBaseID)
|
||||||
|
if knowledgeBase == nil {
|
||||||
|
return fmt.Errorf("knowledge base not found: %d", knowledgeBaseID)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := s.resetKnowledgeBaseIndexStorage(ctx, knowledgeBaseID); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
successCount := 0
|
||||||
|
failedCount := 0
|
||||||
|
if knowledgeBase.KnowledgeType == string(enums.KnowledgeBaseTypeFAQ) {
|
||||||
|
faqs := repositories.KnowledgeFAQRepository.Find(sqls.DB(), sqls.NewCnd().
|
||||||
|
Eq("knowledge_base_id", knowledgeBaseID).
|
||||||
|
Where("status != ?", enums.StatusDeleted))
|
||||||
|
if len(faqs) == 0 {
|
||||||
|
slog.Info("No faqs found in knowledge base, nothing to rebuild", "knowledge_base_id", knowledgeBaseID)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
slog.Info("Rebuilding faq knowledge base index", "knowledge_base_id", knowledgeBaseID, "faq_count", len(faqs))
|
||||||
|
for _, faq := range faqs {
|
||||||
|
if err := s.IndexFAQByID(ctx, faq.ID); err != nil {
|
||||||
|
slog.Error("Failed to index faq", "faq_id", faq.ID, "error", err)
|
||||||
|
failedCount++
|
||||||
|
} else {
|
||||||
|
successCount++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
documents := repositories.KnowledgeDocumentRepository.Find(sqls.DB(), sqls.NewCnd().
|
||||||
|
Eq("knowledge_base_id", knowledgeBaseID).
|
||||||
|
Where("status != ?", enums.StatusDeleted))
|
||||||
|
if len(documents) == 0 {
|
||||||
|
slog.Info("No documents found in knowledge base, nothing to rebuild", "knowledge_base_id", knowledgeBaseID)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
documentIDs := make([]int64, 0, len(documents))
|
||||||
|
for _, doc := range documents {
|
||||||
|
documentIDs = append(documentIDs, doc.ID)
|
||||||
|
}
|
||||||
|
if err := s.markKnowledgeBaseDocumentsIndexPending(knowledgeBaseID, documentIDs); err != nil {
|
||||||
|
slog.Error("Failed to mark knowledge base documents index as pending", "knowledge_base_id", knowledgeBaseID, "error", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
slog.Info("Rebuilding knowledge base index", "knowledge_base_id", knowledgeBaseID, "document_count", len(documents))
|
||||||
|
for _, doc := range documents {
|
||||||
|
if err := s.IndexDocumentByID(ctx, doc.ID); err != nil {
|
||||||
|
slog.Error("Failed to index document", "document_id", doc.ID, "error", err)
|
||||||
|
failedCount++
|
||||||
|
} else {
|
||||||
|
successCount++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
slog.Info("Knowledge base index rebuild completed",
|
||||||
|
"knowledge_base_id", knowledgeBaseID,
|
||||||
|
"success_count", successCount,
|
||||||
|
"failed_count", failedCount)
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildFAQChunkContent(faq models.KnowledgeFAQ) string {
|
||||||
|
parts := []string{fmt.Sprintf("问题:%s", faq.Question)}
|
||||||
|
var similarQuestions []string
|
||||||
|
if faq.SimilarQuestions != "" {
|
||||||
|
_ = json.Unmarshal([]byte(faq.SimilarQuestions), &similarQuestions)
|
||||||
|
}
|
||||||
|
if len(similarQuestions) > 0 {
|
||||||
|
parts = append(parts, fmt.Sprintf("相似问:%s", joinSimilarQuestions(similarQuestions)))
|
||||||
|
}
|
||||||
|
parts = append(parts, fmt.Sprintf("回答:%s", faq.Answer))
|
||||||
|
content := ""
|
||||||
|
for _, part := range parts {
|
||||||
|
if part == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if content != "" {
|
||||||
|
content += "\n"
|
||||||
|
}
|
||||||
|
content += part
|
||||||
|
}
|
||||||
|
return content
|
||||||
|
}
|
||||||
|
|
||||||
|
func joinSimilarQuestions(items []string) string {
|
||||||
|
result := ""
|
||||||
|
for _, item := range items {
|
||||||
|
if item == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if result != "" {
|
||||||
|
result += ";"
|
||||||
|
}
|
||||||
|
result += item
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) resetKnowledgeBaseIndexStorage(ctx context.Context, knowledgeBaseID int64) error {
|
||||||
|
chunks := repositories.KnowledgeChunkRepository.FindByKnowledgeBaseID(sqls.DB(), knowledgeBaseID)
|
||||||
|
return s.cleanupKnowledgeBaseChunks(ctx, knowledgeBaseID, chunks)
|
||||||
|
}
|
||||||
49
internal/ai/rag/index_cleanup_helpers.go
Normal file
49
internal/ai/rag/index_cleanup_helpers.go
Normal file
@@ -0,0 +1,49 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/common/strs"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) collectChunkVectorIDs(chunks []models.KnowledgeChunk) []string {
|
||||||
|
vectorIDs := make([]string, 0, len(chunks))
|
||||||
|
for _, chunk := range chunks {
|
||||||
|
if strs.IsNotBlank(chunk.VectorID) {
|
||||||
|
vectorIDs = append(vectorIDs, chunk.VectorID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return vectorIDs
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) deleteChunkVectors(ctx context.Context, vectorIDs []string) error {
|
||||||
|
if len(vectorIDs) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
provider := vectordb.GetProvider()
|
||||||
|
if provider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return provider.DeleteVectors(ctx, s.getCollectionName(), vectorIDs)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) cleanupKnowledgeBaseChunks(ctx context.Context, knowledgeBaseID int64, chunks []models.KnowledgeChunk) error {
|
||||||
|
vectorIDs := s.collectChunkVectorIDs(chunks)
|
||||||
|
if err := s.deleteChunkVectors(ctx, vectorIDs); err != nil {
|
||||||
|
return fmt.Errorf("failed to delete vectors for knowledge base %d before rebuild: %w", knowledgeBaseID, err)
|
||||||
|
}
|
||||||
|
if err := repositories.KnowledgeChunkRepository.DeleteByKnowledgeBaseID(sqls.DB(), knowledgeBaseID); err != nil {
|
||||||
|
return fmt.Errorf("failed to clear chunks before rebuild: %w", err)
|
||||||
|
}
|
||||||
|
slog.Info("Knowledge base index storage reset",
|
||||||
|
"knowledge_base_id", knowledgeBaseID,
|
||||||
|
"collection", s.getCollectionName())
|
||||||
|
return nil
|
||||||
|
}
|
||||||
110
internal/ai/rag/index_document_helpers.go
Normal file
110
internal/ai/rag/index_document_helpers.go
Normal file
@@ -0,0 +1,110 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
ragchunk "agent-desk/internal/ai/rag/chunk"
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) buildDocumentChunkRequest(document models.KnowledgeDocument, knowledgeBase models.KnowledgeBase) *ragchunk.ChunkRequest {
|
||||||
|
return &ragchunk.ChunkRequest{
|
||||||
|
KnowledgeBaseID: document.KnowledgeBaseID,
|
||||||
|
DocumentID: document.ID,
|
||||||
|
DocumentTitle: document.Title,
|
||||||
|
ContentType: document.ContentType,
|
||||||
|
Content: document.Content,
|
||||||
|
PlainText: ExtractPlainText(document.Content, document.ContentType),
|
||||||
|
Options: ragchunk.ChunkOptions{
|
||||||
|
Provider: firstNonEmptyString(knowledgeBase.ChunkProvider, s.chunkConfig.Provider),
|
||||||
|
TargetTokens: firstPositiveInt(knowledgeBase.ChunkTargetTokens, s.chunkConfig.TargetTokens),
|
||||||
|
MaxTokens: firstPositiveInt(knowledgeBase.ChunkMaxTokens, s.chunkConfig.MaxTokens),
|
||||||
|
OverlapTokens: firstPositiveInt(knowledgeBase.ChunkOverlapTokens, s.chunkConfig.OverlapTokens),
|
||||||
|
EnableFallback: s.chunkConfig.EnableFallback,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) buildDocumentChunks(ctx context.Context, document models.KnowledgeDocument, knowledgeBase models.KnowledgeBase) ([]ragchunk.ChunkResult, error) {
|
||||||
|
chunks, err := s.registry.Chunk(ctx, s.buildDocumentChunkRequest(document, knowledgeBase))
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to chunk document: %w", err)
|
||||||
|
}
|
||||||
|
if len(chunks) == 0 {
|
||||||
|
return nil, fmt.Errorf("no chunks generated from document")
|
||||||
|
}
|
||||||
|
return chunks, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) prepareDocumentVectors(ctx context.Context, knowledgeBase models.KnowledgeBase, document models.KnowledgeDocument, chunks []ragchunk.ChunkResult) ([]vectordb.Vector, []models.KnowledgeChunk, int, error) {
|
||||||
|
vectors := make([]vectordb.Vector, 0, len(chunks))
|
||||||
|
chunkModels := make([]models.KnowledgeChunk, 0, len(chunks))
|
||||||
|
dimension := 0
|
||||||
|
directoryPath := loadKnowledgeDirectoryPath(document.DirectoryID)
|
||||||
|
|
||||||
|
for i, chunk := range chunks {
|
||||||
|
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, chunk.Content)
|
||||||
|
if err != nil {
|
||||||
|
slog.Error("Failed to generate embedding for chunk", "document_id", document.ID, "chunk_index", i, "error", err)
|
||||||
|
return nil, nil, 0, fmt.Errorf("failed to generate embedding for chunk %d: %w", i, err)
|
||||||
|
}
|
||||||
|
if dimension == 0 {
|
||||||
|
dimension = embeddingResult.Dimension
|
||||||
|
}
|
||||||
|
|
||||||
|
chunkID := buildKnowledgeChunkVectorID(knowledgeBase.ID, document.ID, chunk.ChunkNo)
|
||||||
|
providerName := ""
|
||||||
|
if chunk.Metadata != nil {
|
||||||
|
if value, ok := chunk.Metadata["provider"].(string); ok {
|
||||||
|
providerName = value
|
||||||
|
}
|
||||||
|
}
|
||||||
|
sectionPath := joinKnowledgeSectionPath(directoryPath, chunk.SectionPath)
|
||||||
|
now := time.Now()
|
||||||
|
chunkModels = append(chunkModels, models.KnowledgeChunk{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
DocumentID: document.ID,
|
||||||
|
ChunkNo: chunk.ChunkNo,
|
||||||
|
Title: chunk.Title,
|
||||||
|
Content: chunk.Content,
|
||||||
|
ContentHash: buildChunkContentHash(chunk.Content),
|
||||||
|
CharCount: chunk.CharCount,
|
||||||
|
TokenCount: chunk.TokenCount,
|
||||||
|
ChunkType: string(chunk.ChunkType),
|
||||||
|
SectionPath: sectionPath,
|
||||||
|
Provider: providerName,
|
||||||
|
VectorID: chunkID,
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
})
|
||||||
|
|
||||||
|
vectors = append(vectors, vectordb.Vector{
|
||||||
|
ID: chunkID,
|
||||||
|
Vector: embeddingResult.Vector,
|
||||||
|
Payload: vectordb.ChunkPayload{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
DocumentID: document.ID,
|
||||||
|
DocumentTitle: document.Title,
|
||||||
|
ChunkNo: chunk.ChunkNo,
|
||||||
|
ChunkType: string(chunk.ChunkType),
|
||||||
|
SectionPath: sectionPath,
|
||||||
|
Content: chunk.Content,
|
||||||
|
Title: chunk.Title,
|
||||||
|
Provider: providerName,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(vectors) == 0 {
|
||||||
|
return nil, nil, 0, fmt.Errorf("no vectors generated")
|
||||||
|
}
|
||||||
|
return vectors, chunkModels, dimension, nil
|
||||||
|
}
|
||||||
61
internal/ai/rag/index_faq_helpers.go
Normal file
61
internal/ai/rag/index_faq_helpers.go
Normal file
@@ -0,0 +1,61 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai"
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
func buildFAQChunkModel(knowledgeBase models.KnowledgeBase, faq models.KnowledgeFAQ, content string) (models.KnowledgeChunk, string) {
|
||||||
|
chunkID := buildKnowledgeFAQChunkVectorID(knowledgeBase.ID, faq.ID, 0)
|
||||||
|
now := time.Now()
|
||||||
|
sectionPath := loadKnowledgeDirectoryPath(faq.DirectoryID)
|
||||||
|
return models.KnowledgeChunk{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
FaqID: faq.ID,
|
||||||
|
ChunkNo: 0,
|
||||||
|
Title: faq.Question,
|
||||||
|
Content: content,
|
||||||
|
ContentHash: buildChunkContentHash(content),
|
||||||
|
CharCount: len([]rune(content)),
|
||||||
|
TokenCount: len([]rune(content)) / 2,
|
||||||
|
ChunkType: string(enums.KnowledgeChunkTypeFAQ),
|
||||||
|
SectionPath: sectionPath,
|
||||||
|
Provider: string(enums.KnowledgeChunkProviderFAQ),
|
||||||
|
VectorID: chunkID,
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
}, chunkID
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) prepareFAQVector(ctx context.Context, knowledgeBase models.KnowledgeBase, faq models.KnowledgeFAQ, content string) (vectordb.Vector, models.KnowledgeChunk, int, error) {
|
||||||
|
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, content)
|
||||||
|
if err != nil {
|
||||||
|
return vectordb.Vector{}, models.KnowledgeChunk{}, 0, fmt.Errorf("failed to generate embedding for faq %d: %w", faq.ID, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
chunkModel, chunkID := buildFAQChunkModel(knowledgeBase, faq, content)
|
||||||
|
sectionPath := chunkModel.SectionPath
|
||||||
|
vector := vectordb.Vector{
|
||||||
|
ID: chunkID,
|
||||||
|
Vector: embeddingResult.Vector,
|
||||||
|
Payload: vectordb.ChunkPayload{
|
||||||
|
KnowledgeBaseID: knowledgeBase.ID,
|
||||||
|
FaqID: faq.ID,
|
||||||
|
FaqQuestion: faq.Question,
|
||||||
|
ChunkNo: 0,
|
||||||
|
ChunkType: string(enums.KnowledgeChunkTypeFAQ),
|
||||||
|
SectionPath: sectionPath,
|
||||||
|
Content: content,
|
||||||
|
Title: faq.Question,
|
||||||
|
Provider: string(enums.KnowledgeChunkProviderFAQ),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
return vector, chunkModel, embeddingResult.Dimension, nil
|
||||||
|
}
|
||||||
76
internal/ai/rag/index_flow.go
Normal file
76
internal/ai/rag/index_flow.go
Normal file
@@ -0,0 +1,76 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) loadDocumentByID(documentID int64) (*models.KnowledgeDocument, error) {
|
||||||
|
document := repositories.KnowledgeDocumentRepository.Get(sqls.DB(), documentID)
|
||||||
|
if document == nil {
|
||||||
|
return nil, fmt.Errorf("document not found: %d", documentID)
|
||||||
|
}
|
||||||
|
return document, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func loadKnowledgeDirectoryPath(directoryID int64) string {
|
||||||
|
if directoryID <= 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
item := repositories.KnowledgeDirectoryRepository.Get(sqls.DB(), directoryID)
|
||||||
|
if item == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if item.ParentID <= 0 {
|
||||||
|
return item.Name
|
||||||
|
}
|
||||||
|
parent := repositories.KnowledgeDirectoryRepository.Get(sqls.DB(), item.ParentID)
|
||||||
|
if parent == nil {
|
||||||
|
return item.Name
|
||||||
|
}
|
||||||
|
return parent.Name + " / " + item.Name
|
||||||
|
}
|
||||||
|
|
||||||
|
func joinKnowledgeSectionPath(parts ...string) string {
|
||||||
|
ret := make([]string, 0, len(parts))
|
||||||
|
for _, item := range parts {
|
||||||
|
item = strings.TrimSpace(item)
|
||||||
|
if item != "" {
|
||||||
|
ret = append(ret, item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return strings.Join(ret, " / ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) loadFAQByID(faqID int64) (*models.KnowledgeFAQ, error) {
|
||||||
|
faq := repositories.KnowledgeFAQRepository.Get(sqls.DB(), faqID)
|
||||||
|
if faq == nil {
|
||||||
|
return nil, fmt.Errorf("faq not found: %d", faqID)
|
||||||
|
}
|
||||||
|
return faq, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) loadDocumentKnowledgeBase(document models.KnowledgeDocument) (*models.KnowledgeBase, error) {
|
||||||
|
knowledgeBase := repositories.KnowledgeBaseRepository.Get(sqls.DB(), document.KnowledgeBaseID)
|
||||||
|
if knowledgeBase == nil {
|
||||||
|
return nil, fmt.Errorf("knowledge base not found: %d", document.KnowledgeBaseID)
|
||||||
|
}
|
||||||
|
return knowledgeBase, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) loadFAQKnowledgeBase(faq models.KnowledgeFAQ) (*models.KnowledgeBase, error) {
|
||||||
|
knowledgeBase := repositories.KnowledgeBaseRepository.Get(sqls.DB(), faq.KnowledgeBaseID)
|
||||||
|
if knowledgeBase == nil {
|
||||||
|
return nil, fmt.Errorf("knowledge base not found: %d", faq.KnowledgeBaseID)
|
||||||
|
}
|
||||||
|
if knowledgeBase.KnowledgeType != string(enums.KnowledgeBaseTypeFAQ) {
|
||||||
|
return nil, fmt.Errorf("knowledge base %d is not faq type", knowledgeBase.ID)
|
||||||
|
}
|
||||||
|
return knowledgeBase, nil
|
||||||
|
}
|
||||||
91
internal/ai/rag/index_retrieve_helpers_test.go
Normal file
91
internal/ai/rag/index_retrieve_helpers_test.go
Normal file
@@ -0,0 +1,91 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuildFAQChunkContent(t *testing.T) {
|
||||||
|
faq := models.KnowledgeFAQ{
|
||||||
|
Question: "如何退款",
|
||||||
|
SimilarQuestions: `["退款怎么申请","申请售后"]`,
|
||||||
|
Answer: "在订单页发起退款。",
|
||||||
|
}
|
||||||
|
|
||||||
|
content := buildFAQChunkContent(faq)
|
||||||
|
if !strings.Contains(content, "问题:如何退款") {
|
||||||
|
t.Fatalf("expected question in content, got %q", content)
|
||||||
|
}
|
||||||
|
if !strings.Contains(content, "相似问:退款怎么申请;申请售后") {
|
||||||
|
t.Fatalf("expected similar questions in content, got %q", content)
|
||||||
|
}
|
||||||
|
if !strings.Contains(content, "回答:在订单页发起退款。") {
|
||||||
|
t.Fatalf("expected answer in content, got %q", content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildFAQChunkModel(t *testing.T) {
|
||||||
|
knowledgeBase := models.KnowledgeBase{ID: 11}
|
||||||
|
faq := models.KnowledgeFAQ{ID: 22, Question: "如何退款"}
|
||||||
|
|
||||||
|
chunk, chunkID := buildFAQChunkModel(knowledgeBase, faq, "问题:如何退款\n回答:在订单页发起退款。")
|
||||||
|
if chunkID == "" {
|
||||||
|
t.Fatalf("expected chunk id")
|
||||||
|
}
|
||||||
|
if chunk.KnowledgeBaseID != 11 || chunk.FaqID != 22 {
|
||||||
|
t.Fatalf("unexpected chunk identity: %#v", chunk)
|
||||||
|
}
|
||||||
|
if chunk.VectorID != chunkID {
|
||||||
|
t.Fatalf("expected vector id to match chunk id")
|
||||||
|
}
|
||||||
|
if chunk.Title != "如何退款" {
|
||||||
|
t.Fatalf("unexpected title: %q", chunk.Title)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNormalizeContextResultsMergesAndDedupes(t *testing.T) {
|
||||||
|
results := normalizeContextResults([]RetrieveResult{
|
||||||
|
{DocumentID: 1, ChunkNo: 1, SectionPath: "A", Content: "第一段", Score: 0.7},
|
||||||
|
{DocumentID: 1, ChunkNo: 2, SectionPath: "A", Content: "第二段", Score: 0.9},
|
||||||
|
{DocumentID: 1, ChunkNo: 3, SectionPath: "A", Content: "第三段", Score: 0.6},
|
||||||
|
{DocumentID: 2, ChunkNo: 1, Title: "标题", Content: "独立段", Score: 0.5},
|
||||||
|
{FaqID: 9, FaqQuestion: "FAQ", Content: "FAQ内容", Score: 0.8},
|
||||||
|
{FaqID: 9, FaqQuestion: "FAQ", Content: "FAQ重复", Score: 0.7},
|
||||||
|
})
|
||||||
|
|
||||||
|
if len(results) != 3 {
|
||||||
|
t.Fatalf("expected 3 normalized results, got %d", len(results))
|
||||||
|
}
|
||||||
|
if !strings.Contains(results[0].Content, "第一段\n第二段") {
|
||||||
|
t.Fatalf("expected adjacent document chunks to merge, got %q", results[0].Content)
|
||||||
|
}
|
||||||
|
if results[0].Score != 0.9 {
|
||||||
|
t.Fatalf("expected merged score to keep max score, got %v", results[0].Score)
|
||||||
|
}
|
||||||
|
if results[1].Content != "独立段" {
|
||||||
|
t.Fatalf("expected section duplicates to be removed after merge, got %q", results[1].Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildContextChunkText(t *testing.T) {
|
||||||
|
faqText := buildContextChunkText(RetrieveResult{
|
||||||
|
FaqID: 1,
|
||||||
|
FaqQuestion: "如何退款",
|
||||||
|
Content: "在订单页发起退款。",
|
||||||
|
})
|
||||||
|
if !strings.Contains(faqText, "【FAQ:如何退款】") {
|
||||||
|
t.Fatalf("unexpected faq context text: %q", faqText)
|
||||||
|
}
|
||||||
|
|
||||||
|
docText := buildContextChunkText(RetrieveResult{
|
||||||
|
DocumentID: 2,
|
||||||
|
DocumentTitle: "退款文档",
|
||||||
|
SectionPath: "售后/退款",
|
||||||
|
Content: "文档内容",
|
||||||
|
})
|
||||||
|
if !strings.Contains(docText, "【文档:退款文档|章节:售后/退款】") {
|
||||||
|
t.Fatalf("unexpected document context text: %q", docText)
|
||||||
|
}
|
||||||
|
}
|
||||||
110
internal/ai/rag/index_run_helpers.go
Normal file
110
internal/ai/rag/index_run_helpers.go
Normal file
@@ -0,0 +1,110 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai"
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) runDocumentIndex(ctx context.Context, document models.KnowledgeDocument, knowledgeBase models.KnowledgeBase) ([]vectordb.Vector, int, error) {
|
||||||
|
existingChunks := repositories.KnowledgeChunkRepository.FindByDocumentID(sqls.DB(), document.ID)
|
||||||
|
chunks, err := s.buildDocumentChunks(ctx, document, knowledgeBase)
|
||||||
|
if err != nil {
|
||||||
|
return nil, 0, err
|
||||||
|
}
|
||||||
|
|
||||||
|
collectionName := s.getCollectionName()
|
||||||
|
provider := vectordb.GetProvider()
|
||||||
|
if provider == nil {
|
||||||
|
return nil, 0, fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
||||||
|
return nil, 0, fmt.Errorf("failed to get embedding model: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
vectors, chunkModels, dimension, err := s.prepareDocumentVectors(ctx, knowledgeBase, document, chunks)
|
||||||
|
if err != nil {
|
||||||
|
return nil, 0, err
|
||||||
|
}
|
||||||
|
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
|
||||||
|
return nil, 0, err
|
||||||
|
}
|
||||||
|
if err := provider.UpsertVectors(ctx, collectionName, vectors); err != nil {
|
||||||
|
return nil, 0, fmt.Errorf("failed to upsert vectors: %w", err)
|
||||||
|
}
|
||||||
|
if err := repositories.KnowledgeChunkRepository.ReplaceByDocumentID(sqls.DB(), document.ID, chunkModels); err != nil {
|
||||||
|
return nil, 0, fmt.Errorf("failed to save chunks: %w", err)
|
||||||
|
}
|
||||||
|
if staleVectorIDs := s.collectStaleVectorIDs(existingChunks, vectors); len(staleVectorIDs) > 0 {
|
||||||
|
if err := provider.DeleteVectors(ctx, collectionName, staleVectorIDs); err != nil {
|
||||||
|
slog.Error("Failed to delete stale document vectors", "document_id", document.ID, "error", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return vectors, len(chunks), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) runFAQIndex(ctx context.Context, faq models.KnowledgeFAQ, knowledgeBase models.KnowledgeBase) error {
|
||||||
|
existingChunks := repositories.KnowledgeChunkRepository.FindByFaqID(sqls.DB(), faq.ID)
|
||||||
|
content := buildFAQChunkContent(faq)
|
||||||
|
if content == "" {
|
||||||
|
return fmt.Errorf("faq content is empty")
|
||||||
|
}
|
||||||
|
|
||||||
|
provider := vectordb.GetProvider()
|
||||||
|
if provider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
if _, err := ai.Embedding.GetModel(ctx); err != nil {
|
||||||
|
return fmt.Errorf("failed to get embedding model: %w", err)
|
||||||
|
}
|
||||||
|
vector, chunkModel, dimension, err := s.prepareFAQVector(ctx, knowledgeBase, faq, content)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
collectionName := s.getCollectionName()
|
||||||
|
if err := s.ensureCollection(ctx, provider, collectionName, dimension); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := provider.UpsertVectors(ctx, collectionName, []vectordb.Vector{vector}); err != nil {
|
||||||
|
return fmt.Errorf("failed to upsert vectors: %w", err)
|
||||||
|
}
|
||||||
|
if err := repositories.KnowledgeChunkRepository.ReplaceByFaqID(sqls.DB(), faq.ID, &chunkModel); err != nil {
|
||||||
|
return fmt.Errorf("failed to save faq chunk: %w", err)
|
||||||
|
}
|
||||||
|
if staleVectorIDs := s.collectStaleVectorIDs(existingChunks, []vectordb.Vector{vector}); len(staleVectorIDs) > 0 {
|
||||||
|
if err := provider.DeleteVectors(ctx, collectionName, staleVectorIDs); err != nil {
|
||||||
|
slog.Error("Failed to delete stale faq vectors", "faq_id", faq.ID, "error", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) collectStaleVectorIDs(existingChunks []models.KnowledgeChunk, currentVectors []vectordb.Vector) []string {
|
||||||
|
currentVectorIDs := make(map[string]struct{}, len(currentVectors))
|
||||||
|
for _, vector := range currentVectors {
|
||||||
|
if vector.ID == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
currentVectorIDs[vector.ID] = struct{}{}
|
||||||
|
}
|
||||||
|
staleVectorIDs := make([]string, 0, len(existingChunks))
|
||||||
|
for _, chunk := range existingChunks {
|
||||||
|
if chunk.VectorID == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := currentVectorIDs[chunk.VectorID]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
staleVectorIDs = append(staleVectorIDs, chunk.VectorID)
|
||||||
|
}
|
||||||
|
return staleVectorIDs
|
||||||
|
}
|
||||||
91
internal/ai/rag/index_status_helpers.go
Normal file
91
internal/ai/rag/index_status_helpers.go
Normal file
@@ -0,0 +1,91 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) markDocumentIndexPending(documentID int64) error {
|
||||||
|
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
||||||
|
"indexed_at": nil,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markDocumentIndexIndexed(documentID int64) error {
|
||||||
|
now := time.Now()
|
||||||
|
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusIndexed,
|
||||||
|
"indexed_at": &now,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": now,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markDocumentIndexFailed(documentID int64, err error) error {
|
||||||
|
return repositories.KnowledgeDocumentRepository.Updates(sqls.DB(), documentID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusFailed,
|
||||||
|
"index_error": truncateIndexError(err),
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markKnowledgeBaseDocumentsIndexPending(knowledgeBaseID int64, documentIDs []int64) error {
|
||||||
|
if len(documentIDs) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return sqls.DB().Model(&models.KnowledgeDocument{}).
|
||||||
|
Where("knowledge_base_id = ?", knowledgeBaseID).
|
||||||
|
Where("id IN ?", documentIDs).
|
||||||
|
Updates(map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
||||||
|
"indexed_at": nil,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
}).Error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markFAQIndexPending(faqID int64) error {
|
||||||
|
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusPending,
|
||||||
|
"indexed_at": nil,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markFAQIndexIndexed(faqID int64) error {
|
||||||
|
now := time.Now()
|
||||||
|
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusIndexed,
|
||||||
|
"indexed_at": &now,
|
||||||
|
"index_error": "",
|
||||||
|
"updated_at": now,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *index) markFAQIndexFailed(faqID int64, err error) error {
|
||||||
|
return repositories.KnowledgeFAQRepository.Updates(sqls.DB(), faqID, map[string]any{
|
||||||
|
"index_status": enums.KnowledgeDocumentIndexStatusFailed,
|
||||||
|
"index_error": truncateIndexError(err),
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateIndexError(err error) string {
|
||||||
|
if err == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
message := err.Error()
|
||||||
|
if len(message) <= 1000 {
|
||||||
|
return message
|
||||||
|
}
|
||||||
|
return message[:1000]
|
||||||
|
}
|
||||||
24
internal/ai/rag/index_storage_helpers.go
Normal file
24
internal/ai/rag/index_storage_helpers.go
Normal file
@@ -0,0 +1,24 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *index) ensureCollection(ctx context.Context, provider vectordb.Provider, collectionName string, dimension int) error {
|
||||||
|
collectionInfo, err := provider.GetCollection(ctx, collectionName)
|
||||||
|
if err == nil && collectionInfo != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if dimension <= 0 {
|
||||||
|
return fmt.Errorf("invalid embedding dimension: %d", dimension)
|
||||||
|
}
|
||||||
|
if err := provider.CreateCollection(ctx, collectionName, dimension); err != nil {
|
||||||
|
return fmt.Errorf("failed to create collection: %w", err)
|
||||||
|
}
|
||||||
|
slog.Info("Created collection for knowledge base", "collection", collectionName, "dimension", dimension)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
140
internal/ai/rag/rerank.go
Normal file
140
internal/ai/rag/rerank.go
Normal file
@@ -0,0 +1,140 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"sort"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
type rerank struct{}
|
||||||
|
|
||||||
|
var Rerank = &rerank{}
|
||||||
|
|
||||||
|
func (s *rerank) Rerank(ctx context.Context, query string, documents []string, topN int) ([]RerankResult, error) {
|
||||||
|
if len(documents) == 0 {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if topN <= 0 {
|
||||||
|
topN = len(documents)
|
||||||
|
}
|
||||||
|
|
||||||
|
results, err := s.callRerankAPI(ctx, query, documents, topN)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *rerank) callRerankAPI(ctx context.Context, query string, documents []string, topN int) ([]RerankResult, error) {
|
||||||
|
config, err := ai.GetEnabledAIConfig(enums.AIModelTypeRerank)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
reqBody := RerankRequest{
|
||||||
|
Model: config.ModelName,
|
||||||
|
Query: query,
|
||||||
|
Documents: documents,
|
||||||
|
TopN: topN,
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonBody, err := json.Marshal(reqBody)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, "POST", config.BaseURL+"/v1/rerank", bytes.NewBuffer(jsonBody))
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
req.Header.Set("Authorization", "Bearer "+config.APIKey)
|
||||||
|
|
||||||
|
client := &http.Client{
|
||||||
|
Timeout: time.Duration(config.TimeoutMS) * time.Millisecond,
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to call rerank API: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
body, err := io.ReadAll(resp.Body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var rerankResp RerankResponse
|
||||||
|
if err := json.Unmarshal(body, &rerankResp); err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
results := make([]RerankResult, 0, len(rerankResp.Results))
|
||||||
|
for _, r := range rerankResp.Results {
|
||||||
|
results = append(results, RerankResult{
|
||||||
|
Index: r.Index,
|
||||||
|
RelevanceScore: r.RelevanceScore,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *rerank) RerankResults(ctx context.Context, query string, results []RetrieveResult, topN int) ([]RetrieveResult, error) {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if topN <= 0 {
|
||||||
|
topN = len(results)
|
||||||
|
}
|
||||||
|
|
||||||
|
documents := make([]string, 0, len(results))
|
||||||
|
for _, r := range results {
|
||||||
|
documents = append(documents, r.Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
rerankResults, err := s.Rerank(ctx, query, documents, topN)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
rerankedResults := make([]RetrieveResult, 0, len(rerankResults))
|
||||||
|
for _, rr := range rerankResults {
|
||||||
|
if rr.Index < len(results) {
|
||||||
|
result := results[rr.Index]
|
||||||
|
result.Score = float32(rr.RelevanceScore)
|
||||||
|
rerankedResults = append(rerankedResults, result)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return rerankedResults, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *rerank) SimpleRerank(query string, results []RetrieveResult, topN int) []RetrieveResult {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
sort.Slice(results, func(i, j int) bool {
|
||||||
|
return results[i].Score > results[j].Score
|
||||||
|
})
|
||||||
|
|
||||||
|
if topN > 0 && len(results) > topN {
|
||||||
|
return results[:topN]
|
||||||
|
}
|
||||||
|
|
||||||
|
return results
|
||||||
|
}
|
||||||
230
internal/ai/rag/retrieve.go
Normal file
230
internal/ai/rag/retrieve.go
Normal file
@@ -0,0 +1,230 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
type retrieve struct {
|
||||||
|
}
|
||||||
|
|
||||||
|
var Retrieve = &retrieve{}
|
||||||
|
|
||||||
|
func (s *retrieve) Retrieve(ctx context.Context, req RetrieveRequest) ([]RetrieveResult, error) {
|
||||||
|
results, _, err := s.RetrieveWithTrace(ctx, req)
|
||||||
|
return results, err
|
||||||
|
}
|
||||||
|
|
||||||
|
type RetrieveTrace struct {
|
||||||
|
EmbeddingMs int64
|
||||||
|
VectorSearchMs int64
|
||||||
|
HydrateMs int64
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) RetrieveWithTrace(ctx context.Context, req RetrieveRequest) ([]RetrieveResult, *RetrieveTrace, error) {
|
||||||
|
trace := newRetrieveTrace()
|
||||||
|
retrievableKnowledgeBases, _, ok := s.prepareRetrievableKnowledgeBases(req, trace)
|
||||||
|
if !ok {
|
||||||
|
return nil, trace, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
searchResults, searchTrace, err := s.searchKnowledgeBaseVectors(ctx, req, retrievableKnowledgeBases)
|
||||||
|
if err != nil {
|
||||||
|
applySearchTrace(trace, searchTrace)
|
||||||
|
return nil, trace, err
|
||||||
|
}
|
||||||
|
applySearchTrace(trace, searchTrace)
|
||||||
|
|
||||||
|
if len(searchResults) == 0 {
|
||||||
|
return nil, trace, nil
|
||||||
|
}
|
||||||
|
results, hydrateMs := s.hydrateRetrieveResults(searchResults)
|
||||||
|
trace.HydrateMs = hydrateMs
|
||||||
|
|
||||||
|
return results, trace, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func extractChunkType(payload vectordb.ChunkPayload) string {
|
||||||
|
if payload.ChunkType != "" {
|
||||||
|
return payload.ChunkType
|
||||||
|
}
|
||||||
|
return string(enums.KnowledgeChunkTypeText)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) logEmptySearchDiagnostics(ctx context.Context, provider vectordb.Provider, collectionName string, vector []float32, topK int, scoreThreshold float32, knowledgeBaseIDs []int64, req RetrieveRequest) {
|
||||||
|
rawResults, err := provider.Search(ctx, &vectordb.SearchRequest{
|
||||||
|
CollectionName: collectionName,
|
||||||
|
Vector: vector,
|
||||||
|
TopK: topK,
|
||||||
|
ScoreThreshold: 0,
|
||||||
|
Filter: &vectordb.SearchFilter{
|
||||||
|
KnowledgeBaseIDs: knowledgeBaseIDs,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
slog.Warn("Knowledge retrieve diagnostics failed",
|
||||||
|
"knowledge_base_ids", fmt.Sprint(knowledgeBaseIDs),
|
||||||
|
"collection", collectionName,
|
||||||
|
"query", truncateForLog(req.Query, 80),
|
||||||
|
"score_threshold", scoreThreshold,
|
||||||
|
"error", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if len(rawResults) == 0 {
|
||||||
|
slog.Info("Knowledge retrieve returned no candidates even without threshold",
|
||||||
|
"knowledge_base_ids", fmt.Sprint(knowledgeBaseIDs),
|
||||||
|
"collection", collectionName,
|
||||||
|
"query", truncateForLog(req.Query, 80),
|
||||||
|
"score_threshold", scoreThreshold)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
candidates := make([]string, 0, len(rawResults))
|
||||||
|
for _, item := range rawResults {
|
||||||
|
candidates = append(candidates, fmt.Sprintf("%s:%.4f", item.ID, item.Score))
|
||||||
|
}
|
||||||
|
|
||||||
|
slog.Info("Knowledge retrieve filtered all candidates by score threshold",
|
||||||
|
"knowledge_base_ids", fmt.Sprint(knowledgeBaseIDs),
|
||||||
|
"collection", collectionName,
|
||||||
|
"query", truncateForLog(req.Query, 80),
|
||||||
|
"score_threshold", scoreThreshold,
|
||||||
|
"top_candidates", strings.Join(candidates, ","))
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateForLog(text string, limit int) string {
|
||||||
|
if limit <= 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
runes := []rune(strings.TrimSpace(text))
|
||||||
|
if len(runes) <= limit {
|
||||||
|
return string(runes)
|
||||||
|
}
|
||||||
|
return string(runes[:limit]) + "..."
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) RetrieveWithRerank(ctx context.Context, req RetrieveRequest, rerankLimit int) ([]RetrieveResult, error) {
|
||||||
|
results, err := s.Retrieve(ctx, req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(results) <= rerankLimit {
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
rerankedResults, err := s.rerank(ctx, req.Query, results, rerankLimit)
|
||||||
|
if err != nil {
|
||||||
|
slog.Warn("Rerank failed, returning original results", "error", err)
|
||||||
|
if len(results) > rerankLimit {
|
||||||
|
return results[:rerankLimit], nil
|
||||||
|
}
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return rerankedResults, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) rerank(ctx context.Context, query string, results []RetrieveResult, limit int) ([]RetrieveResult, error) {
|
||||||
|
return Rerank.RerankResults(ctx, query, results, limit)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) GetKnowledgeBaseStats(ctx context.Context, knowledgeBaseID int64) (*KnowledgeBaseStats, error) {
|
||||||
|
knowledgeBase := repositories.KnowledgeBaseRepository.Get(sqls.DB(), knowledgeBaseID)
|
||||||
|
if knowledgeBase == nil {
|
||||||
|
return nil, fmt.Errorf("knowledge base not found")
|
||||||
|
}
|
||||||
|
|
||||||
|
documentCount := repositories.KnowledgeDocumentRepository.CountByKnowledgeBaseID(sqls.DB(), knowledgeBaseID)
|
||||||
|
chunkCount := repositories.KnowledgeChunkRepository.CountByKnowledgeBaseID(sqls.DB(), knowledgeBaseID)
|
||||||
|
|
||||||
|
publishedCount := repositories.KnowledgeDocumentRepository.Count(sqls.DB(), sqls.NewCnd().
|
||||||
|
Eq("knowledge_base_id", knowledgeBaseID).
|
||||||
|
Eq("status", enums.StatusOk))
|
||||||
|
|
||||||
|
return &KnowledgeBaseStats{
|
||||||
|
KnowledgeBaseID: knowledgeBaseID,
|
||||||
|
DocumentCount: documentCount,
|
||||||
|
PublishedCount: publishedCount,
|
||||||
|
ChunkCount: chunkCount,
|
||||||
|
VectorCount: int(chunkCount),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeKnowledgeBaseIDs(ids []int64) []int64 {
|
||||||
|
if len(ids) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
seen := make(map[int64]struct{}, len(ids))
|
||||||
|
normalized := make([]int64, 0, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
if id <= 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[id]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[id] = struct{}{}
|
||||||
|
normalized = append(normalized, id)
|
||||||
|
}
|
||||||
|
return normalized
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveKnowledgeBaseSearchOptions(req RetrieveRequest, knowledgeBase *models.KnowledgeBase) (int, float32) {
|
||||||
|
topK := req.TopK
|
||||||
|
if topK <= 0 && knowledgeBase != nil && knowledgeBase.DefaultTopK > 0 {
|
||||||
|
topK = knowledgeBase.DefaultTopK
|
||||||
|
}
|
||||||
|
if topK <= 0 {
|
||||||
|
topK = 8
|
||||||
|
}
|
||||||
|
|
||||||
|
scoreThreshold := float32(req.ScoreThreshold)
|
||||||
|
if scoreThreshold <= 0 && knowledgeBase != nil && knowledgeBase.DefaultScoreThreshold > 0 {
|
||||||
|
scoreThreshold = float32(knowledgeBase.DefaultScoreThreshold)
|
||||||
|
}
|
||||||
|
if scoreThreshold <= 0 {
|
||||||
|
scoreThreshold = 0.3
|
||||||
|
}
|
||||||
|
return topK, scoreThreshold
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) loadRetrievableKnowledgeBases(ids []int64) []models.KnowledgeBase {
|
||||||
|
if len(ids) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
items := repositories.KnowledgeBaseRepository.Find(sqls.DB(), sqls.NewCnd().In("id", ids))
|
||||||
|
if len(items) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
allowed := make(map[int64]models.KnowledgeBase, len(items))
|
||||||
|
for _, item := range items {
|
||||||
|
if item.Status == enums.StatusOk {
|
||||||
|
allowed[item.ID] = item
|
||||||
|
}
|
||||||
|
}
|
||||||
|
filtered := make([]models.KnowledgeBase, 0, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
if item, ok := allowed[id]; ok {
|
||||||
|
filtered = append(filtered, item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return filtered
|
||||||
|
}
|
||||||
|
|
||||||
|
type KnowledgeBaseStats struct {
|
||||||
|
KnowledgeBaseID int64 `json:"knowledgeBaseId"`
|
||||||
|
DocumentCount int64 `json:"documentCount"`
|
||||||
|
PublishedCount int64 `json:"publishedCount"`
|
||||||
|
ChunkCount int64 `json:"chunkCount"`
|
||||||
|
VectorCount int `json:"vectorCount"`
|
||||||
|
}
|
||||||
150
internal/ai/rag/retrieve_context.go
Normal file
150
internal/ai/rag/retrieve_context.go
Normal file
@@ -0,0 +1,150 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *retrieve) SelectContextResults(results []RetrieveResult, maxTokens int) []RetrieveResult {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
normalizedResults := normalizeContextResults(results)
|
||||||
|
selected := make([]RetrieveResult, 0, len(normalizedResults))
|
||||||
|
totalTokens := 0
|
||||||
|
documentUsage := make(map[int64]int)
|
||||||
|
|
||||||
|
for _, item := range normalizedResults {
|
||||||
|
if documentUsage[item.DocumentID] >= 2 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
chunkText := buildContextChunkText(item)
|
||||||
|
estimatedTokens := len(chunkText) / 2
|
||||||
|
if totalTokens+estimatedTokens > maxTokens {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
selected = append(selected, item)
|
||||||
|
totalTokens += estimatedTokens
|
||||||
|
documentUsage[item.DocumentID]++
|
||||||
|
}
|
||||||
|
return selected
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) BuildContext(_ context.Context, results []RetrieveResult, maxTokens int) string {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
normalizedResults := s.SelectContextResults(results, maxTokens)
|
||||||
|
var builder strings.Builder
|
||||||
|
for _, r := range normalizedResults {
|
||||||
|
builder.WriteString(buildContextChunkText(r))
|
||||||
|
}
|
||||||
|
|
||||||
|
return builder.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeContextResults(results []RetrieveResult) []RetrieveResult {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
merged := mergeAdjacentResults(results)
|
||||||
|
return dedupeSectionResults(merged)
|
||||||
|
}
|
||||||
|
|
||||||
|
func dedupeSectionResults(results []RetrieveResult) []RetrieveResult {
|
||||||
|
seen := make(map[string]struct{})
|
||||||
|
deduped := make([]RetrieveResult, 0, len(results))
|
||||||
|
for _, item := range results {
|
||||||
|
key := buildSectionKey(item)
|
||||||
|
if _, ok := seen[key]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[key] = struct{}{}
|
||||||
|
deduped = append(deduped, item)
|
||||||
|
}
|
||||||
|
return deduped
|
||||||
|
}
|
||||||
|
|
||||||
|
func mergeAdjacentResults(results []RetrieveResult) []RetrieveResult {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
merged := make([]RetrieveResult, 0, len(results))
|
||||||
|
for _, item := range results {
|
||||||
|
if len(merged) == 0 {
|
||||||
|
merged = append(merged, item)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
last := &merged[len(merged)-1]
|
||||||
|
if canMergeContextResult(*last, item) {
|
||||||
|
last.Content = strings.TrimSpace(last.Content + "\n" + item.Content)
|
||||||
|
if item.Score > last.Score {
|
||||||
|
last.Score = item.Score
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
merged = append(merged, item)
|
||||||
|
}
|
||||||
|
return merged
|
||||||
|
}
|
||||||
|
|
||||||
|
func canMergeContextResult(left, right RetrieveResult) bool {
|
||||||
|
if left.FaqID > 0 || right.FaqID > 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if left.DocumentID != right.DocumentID {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if left.SectionPath == "" || right.SectionPath == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if left.SectionPath != right.SectionPath {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return right.ChunkNo == left.ChunkNo+1
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildSectionKey(item RetrieveResult) string {
|
||||||
|
if item.FaqID > 0 {
|
||||||
|
return fmt.Sprintf("faq:%d", item.FaqID)
|
||||||
|
}
|
||||||
|
sectionPath := strings.TrimSpace(item.SectionPath)
|
||||||
|
if sectionPath != "" {
|
||||||
|
return fmt.Sprintf("%d|%s", item.DocumentID, sectionPath)
|
||||||
|
}
|
||||||
|
title := strings.TrimSpace(item.Title)
|
||||||
|
if title != "" {
|
||||||
|
return fmt.Sprintf("%d|%s", item.DocumentID, title)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%d|chunk:%d", item.DocumentID, item.ChunkNo)
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildContextChunkText(item RetrieveResult) string {
|
||||||
|
if item.FaqID > 0 {
|
||||||
|
title := strings.TrimSpace(item.FaqQuestion)
|
||||||
|
if title == "" {
|
||||||
|
title = strings.TrimSpace(item.Title)
|
||||||
|
}
|
||||||
|
if title == "" {
|
||||||
|
title = fmt.Sprintf("FAQ#%d", item.FaqID)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("【FAQ:%s】\n%s\n\n", title, item.Content)
|
||||||
|
}
|
||||||
|
title := strings.TrimSpace(item.DocumentTitle)
|
||||||
|
if title == "" {
|
||||||
|
title = fmt.Sprintf("文档#%d", item.DocumentID)
|
||||||
|
}
|
||||||
|
if item.SectionPath != "" {
|
||||||
|
return fmt.Sprintf("【文档:%s|章节:%s】\n%s\n\n", title, item.SectionPath, item.Content)
|
||||||
|
}
|
||||||
|
if item.Title != "" {
|
||||||
|
return fmt.Sprintf("【文档:%s|标题:%s】\n%s\n\n", title, item.Title, item.Content)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("【文档:%s】\n%s\n\n", title, item.Content)
|
||||||
|
}
|
||||||
38
internal/ai/rag/retrieve_flow.go
Normal file
38
internal/ai/rag/retrieve_flow.go
Normal file
@@ -0,0 +1,38 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
func newRetrieveTrace() *RetrieveTrace {
|
||||||
|
return &RetrieveTrace{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) prepareRetrievableKnowledgeBases(req RetrieveRequest, trace *RetrieveTrace) ([]models.KnowledgeBase, []int64, bool) {
|
||||||
|
if req.Query == "" {
|
||||||
|
return nil, nil, false
|
||||||
|
}
|
||||||
|
knowledgeBaseIDs := normalizeKnowledgeBaseIDs(req.KnowledgeBaseIDs)
|
||||||
|
if len(knowledgeBaseIDs) == 0 {
|
||||||
|
return nil, nil, false
|
||||||
|
}
|
||||||
|
|
||||||
|
retrievableKnowledgeBases := s.loadRetrievableKnowledgeBases(knowledgeBaseIDs)
|
||||||
|
if len(retrievableKnowledgeBases) == 0 {
|
||||||
|
slog.Info("Skip retrieve for non-enabled knowledge bases",
|
||||||
|
"knowledge_base_ids", fmt.Sprint(knowledgeBaseIDs))
|
||||||
|
return nil, knowledgeBaseIDs, false
|
||||||
|
}
|
||||||
|
return retrievableKnowledgeBases, knowledgeBaseIDs, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func applySearchTrace(target *RetrieveTrace, source *RetrieveTrace) {
|
||||||
|
if target == nil || source == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
target.EmbeddingMs = source.EmbeddingMs
|
||||||
|
target.VectorSearchMs = source.VectorSearchMs
|
||||||
|
}
|
||||||
300
internal/ai/rag/retrieve_log.go
Normal file
300
internal/ai/rag/retrieve_log.go
Normal file
@@ -0,0 +1,300 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/dto"
|
||||||
|
"agent-desk/internal/pkg/dto/response"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
var RetrieveLog = &retrieveLog{}
|
||||||
|
|
||||||
|
type retrieveLog struct {
|
||||||
|
}
|
||||||
|
|
||||||
|
type CreateRetrieveLogRequest struct {
|
||||||
|
KnowledgeBaseID int64
|
||||||
|
Channel string
|
||||||
|
Scene string
|
||||||
|
SessionID string
|
||||||
|
ConversationID int64
|
||||||
|
Question string
|
||||||
|
RewriteQuestion string
|
||||||
|
Answer string
|
||||||
|
AnswerStatus int
|
||||||
|
ChunkProvider string
|
||||||
|
ChunkTargetTokens int
|
||||||
|
ChunkMaxTokens int
|
||||||
|
ChunkOverlapTokens int
|
||||||
|
RerankEnabled bool
|
||||||
|
RerankLimit int
|
||||||
|
Hits []response.KnowledgeSearchResult
|
||||||
|
UsedHits []response.KnowledgeSearchResult
|
||||||
|
Citations []response.KnowledgeCitation
|
||||||
|
LatencyMs int64
|
||||||
|
RetrieveMs int64
|
||||||
|
GenerateMs int64
|
||||||
|
PromptTokens int
|
||||||
|
CompletionTokens int
|
||||||
|
ModelName string
|
||||||
|
}
|
||||||
|
|
||||||
|
type retrieveTraceData struct {
|
||||||
|
Retrieve retrieveTraceRetrieve `json:"retrieve"`
|
||||||
|
ChunkConfig retrieveTraceChunkConfig `json:"chunkConfig"`
|
||||||
|
Context retrieveTraceContext `json:"context"`
|
||||||
|
Citations []retrieveTraceCitation `json:"citations"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type retrieveTraceRetrieve struct {
|
||||||
|
Provider string `json:"provider"`
|
||||||
|
RerankEnabled bool `json:"rerankEnabled"`
|
||||||
|
RerankLimit int `json:"rerankLimit"`
|
||||||
|
RawHitCount int `json:"rawHitCount"`
|
||||||
|
ContextHitCount int `json:"contextHitCount"`
|
||||||
|
CitationCount int `json:"citationCount"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type retrieveTraceChunkConfig struct {
|
||||||
|
Provider string `json:"provider"`
|
||||||
|
TargetTokens int `json:"targetTokens"`
|
||||||
|
MaxTokens int `json:"maxTokens"`
|
||||||
|
OverlapTokens int `json:"overlapTokens"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type retrieveTraceContext struct {
|
||||||
|
KnowledgeBaseIDs []int64 `json:"knowledgeBaseIds"`
|
||||||
|
DocumentIDs []int64 `json:"documentIds"`
|
||||||
|
SectionPaths []string `json:"sectionPaths"`
|
||||||
|
UsedChunkKeys []string `json:"usedChunkKeys"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type retrieveTraceCitation struct {
|
||||||
|
DocumentID int64 `json:"documentId"`
|
||||||
|
ChunkNo int `json:"chunkNo"`
|
||||||
|
SectionPath string `json:"sectionPath"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieveLog) FindHitsByRetrieveLogID(retrieveLogID int64) []models.KnowledgeRetrieveHit {
|
||||||
|
if retrieveLogID <= 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var list []models.KnowledgeRetrieveHit
|
||||||
|
sqls.DB().Where("retrieve_log_id = ?", retrieveLogID).Order("rank_no asc, id asc").Find(&list)
|
||||||
|
return list
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieveLog) CreateRetrieveLog(req *CreateRetrieveLogRequest, _ *dto.AuthPrincipal) (*models.KnowledgeRetrieveLog, error) {
|
||||||
|
if req == nil {
|
||||||
|
return nil, fmt.Errorf("retrieve log request is nil")
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
topScore := 0.0
|
||||||
|
if len(req.Hits) > 0 {
|
||||||
|
topScore = req.Hits[0].Score
|
||||||
|
}
|
||||||
|
traceData := buildRetrieveTraceData(req)
|
||||||
|
|
||||||
|
log := &models.KnowledgeRetrieveLog{
|
||||||
|
KnowledgeBaseID: req.KnowledgeBaseID,
|
||||||
|
Channel: req.Channel,
|
||||||
|
Scene: req.Scene,
|
||||||
|
SessionID: req.SessionID,
|
||||||
|
ConversationID: req.ConversationID,
|
||||||
|
RequestID: uuid.New().String(),
|
||||||
|
Question: req.Question,
|
||||||
|
RewriteQuestion: req.RewriteQuestion,
|
||||||
|
Answer: req.Answer,
|
||||||
|
AnswerStatus: req.AnswerStatus,
|
||||||
|
HitCount: len(req.Hits),
|
||||||
|
TopScore: topScore,
|
||||||
|
ChunkProvider: req.ChunkProvider,
|
||||||
|
ChunkTargetTokens: req.ChunkTargetTokens,
|
||||||
|
ChunkMaxTokens: req.ChunkMaxTokens,
|
||||||
|
ChunkOverlapTokens: req.ChunkOverlapTokens,
|
||||||
|
RerankEnabled: req.RerankEnabled,
|
||||||
|
RerankLimit: req.RerankLimit,
|
||||||
|
CitationCount: len(req.Citations),
|
||||||
|
UsedChunkCount: len(req.UsedHits),
|
||||||
|
LatencyMs: req.LatencyMs,
|
||||||
|
RetrieveMs: req.RetrieveMs,
|
||||||
|
GenerateMs: req.GenerateMs,
|
||||||
|
PromptTokens: req.PromptTokens,
|
||||||
|
CompletionTokens: req.CompletionTokens,
|
||||||
|
ModelName: req.ModelName,
|
||||||
|
TraceData: traceData,
|
||||||
|
CreatedAt: now,
|
||||||
|
}
|
||||||
|
|
||||||
|
usedHitKeys := make(map[string]struct{}, len(req.UsedHits))
|
||||||
|
for _, item := range req.UsedHits {
|
||||||
|
usedHitKeys[buildKnowledgeSearchResultKey(item)] = struct{}{}
|
||||||
|
}
|
||||||
|
citationKeys := make(map[string]struct{}, len(req.Citations))
|
||||||
|
for _, item := range req.Citations {
|
||||||
|
citationKeys[buildKnowledgeCitationKey(item)] = struct{}{}
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := sqls.WithTransaction(func(ctx *sqls.TxContext) error {
|
||||||
|
if err := ctx.Tx.Create(log).Error; err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
for i, hit := range req.Hits {
|
||||||
|
hitKey := buildKnowledgeSearchResultKey(hit)
|
||||||
|
hitRecord := &models.KnowledgeRetrieveHit{
|
||||||
|
RetrieveLogID: log.ID,
|
||||||
|
KnowledgeBaseID: hit.KnowledgeBaseID,
|
||||||
|
ChunkID: hit.ChunkID,
|
||||||
|
DocumentID: hit.DocumentID,
|
||||||
|
DocumentTitle: hit.DocumentTitle,
|
||||||
|
FaqID: hit.FaqID,
|
||||||
|
FaqQuestion: hit.FaqQuestion,
|
||||||
|
ChunkNo: hit.ChunkNo,
|
||||||
|
Title: hit.Title,
|
||||||
|
SectionPath: hit.SectionPath,
|
||||||
|
ChunkType: "",
|
||||||
|
Provider: req.ChunkProvider,
|
||||||
|
RankNo: i + 1,
|
||||||
|
Score: hit.Score,
|
||||||
|
RerankScore: hit.RerankScore,
|
||||||
|
UsedInAnswer: hasHitKey(usedHitKeys, hitKey),
|
||||||
|
IsCitation: hasHitKey(citationKeys, hitKey),
|
||||||
|
Snippet: hit.Content,
|
||||||
|
CreatedAt: now,
|
||||||
|
}
|
||||||
|
if err := ctx.Tx.Create(hitRecord).Error; err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return log, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildRetrieveTraceData(req *CreateRetrieveLogRequest) string {
|
||||||
|
trace := retrieveTraceData{
|
||||||
|
Retrieve: retrieveTraceRetrieve{
|
||||||
|
Provider: req.ChunkProvider,
|
||||||
|
RerankEnabled: req.RerankEnabled,
|
||||||
|
RerankLimit: req.RerankLimit,
|
||||||
|
RawHitCount: len(req.Hits),
|
||||||
|
ContextHitCount: len(req.UsedHits),
|
||||||
|
CitationCount: len(req.Citations),
|
||||||
|
},
|
||||||
|
ChunkConfig: retrieveTraceChunkConfig{
|
||||||
|
Provider: req.ChunkProvider,
|
||||||
|
TargetTokens: req.ChunkTargetTokens,
|
||||||
|
MaxTokens: req.ChunkMaxTokens,
|
||||||
|
OverlapTokens: req.ChunkOverlapTokens,
|
||||||
|
},
|
||||||
|
Context: retrieveTraceContext{
|
||||||
|
KnowledgeBaseIDs: distinctKnowledgeBaseIDs(req.UsedHits),
|
||||||
|
DocumentIDs: distinctDocumentIDs(req.UsedHits),
|
||||||
|
SectionPaths: distinctSectionPaths(req.UsedHits),
|
||||||
|
UsedChunkKeys: buildUsedChunkKeys(req.UsedHits),
|
||||||
|
},
|
||||||
|
Citations: buildTraceCitations(req.Citations),
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(trace)
|
||||||
|
if err != nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return string(data)
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildTraceCitations(citations []response.KnowledgeCitation) []retrieveTraceCitation {
|
||||||
|
items := make([]retrieveTraceCitation, 0, len(citations))
|
||||||
|
for _, item := range citations {
|
||||||
|
items = append(items, retrieveTraceCitation{
|
||||||
|
DocumentID: item.DocumentID,
|
||||||
|
ChunkNo: item.ChunkNo,
|
||||||
|
SectionPath: item.SectionPath,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return items
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildUsedChunkKeys(hits []response.KnowledgeSearchResult) []string {
|
||||||
|
keys := make([]string, 0, len(hits))
|
||||||
|
for _, item := range hits {
|
||||||
|
keys = append(keys, buildKnowledgeSearchResultKey(item))
|
||||||
|
}
|
||||||
|
return keys
|
||||||
|
}
|
||||||
|
|
||||||
|
func distinctKnowledgeBaseIDs(hits []response.KnowledgeSearchResult) []int64 {
|
||||||
|
ids := make([]int64, 0)
|
||||||
|
seen := make(map[int64]struct{})
|
||||||
|
for _, item := range hits {
|
||||||
|
if item.KnowledgeBaseID <= 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[item.KnowledgeBaseID]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[item.KnowledgeBaseID] = struct{}{}
|
||||||
|
ids = append(ids, item.KnowledgeBaseID)
|
||||||
|
}
|
||||||
|
return ids
|
||||||
|
}
|
||||||
|
|
||||||
|
func distinctDocumentIDs(hits []response.KnowledgeSearchResult) []int64 {
|
||||||
|
seen := make(map[int64]struct{})
|
||||||
|
items := make([]int64, 0)
|
||||||
|
for _, item := range hits {
|
||||||
|
if item.DocumentID <= 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[item.DocumentID]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[item.DocumentID] = struct{}{}
|
||||||
|
items = append(items, item.DocumentID)
|
||||||
|
}
|
||||||
|
return items
|
||||||
|
}
|
||||||
|
|
||||||
|
func distinctSectionPaths(hits []response.KnowledgeSearchResult) []string {
|
||||||
|
seen := make(map[string]struct{})
|
||||||
|
items := make([]string, 0)
|
||||||
|
for _, item := range hits {
|
||||||
|
sectionPath := item.SectionPath
|
||||||
|
if sectionPath == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[sectionPath]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[sectionPath] = struct{}{}
|
||||||
|
items = append(items, sectionPath)
|
||||||
|
}
|
||||||
|
return items
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildKnowledgeSearchResultKey(item response.KnowledgeSearchResult) string {
|
||||||
|
if item.FaqID > 0 {
|
||||||
|
return fmt.Sprintf("faq:%d|%d", item.FaqID, item.ChunkNo)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%d|%s|%d", item.DocumentID, item.SectionPath, item.ChunkNo)
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildKnowledgeCitationKey(item response.KnowledgeCitation) string {
|
||||||
|
if item.FaqID > 0 {
|
||||||
|
return fmt.Sprintf("faq:%d|%d", item.FaqID, item.ChunkNo)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%d|%s|%d", item.DocumentID, item.SectionPath, item.ChunkNo)
|
||||||
|
}
|
||||||
|
|
||||||
|
func hasHitKey(items map[string]struct{}, key string) bool {
|
||||||
|
_, ok := items[key]
|
||||||
|
return ok
|
||||||
|
}
|
||||||
163
internal/ai/rag/retrieve_search.go
Normal file
163
internal/ai/rag/retrieve_search.go
Normal file
@@ -0,0 +1,163 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai"
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *retrieve) searchKnowledgeBaseVectors(ctx context.Context, req RetrieveRequest, knowledgeBases []models.KnowledgeBase) ([]vectordb.SearchResult, *RetrieveTrace, error) {
|
||||||
|
trace := &RetrieveTrace{}
|
||||||
|
|
||||||
|
embeddingStartedAt := time.Now()
|
||||||
|
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, req.Query)
|
||||||
|
trace.EmbeddingMs = time.Since(embeddingStartedAt).Milliseconds()
|
||||||
|
if err != nil {
|
||||||
|
return nil, trace, fmt.Errorf("failed to generate query embedding: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
collectionName := knowledgeCollectionName
|
||||||
|
provider := vectordb.GetProvider()
|
||||||
|
if provider == nil {
|
||||||
|
return nil, trace, fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
|
||||||
|
searchResults := make([]vectordb.SearchResult, 0)
|
||||||
|
vectorSearchStartedAt := time.Now()
|
||||||
|
for _, knowledgeBase := range knowledgeBases {
|
||||||
|
topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(req, &knowledgeBase)
|
||||||
|
kbResults, searchErr := provider.Search(ctx, &vectordb.SearchRequest{
|
||||||
|
CollectionName: collectionName,
|
||||||
|
Vector: embeddingResult.Vector,
|
||||||
|
TopK: topK,
|
||||||
|
ScoreThreshold: scoreThreshold,
|
||||||
|
Filter: &vectordb.SearchFilter{
|
||||||
|
KnowledgeBaseIDs: []int64{knowledgeBase.ID},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if searchErr != nil {
|
||||||
|
slog.Error("Failed to search vectors",
|
||||||
|
"knowledge_base_id", knowledgeBase.ID,
|
||||||
|
"error", searchErr)
|
||||||
|
trace.VectorSearchMs = time.Since(vectorSearchStartedAt).Milliseconds()
|
||||||
|
return nil, trace, fmt.Errorf("failed to search vectors: %w", searchErr)
|
||||||
|
}
|
||||||
|
if len(kbResults) == 0 && scoreThreshold > 0 {
|
||||||
|
s.logEmptySearchDiagnostics(ctx, provider, collectionName, embeddingResult.Vector, topK, scoreThreshold, []int64{knowledgeBase.ID}, req)
|
||||||
|
}
|
||||||
|
searchResults = append(searchResults, kbResults...)
|
||||||
|
}
|
||||||
|
trace.VectorSearchMs = time.Since(vectorSearchStartedAt).Milliseconds()
|
||||||
|
|
||||||
|
if len(searchResults) > 0 {
|
||||||
|
sort.SliceStable(searchResults, func(i, j int) bool {
|
||||||
|
if searchResults[i].Score == searchResults[j].Score {
|
||||||
|
return searchResults[i].ID < searchResults[j].ID
|
||||||
|
}
|
||||||
|
return searchResults[i].Score > searchResults[j].Score
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return searchResults, trace, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *retrieve) hydrateRetrieveResults(searchResults []vectordb.SearchResult) ([]RetrieveResult, int64) {
|
||||||
|
if len(searchResults) == 0 {
|
||||||
|
return nil, 0
|
||||||
|
}
|
||||||
|
|
||||||
|
hydrateStartedAt := time.Now()
|
||||||
|
results := make([]RetrieveResult, 0, len(searchResults))
|
||||||
|
vectorIDs := make([]string, 0, len(searchResults))
|
||||||
|
for _, sr := range searchResults {
|
||||||
|
if strings.TrimSpace(sr.ID) == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
vectorIDs = append(vectorIDs, sr.ID)
|
||||||
|
}
|
||||||
|
chunks := repositories.KnowledgeChunkRepository.FindByVectorIDs(sqls.DB(), vectorIDs)
|
||||||
|
chunkByVectorID := make(map[string]*models.KnowledgeChunk, len(chunks))
|
||||||
|
documentIDs := make([]int64, 0)
|
||||||
|
faqIDs := make([]int64, 0)
|
||||||
|
documentSeen := make(map[int64]struct{})
|
||||||
|
faqSeen := make(map[int64]struct{})
|
||||||
|
for i := range chunks {
|
||||||
|
chunk := &chunks[i]
|
||||||
|
chunkByVectorID[chunk.VectorID] = chunk
|
||||||
|
if chunk.DocumentID > 0 {
|
||||||
|
if _, ok := documentSeen[chunk.DocumentID]; !ok {
|
||||||
|
documentSeen[chunk.DocumentID] = struct{}{}
|
||||||
|
documentIDs = append(documentIDs, chunk.DocumentID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if chunk.FaqID > 0 {
|
||||||
|
if _, ok := faqSeen[chunk.FaqID]; !ok {
|
||||||
|
faqSeen[chunk.FaqID] = struct{}{}
|
||||||
|
faqIDs = append(faqIDs, chunk.FaqID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
documents := repositories.KnowledgeDocumentRepository.FindByIDs(sqls.DB(), documentIDs)
|
||||||
|
documentByID := make(map[int64]*models.KnowledgeDocument, len(documents))
|
||||||
|
for i := range documents {
|
||||||
|
document := &documents[i]
|
||||||
|
documentByID[document.ID] = document
|
||||||
|
}
|
||||||
|
faqs := repositories.KnowledgeFAQRepository.FindByIDs(sqls.DB(), faqIDs)
|
||||||
|
faqByID := make(map[int64]*models.KnowledgeFAQ, len(faqs))
|
||||||
|
for i := range faqs {
|
||||||
|
faq := &faqs[i]
|
||||||
|
faqByID[faq.ID] = faq
|
||||||
|
}
|
||||||
|
for _, sr := range searchResults {
|
||||||
|
chunk := chunkByVectorID[sr.ID]
|
||||||
|
if chunk == nil || chunk.Status != enums.StatusOk {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
documentTitle := ""
|
||||||
|
faqQuestion := ""
|
||||||
|
if chunk.DocumentID > 0 {
|
||||||
|
document := documentByID[chunk.DocumentID]
|
||||||
|
if document == nil || document.Status != enums.StatusOk {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
documentTitle = document.Title
|
||||||
|
}
|
||||||
|
if chunk.FaqID > 0 {
|
||||||
|
faq := faqByID[chunk.FaqID]
|
||||||
|
if faq == nil || faq.Status != enums.StatusOk {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
faqQuestion = faq.Question
|
||||||
|
}
|
||||||
|
|
||||||
|
results = append(results, RetrieveResult{
|
||||||
|
KnowledgeBaseID: chunk.KnowledgeBaseID,
|
||||||
|
ChunkID: chunk.ID,
|
||||||
|
DocumentID: chunk.DocumentID,
|
||||||
|
DocumentTitle: documentTitle,
|
||||||
|
FaqID: chunk.FaqID,
|
||||||
|
FaqQuestion: faqQuestion,
|
||||||
|
ChunkNo: chunk.ChunkNo,
|
||||||
|
Title: chunk.Title,
|
||||||
|
SectionPath: chunk.SectionPath,
|
||||||
|
Content: chunk.Content,
|
||||||
|
Score: sr.Score,
|
||||||
|
ChunkType: extractChunkType(sr.Payload),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return results, time.Since(hydrateStartedAt).Milliseconds()
|
||||||
|
}
|
||||||
49
internal/ai/rag/retrieve_test.go
Normal file
49
internal/ai/rag/retrieve_test.go
Normal file
@@ -0,0 +1,49 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestResolveKnowledgeBaseSearchOptionsUsesKnowledgeBaseDefaults(t *testing.T) {
|
||||||
|
topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(RetrieveRequest{}, &models.KnowledgeBase{
|
||||||
|
DefaultTopK: 6,
|
||||||
|
DefaultScoreThreshold: 0.42,
|
||||||
|
})
|
||||||
|
|
||||||
|
if topK != 6 {
|
||||||
|
t.Fatalf("expected topK 6, got %d", topK)
|
||||||
|
}
|
||||||
|
if scoreThreshold != float32(0.42) {
|
||||||
|
t.Fatalf("expected score threshold 0.42, got %v", scoreThreshold)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveKnowledgeBaseSearchOptionsRequestOverridesKnowledgeBaseDefaults(t *testing.T) {
|
||||||
|
topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(RetrieveRequest{
|
||||||
|
TopK: 9,
|
||||||
|
ScoreThreshold: 0.55,
|
||||||
|
}, &models.KnowledgeBase{
|
||||||
|
DefaultTopK: 6,
|
||||||
|
DefaultScoreThreshold: 0.42,
|
||||||
|
})
|
||||||
|
|
||||||
|
if topK != 9 {
|
||||||
|
t.Fatalf("expected request topK 9, got %d", topK)
|
||||||
|
}
|
||||||
|
if scoreThreshold != float32(0.55) {
|
||||||
|
t.Fatalf("expected request score threshold 0.55, got %v", scoreThreshold)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveKnowledgeBaseSearchOptionsUsesSystemDefaults(t *testing.T) {
|
||||||
|
topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(RetrieveRequest{}, nil)
|
||||||
|
|
||||||
|
if topK != 8 {
|
||||||
|
t.Fatalf("expected fallback topK 8, got %d", topK)
|
||||||
|
}
|
||||||
|
if scoreThreshold != float32(0.3) {
|
||||||
|
t.Fatalf("expected fallback score threshold 0.3, got %v", scoreThreshold)
|
||||||
|
}
|
||||||
|
}
|
||||||
48
internal/ai/rag/types.go
Normal file
48
internal/ai/rag/types.go
Normal file
@@ -0,0 +1,48 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
type RetrieveRequest struct {
|
||||||
|
KnowledgeBaseIDs []int64
|
||||||
|
Query string
|
||||||
|
TopK int
|
||||||
|
ScoreThreshold float64
|
||||||
|
}
|
||||||
|
|
||||||
|
type RetrieveResult struct {
|
||||||
|
KnowledgeBaseID int64 `json:"knowledgeBaseId"`
|
||||||
|
ChunkID int64 `json:"chunkId"`
|
||||||
|
DocumentID int64 `json:"documentId"`
|
||||||
|
DocumentTitle string `json:"documentTitle"`
|
||||||
|
FaqID int64 `json:"faqId"`
|
||||||
|
FaqQuestion string `json:"faqQuestion"`
|
||||||
|
ChunkNo int `json:"chunkNo"`
|
||||||
|
Title string `json:"title"`
|
||||||
|
SectionPath string `json:"sectionPath"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
Score float32 `json:"score"`
|
||||||
|
ChunkType string `json:"chunkType"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type RerankRequest struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Query string `json:"query"`
|
||||||
|
Documents []string `json:"documents"`
|
||||||
|
TopN int `json:"top_n"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type RerankResponse struct {
|
||||||
|
Results []struct {
|
||||||
|
Document string `json:"document"`
|
||||||
|
Index int `json:"index"`
|
||||||
|
RelevanceScore float64 `json:"relevance_score"`
|
||||||
|
} `json:"results"`
|
||||||
|
Meta struct {
|
||||||
|
APIVersion struct {
|
||||||
|
Version string `json:"version"`
|
||||||
|
} `json:"api_version"`
|
||||||
|
} `json:"meta"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type RerankResult struct {
|
||||||
|
Index int `json:"index"`
|
||||||
|
RelevanceScore float64 `json:"relevanceScore"`
|
||||||
|
}
|
||||||
112
internal/ai/rag/utils.go
Normal file
112
internal/ai/rag/utils.go
Normal file
@@ -0,0 +1,112 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/yuin/goldmark"
|
||||||
|
"github.com/yuin/goldmark/extension"
|
||||||
|
"github.com/yuin/goldmark/parser"
|
||||||
|
"github.com/yuin/goldmark/renderer/html"
|
||||||
|
htmlparser "golang.org/x/net/html"
|
||||||
|
)
|
||||||
|
|
||||||
|
var plainTextMarkdown = goldmark.New(
|
||||||
|
goldmark.WithExtensions(extension.GFM),
|
||||||
|
goldmark.WithParserOptions(
|
||||||
|
parser.WithAutoHeadingID(),
|
||||||
|
),
|
||||||
|
goldmark.WithRendererOptions(
|
||||||
|
html.WithHardWraps(),
|
||||||
|
html.WithXHTML(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
func ExtractPlainText(content string, contentType enums.KnowledgeDocumentContentType) string {
|
||||||
|
switch contentType {
|
||||||
|
case enums.KnowledgeDocumentContentTypeMarkdown:
|
||||||
|
return ExtractPlainTextFromMarkdown(content)
|
||||||
|
case enums.KnowledgeDocumentContentTypeHTML:
|
||||||
|
return ExtractPlainTextFromHTML(content)
|
||||||
|
default:
|
||||||
|
return normalizeWhitespace(content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func ExtractPlainTextFromMarkdown(content string) string {
|
||||||
|
content = strings.TrimSpace(content)
|
||||||
|
if content == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var buf strings.Builder
|
||||||
|
if err := plainTextMarkdown.Convert([]byte(content), &buf); err != nil {
|
||||||
|
return normalizeWhitespace(content)
|
||||||
|
}
|
||||||
|
return ExtractPlainTextFromHTML(buf.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
func ExtractPlainTextFromHTML(content string) string {
|
||||||
|
content = strings.TrimSpace(content)
|
||||||
|
if content == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var builder strings.Builder
|
||||||
|
parent := &htmlparser.Node{
|
||||||
|
Type: htmlparser.ElementNode,
|
||||||
|
Data: "div",
|
||||||
|
}
|
||||||
|
nodes, err := htmlparser.ParseFragment(strings.NewReader(content), parent)
|
||||||
|
if err == nil {
|
||||||
|
for _, node := range nodes {
|
||||||
|
writeHTMLNodeText(&builder, node)
|
||||||
|
}
|
||||||
|
return normalizeWhitespace(builder.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
// 兜底:部分输入在 ParseFragment 下会失败(例如不符合 fragment 规则或上下文不匹配)。
|
||||||
|
// 这里用完整 HTML 解析保证可用性。
|
||||||
|
doc, err := htmlparser.Parse(strings.NewReader("<div>" + content + "</div>"))
|
||||||
|
if err != nil {
|
||||||
|
return normalizeWhitespace(content)
|
||||||
|
}
|
||||||
|
writeHTMLNodeText(&builder, doc)
|
||||||
|
return normalizeWhitespace(builder.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeHTMLNodeText(builder *strings.Builder, node *htmlparser.Node) {
|
||||||
|
if node == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
switch node.Type {
|
||||||
|
case htmlparser.TextNode:
|
||||||
|
builder.WriteString(node.Data)
|
||||||
|
case htmlparser.ElementNode:
|
||||||
|
if shouldSeparateHTMLText(node.Data) {
|
||||||
|
builder.WriteByte(' ')
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for child := node.FirstChild; child != nil; child = child.NextSibling {
|
||||||
|
writeHTMLNodeText(builder, child)
|
||||||
|
}
|
||||||
|
|
||||||
|
if node.Type == htmlparser.ElementNode && shouldSeparateHTMLText(node.Data) {
|
||||||
|
builder.WriteByte(' ')
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func shouldSeparateHTMLText(tag string) bool {
|
||||||
|
switch tag {
|
||||||
|
case "p", "div", "br", "li", "ul", "ol", "blockquote", "pre", "table", "tr", "td", "th", "h1", "h2", "h3", "h4", "h5", "h6":
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeWhitespace(content string) string {
|
||||||
|
return strings.Join(strings.Fields(strings.TrimSpace(content)), " ")
|
||||||
|
}
|
||||||
43
internal/ai/rag/utils_test.go
Normal file
43
internal/ai/rag/utils_test.go
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
package rag
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/ai/rag/vectordb"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestExtractPlainTextFromHTMLSeparatesBlockContent(t *testing.T) {
|
||||||
|
got := ExtractPlainTextFromHTML("<div>Hello</div><div>World</div><p>Again</p>")
|
||||||
|
want := "Hello World Again"
|
||||||
|
if got != want {
|
||||||
|
t.Fatalf("expected %q, got %q", want, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExtractPlainTextMarkdownUsesGoldmark(t *testing.T) {
|
||||||
|
got := ExtractPlainText("# Title\n\n- one\n- two", enums.KnowledgeDocumentContentTypeMarkdown)
|
||||||
|
want := "Title one two"
|
||||||
|
if got != want {
|
||||||
|
t.Fatalf("expected %q, got %q", want, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestChunkPayloadFromMapSupportsTypedConversion(t *testing.T) {
|
||||||
|
got := vectordb.ChunkPayloadFromMap(map[string]any{
|
||||||
|
"knowledge_base_id": "1",
|
||||||
|
"document_id": "123",
|
||||||
|
"document_title": "Doc",
|
||||||
|
"chunk_no": "2",
|
||||||
|
"chunk_type": "text",
|
||||||
|
"section_path": "A > B",
|
||||||
|
"title": "hello",
|
||||||
|
"content": "world",
|
||||||
|
"provider": "structured",
|
||||||
|
})
|
||||||
|
if got.KnowledgeBaseID != 1 || got.DocumentID != 123 || got.ChunkNo != 2 {
|
||||||
|
t.Fatalf("unexpected numeric conversion result: %+v", got)
|
||||||
|
}
|
||||||
|
if got.DocumentTitle != "Doc" || got.SectionPath != "A > B" || got.Provider != "structured" {
|
||||||
|
t.Fatalf("unexpected string conversion result: %+v", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
489
internal/ai/rag/vectordb/lancedb.go
Normal file
489
internal/ai/rag/vectordb/lancedb.go
Normal file
@@ -0,0 +1,489 @@
|
|||||||
|
//go:build lancedb
|
||||||
|
|
||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"math"
|
||||||
|
"os"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
|
||||||
|
"github.com/apache/arrow/go/v17/arrow"
|
||||||
|
"github.com/apache/arrow/go/v17/arrow/array"
|
||||||
|
"github.com/apache/arrow/go/v17/arrow/memory"
|
||||||
|
"github.com/lancedb/lancedb-go/pkg/contracts"
|
||||||
|
"github.com/lancedb/lancedb-go/pkg/lancedb"
|
||||||
|
)
|
||||||
|
|
||||||
|
const lanceDBVectorColumn = "vector"
|
||||||
|
|
||||||
|
type LanceDBProvider struct {
|
||||||
|
conn contracts.IConnection
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewLanceDBProvider(cfg *config.LanceDBVectorDBConfig) (Provider, error) {
|
||||||
|
if cfg == nil {
|
||||||
|
return nil, fmt.Errorf("lancedb config is nil")
|
||||||
|
}
|
||||||
|
path := strings.TrimSpace(cfg.Path)
|
||||||
|
if path == "" {
|
||||||
|
path = "data/lancedb"
|
||||||
|
}
|
||||||
|
if err := os.MkdirAll(path, 0o755); err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create lancedb directory %s: %w", path, err)
|
||||||
|
}
|
||||||
|
conn, err := lancedb.Connect(context.Background(), path, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return &LanceDBProvider{conn: conn}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) Close() error {
|
||||||
|
if p.conn == nil || p.conn.IsClosed() {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return p.conn.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) CreateCollection(ctx context.Context, name string, dimension int) error {
|
||||||
|
if dimension <= 0 {
|
||||||
|
return fmt.Errorf("invalid lancedb vector dimension: %d", dimension)
|
||||||
|
}
|
||||||
|
if err := p.ensureOpen(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
schema, err := newLanceDBSchema(dimension)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
table, err := p.conn.CreateTable(ctx, name, schema)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to create lancedb table %s: %w", name, err)
|
||||||
|
}
|
||||||
|
return table.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) DeleteCollection(ctx context.Context, name string) error {
|
||||||
|
if err := p.ensureOpen(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := p.conn.DropTable(ctx, name); err != nil {
|
||||||
|
return fmt.Errorf("failed to delete lancedb table %s: %w", name, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) GetCollection(ctx context.Context, name string) (*CollectionInfo, error) {
|
||||||
|
table, err := p.openTable(ctx, name)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer table.Close()
|
||||||
|
|
||||||
|
schema, err := table.Schema(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to get lancedb table schema %s: %w", name, err)
|
||||||
|
}
|
||||||
|
count, err := table.Count(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to count lancedb table %s: %w", name, err)
|
||||||
|
}
|
||||||
|
return &CollectionInfo{
|
||||||
|
Name: name,
|
||||||
|
Dimension: lanceDBVectorDimension(schema),
|
||||||
|
PointCount: int(count),
|
||||||
|
Status: "ok",
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) ListCollections(ctx context.Context) ([]string, error) {
|
||||||
|
if err := p.ensureOpen(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
names, err := p.conn.TableNames(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to list lancedb tables: %w", err)
|
||||||
|
}
|
||||||
|
return names, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error {
|
||||||
|
if len(vectors) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
table, err := p.openTable(ctx, collectionName)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer table.Close()
|
||||||
|
|
||||||
|
ids := make([]string, 0, len(vectors))
|
||||||
|
for _, vector := range vectors {
|
||||||
|
if strings.TrimSpace(vector.ID) != "" {
|
||||||
|
ids = append(ids, vector.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(ids) > 0 {
|
||||||
|
if err := table.Delete(ctx, lanceDBStringInFilter("id", ids)); err != nil {
|
||||||
|
return fmt.Errorf("failed to delete existing lancedb vectors from %s: %w", collectionName, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
record, release, err := newLanceDBVectorRecord(vectors)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer release()
|
||||||
|
|
||||||
|
if err := table.AddRecords(ctx, []arrow.Record{record}, nil); err != nil {
|
||||||
|
return fmt.Errorf("failed to add lancedb vectors to %s: %w", collectionName, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) DeleteVectors(ctx context.Context, collectionName string, ids []string) error {
|
||||||
|
if len(ids) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
table, err := p.openTable(ctx, collectionName)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer table.Close()
|
||||||
|
|
||||||
|
if err := table.Delete(ctx, lanceDBStringInFilter("id", ids)); err != nil {
|
||||||
|
return fmt.Errorf("failed to delete lancedb vectors from %s: %w", collectionName, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) {
|
||||||
|
table, err := p.openTable(ctx, req.CollectionName)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer table.Close()
|
||||||
|
|
||||||
|
filter := lanceDBSearchFilter(req.Filter)
|
||||||
|
var rows []map[string]interface{}
|
||||||
|
if filter == "" {
|
||||||
|
rows, err = table.VectorSearch(ctx, lanceDBVectorColumn, req.Vector, req.TopK)
|
||||||
|
} else {
|
||||||
|
rows, err = table.VectorSearchWithFilter(ctx, lanceDBVectorColumn, req.Vector, req.TopK, filter)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to search lancedb table %s: %w", req.CollectionName, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
results := make([]SearchResult, 0, len(rows))
|
||||||
|
for _, row := range rows {
|
||||||
|
score := lanceDBScoreFromRow(row)
|
||||||
|
if req.ScoreThreshold > 0 && score < req.ScoreThreshold {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
results = append(results, SearchResult{
|
||||||
|
ID: valueToString(row["id"]),
|
||||||
|
Score: score,
|
||||||
|
Payload: lanceDBPayloadFromRow(row),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return results, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) ensureOpen() error {
|
||||||
|
if p == nil || p.conn == nil || p.conn.IsClosed() {
|
||||||
|
return fmt.Errorf("lancedb provider is closed")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *LanceDBProvider) openTable(ctx context.Context, name string) (contracts.ITable, error) {
|
||||||
|
if err := p.ensureOpen(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
table, err := p.conn.OpenTable(ctx, name)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to open lancedb table %s: %w", name, err)
|
||||||
|
}
|
||||||
|
return table, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func newLanceDBSchema(dimension int) (contracts.ISchema, error) {
|
||||||
|
schema := arrow.NewSchema([]arrow.Field{
|
||||||
|
{Name: "id", Type: arrow.BinaryTypes.String, Nullable: false},
|
||||||
|
{Name: lanceDBVectorColumn, Type: arrow.FixedSizeListOf(int32(dimension), arrow.PrimitiveTypes.Float32), Nullable: false},
|
||||||
|
{Name: "knowledge_base_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false},
|
||||||
|
{Name: "document_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false},
|
||||||
|
{Name: "document_title", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "faq_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false},
|
||||||
|
{Name: "faq_question", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "chunk_no", Type: arrow.PrimitiveTypes.Int32, Nullable: false},
|
||||||
|
{Name: "chunk_type", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "section_path", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "title", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "content", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "provider", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
}, nil)
|
||||||
|
return lancedb.NewSchema(schema)
|
||||||
|
}
|
||||||
|
|
||||||
|
func newLanceDBVectorRecord(vectors []Vector) (arrow.Record, func(), error) {
|
||||||
|
dimension := 0
|
||||||
|
for _, item := range vectors {
|
||||||
|
if len(item.Vector) > 0 {
|
||||||
|
dimension = len(item.Vector)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if dimension <= 0 {
|
||||||
|
return nil, nil, fmt.Errorf("lancedb vector dimension is empty")
|
||||||
|
}
|
||||||
|
for _, item := range vectors {
|
||||||
|
if len(item.Vector) != dimension {
|
||||||
|
return nil, nil, fmt.Errorf("inconsistent lancedb vector dimension for %s: got %d, want %d", item.ID, len(item.Vector), dimension)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pool := memory.NewGoAllocator()
|
||||||
|
idBuilder := array.NewStringBuilder(pool)
|
||||||
|
kbIDBuilder := array.NewInt64Builder(pool)
|
||||||
|
documentIDBuilder := array.NewInt64Builder(pool)
|
||||||
|
documentTitleBuilder := array.NewStringBuilder(pool)
|
||||||
|
faqIDBuilder := array.NewInt64Builder(pool)
|
||||||
|
faqQuestionBuilder := array.NewStringBuilder(pool)
|
||||||
|
chunkNoBuilder := array.NewInt32Builder(pool)
|
||||||
|
chunkTypeBuilder := array.NewStringBuilder(pool)
|
||||||
|
sectionPathBuilder := array.NewStringBuilder(pool)
|
||||||
|
titleBuilder := array.NewStringBuilder(pool)
|
||||||
|
contentBuilder := array.NewStringBuilder(pool)
|
||||||
|
providerBuilder := array.NewStringBuilder(pool)
|
||||||
|
vectorBuilder := array.NewFloat32Builder(pool)
|
||||||
|
|
||||||
|
for _, item := range vectors {
|
||||||
|
payload := item.Payload
|
||||||
|
idBuilder.Append(item.ID)
|
||||||
|
vectorBuilder.AppendValues(item.Vector, nil)
|
||||||
|
kbIDBuilder.Append(payload.KnowledgeBaseID)
|
||||||
|
documentIDBuilder.Append(payload.DocumentID)
|
||||||
|
documentTitleBuilder.Append(payload.DocumentTitle)
|
||||||
|
faqIDBuilder.Append(payload.FaqID)
|
||||||
|
faqQuestionBuilder.Append(payload.FaqQuestion)
|
||||||
|
chunkNoBuilder.Append(int32(payload.ChunkNo))
|
||||||
|
chunkTypeBuilder.Append(payload.ChunkType)
|
||||||
|
sectionPathBuilder.Append(payload.SectionPath)
|
||||||
|
titleBuilder.Append(payload.Title)
|
||||||
|
contentBuilder.Append(payload.Content)
|
||||||
|
providerBuilder.Append(payload.Provider)
|
||||||
|
}
|
||||||
|
|
||||||
|
idArray := idBuilder.NewArray()
|
||||||
|
vectorValues := vectorBuilder.NewArray()
|
||||||
|
kbIDArray := kbIDBuilder.NewArray()
|
||||||
|
documentIDArray := documentIDBuilder.NewArray()
|
||||||
|
documentTitleArray := documentTitleBuilder.NewArray()
|
||||||
|
faqIDArray := faqIDBuilder.NewArray()
|
||||||
|
faqQuestionArray := faqQuestionBuilder.NewArray()
|
||||||
|
chunkNoArray := chunkNoBuilder.NewArray()
|
||||||
|
chunkTypeArray := chunkTypeBuilder.NewArray()
|
||||||
|
sectionPathArray := sectionPathBuilder.NewArray()
|
||||||
|
titleArray := titleBuilder.NewArray()
|
||||||
|
contentArray := contentBuilder.NewArray()
|
||||||
|
providerArray := providerBuilder.NewArray()
|
||||||
|
|
||||||
|
vectorType := arrow.FixedSizeListOf(int32(dimension), arrow.PrimitiveTypes.Float32)
|
||||||
|
vectorArray := array.NewFixedSizeListData(
|
||||||
|
array.NewData(vectorType, len(vectors), []*memory.Buffer{nil}, []arrow.ArrayData{vectorValues.Data()}, 0, 0),
|
||||||
|
)
|
||||||
|
schema := arrow.NewSchema([]arrow.Field{
|
||||||
|
{Name: "id", Type: arrow.BinaryTypes.String, Nullable: false},
|
||||||
|
{Name: lanceDBVectorColumn, Type: vectorType, Nullable: false},
|
||||||
|
{Name: "knowledge_base_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false},
|
||||||
|
{Name: "document_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false},
|
||||||
|
{Name: "document_title", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "faq_id", Type: arrow.PrimitiveTypes.Int64, Nullable: false},
|
||||||
|
{Name: "faq_question", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "chunk_no", Type: arrow.PrimitiveTypes.Int32, Nullable: false},
|
||||||
|
{Name: "chunk_type", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "section_path", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "title", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "content", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
{Name: "provider", Type: arrow.BinaryTypes.String, Nullable: true},
|
||||||
|
}, nil)
|
||||||
|
columns := []arrow.Array{
|
||||||
|
idArray,
|
||||||
|
vectorArray,
|
||||||
|
kbIDArray,
|
||||||
|
documentIDArray,
|
||||||
|
documentTitleArray,
|
||||||
|
faqIDArray,
|
||||||
|
faqQuestionArray,
|
||||||
|
chunkNoArray,
|
||||||
|
chunkTypeArray,
|
||||||
|
sectionPathArray,
|
||||||
|
titleArray,
|
||||||
|
contentArray,
|
||||||
|
providerArray,
|
||||||
|
}
|
||||||
|
record := array.NewRecord(schema, columns, int64(len(vectors)))
|
||||||
|
release := func() {
|
||||||
|
record.Release()
|
||||||
|
for _, column := range columns {
|
||||||
|
column.Release()
|
||||||
|
}
|
||||||
|
vectorValues.Release()
|
||||||
|
}
|
||||||
|
return record, release, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func lanceDBVectorDimension(schema *arrow.Schema) int {
|
||||||
|
if schema == nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
for i := 0; i < schema.NumFields(); i++ {
|
||||||
|
field := schema.Field(i)
|
||||||
|
if field.Name != lanceDBVectorColumn {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
listType, ok := field.Type.(*arrow.FixedSizeListType)
|
||||||
|
if !ok {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return int(listType.Len())
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func lanceDBSearchFilter(filter *SearchFilter) string {
|
||||||
|
if filter == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
parts := make([]string, 0, 2)
|
||||||
|
if len(filter.KnowledgeBaseIDs) > 0 {
|
||||||
|
parts = append(parts, lanceDBIntInFilter("knowledge_base_id", filter.KnowledgeBaseIDs))
|
||||||
|
}
|
||||||
|
if len(filter.DocumentIDs) > 0 {
|
||||||
|
parts = append(parts, lanceDBIntInFilter("document_id", filter.DocumentIDs))
|
||||||
|
}
|
||||||
|
return strings.Join(parts, " AND ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func lanceDBIntInFilter(column string, values []int64) string {
|
||||||
|
items := make([]string, 0, len(values))
|
||||||
|
for _, value := range values {
|
||||||
|
items = append(items, strconv.FormatInt(value, 10))
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%s IN (%s)", column, strings.Join(items, ","))
|
||||||
|
}
|
||||||
|
|
||||||
|
func lanceDBStringInFilter(column string, values []string) string {
|
||||||
|
items := make([]string, 0, len(values))
|
||||||
|
for _, value := range values {
|
||||||
|
items = append(items, "'"+strings.ReplaceAll(value, "'", "''")+"'")
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%s IN (%s)", column, strings.Join(items, ","))
|
||||||
|
}
|
||||||
|
|
||||||
|
func lanceDBScoreFromRow(row map[string]interface{}) float32 {
|
||||||
|
for _, key := range []string{"_distance", "distance"} {
|
||||||
|
if value, ok := row[key]; ok {
|
||||||
|
distance := valueToFloat64(value)
|
||||||
|
if math.IsNaN(distance) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
score := 1 - distance
|
||||||
|
if score < 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
if score > 1 {
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
return float32(score)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, key := range []string{"_score", "score"} {
|
||||||
|
if value, ok := row[key]; ok {
|
||||||
|
score := valueToFloat64(value)
|
||||||
|
if !math.IsNaN(score) {
|
||||||
|
return float32(score)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func lanceDBPayloadFromRow(row map[string]interface{}) ChunkPayload {
|
||||||
|
return ChunkPayload{
|
||||||
|
KnowledgeBaseID: valueToInt64(row["knowledge_base_id"]),
|
||||||
|
DocumentID: valueToInt64(row["document_id"]),
|
||||||
|
DocumentTitle: valueToString(row["document_title"]),
|
||||||
|
FaqID: valueToInt64(row["faq_id"]),
|
||||||
|
FaqQuestion: valueToString(row["faq_question"]),
|
||||||
|
ChunkNo: int(valueToInt64(row["chunk_no"])),
|
||||||
|
ChunkType: valueToString(row["chunk_type"]),
|
||||||
|
SectionPath: valueToString(row["section_path"]),
|
||||||
|
Title: valueToString(row["title"]),
|
||||||
|
Content: valueToString(row["content"]),
|
||||||
|
Provider: valueToString(row["provider"]),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func valueToString(value interface{}) string {
|
||||||
|
switch v := value.(type) {
|
||||||
|
case nil:
|
||||||
|
return ""
|
||||||
|
case string:
|
||||||
|
return v
|
||||||
|
case []byte:
|
||||||
|
return string(v)
|
||||||
|
default:
|
||||||
|
return fmt.Sprint(value)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func valueToInt64(value interface{}) int64 {
|
||||||
|
switch v := value.(type) {
|
||||||
|
case int:
|
||||||
|
return int64(v)
|
||||||
|
case int32:
|
||||||
|
return int64(v)
|
||||||
|
case int64:
|
||||||
|
return v
|
||||||
|
case uint64:
|
||||||
|
return int64(v)
|
||||||
|
case float32:
|
||||||
|
return int64(v)
|
||||||
|
case float64:
|
||||||
|
return int64(v)
|
||||||
|
case string:
|
||||||
|
ret, _ := strconv.ParseInt(v, 10, 64)
|
||||||
|
return ret
|
||||||
|
default:
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func valueToFloat64(value interface{}) float64 {
|
||||||
|
switch v := value.(type) {
|
||||||
|
case float32:
|
||||||
|
return float64(v)
|
||||||
|
case float64:
|
||||||
|
return v
|
||||||
|
case int:
|
||||||
|
return float64(v)
|
||||||
|
case int32:
|
||||||
|
return float64(v)
|
||||||
|
case int64:
|
||||||
|
return float64(v)
|
||||||
|
case string:
|
||||||
|
ret, err := strconv.ParseFloat(v, 64)
|
||||||
|
if err == nil {
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return math.NaN()
|
||||||
|
}
|
||||||
13
internal/ai/rag/vectordb/lancedb_stub.go
Normal file
13
internal/ai/rag/vectordb/lancedb_stub.go
Normal file
@@ -0,0 +1,13 @@
|
|||||||
|
//go:build !lancedb
|
||||||
|
|
||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
func NewLanceDBProvider(_ *config.LanceDBVectorDBConfig) (Provider, error) {
|
||||||
|
return nil, fmt.Errorf("LanceDB provider is not built. Rebuild with -tags lancedb and configure LanceDB native libraries")
|
||||||
|
}
|
||||||
102
internal/ai/rag/vectordb/lancedb_test.go
Normal file
102
internal/ai/rag/vectordb/lancedb_test.go
Normal file
@@ -0,0 +1,102 @@
|
|||||||
|
//go:build lancedb
|
||||||
|
|
||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestLanceDBProviderVectorLifecycle(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
provider, err := NewLanceDBProvider(&config.LanceDBVectorDBConfig{Path: t.TempDir()})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewLanceDBProvider() error = %v", err)
|
||||||
|
}
|
||||||
|
defer provider.Close()
|
||||||
|
|
||||||
|
const collectionName = "knowledge_chunks"
|
||||||
|
if err := provider.CreateCollection(ctx, collectionName, 3); err != nil {
|
||||||
|
t.Fatalf("CreateCollection() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
vectors := []Vector{
|
||||||
|
{
|
||||||
|
ID: "a",
|
||||||
|
Vector: []float32{1, 0, 0},
|
||||||
|
Payload: ChunkPayload{
|
||||||
|
KnowledgeBaseID: 10,
|
||||||
|
DocumentID: 100,
|
||||||
|
Title: "A",
|
||||||
|
Content: "alpha",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
ID: "b",
|
||||||
|
Vector: []float32{0, 1, 0},
|
||||||
|
Payload: ChunkPayload{
|
||||||
|
KnowledgeBaseID: 20,
|
||||||
|
DocumentID: 200,
|
||||||
|
Title: "B",
|
||||||
|
Content: "beta",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := provider.UpsertVectors(ctx, collectionName, vectors); err != nil {
|
||||||
|
t.Fatalf("UpsertVectors() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
info, err := provider.GetCollection(ctx, collectionName)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetCollection() error = %v", err)
|
||||||
|
}
|
||||||
|
if info.Dimension != 3 {
|
||||||
|
t.Fatalf("CollectionInfo.Dimension = %d, want 3", info.Dimension)
|
||||||
|
}
|
||||||
|
if info.PointCount != 2 {
|
||||||
|
t.Fatalf("CollectionInfo.PointCount = %d, want 2", info.PointCount)
|
||||||
|
}
|
||||||
|
|
||||||
|
results, err := provider.Search(ctx, &SearchRequest{
|
||||||
|
CollectionName: collectionName,
|
||||||
|
Vector: []float32{1, 0, 0},
|
||||||
|
TopK: 5,
|
||||||
|
ScoreThreshold: 0,
|
||||||
|
Filter: &SearchFilter{
|
||||||
|
KnowledgeBaseIDs: []int64{10},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Search() error = %v", err)
|
||||||
|
}
|
||||||
|
if len(results) != 1 {
|
||||||
|
t.Fatalf("Search() returned %d results, want 1: %#v", len(results), results)
|
||||||
|
}
|
||||||
|
if results[0].ID != "a" {
|
||||||
|
t.Fatalf("Search()[0].ID = %q, want %q", results[0].ID, "a")
|
||||||
|
}
|
||||||
|
if results[0].Payload.KnowledgeBaseID != 10 {
|
||||||
|
t.Fatalf("Search()[0].Payload.KnowledgeBaseID = %d, want 10", results[0].Payload.KnowledgeBaseID)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := provider.DeleteVectors(ctx, collectionName, []string{"a"}); err != nil {
|
||||||
|
t.Fatalf("DeleteVectors() error = %v", err)
|
||||||
|
}
|
||||||
|
results, err = provider.Search(ctx, &SearchRequest{
|
||||||
|
CollectionName: collectionName,
|
||||||
|
Vector: []float32{1, 0, 0},
|
||||||
|
TopK: 5,
|
||||||
|
ScoreThreshold: 0,
|
||||||
|
Filter: &SearchFilter{
|
||||||
|
KnowledgeBaseIDs: []int64{10},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Search() after delete error = %v", err)
|
||||||
|
}
|
||||||
|
if len(results) != 0 {
|
||||||
|
t.Fatalf("Search() after delete returned %d results, want 0: %#v", len(results), results)
|
||||||
|
}
|
||||||
|
}
|
||||||
43
internal/ai/rag/vectordb/payload.go
Normal file
43
internal/ai/rag/vectordb/payload.go
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"github.com/mlogclub/simple/common/structs"
|
||||||
|
"github.com/spf13/cast"
|
||||||
|
)
|
||||||
|
|
||||||
|
type ChunkPayload struct {
|
||||||
|
KnowledgeBaseID int64 `json:"knowledge_base_id"`
|
||||||
|
DocumentID int64 `json:"document_id"`
|
||||||
|
DocumentTitle string `json:"document_title"`
|
||||||
|
FaqID int64 `json:"faq_id"`
|
||||||
|
FaqQuestion string `json:"faq_question"`
|
||||||
|
ChunkNo int `json:"chunk_no"`
|
||||||
|
ChunkType string `json:"chunk_type"`
|
||||||
|
SectionPath string `json:"section_path"`
|
||||||
|
Title string `json:"title"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
Provider string `json:"provider"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p ChunkPayload) ToMap() map[string]any {
|
||||||
|
return structs.StructToMap(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
func ChunkPayloadFromMap(data map[string]any) ChunkPayload {
|
||||||
|
if data == nil {
|
||||||
|
return ChunkPayload{}
|
||||||
|
}
|
||||||
|
return ChunkPayload{
|
||||||
|
KnowledgeBaseID: cast.ToInt64(data["knowledge_base_id"]),
|
||||||
|
DocumentID: cast.ToInt64(data["document_id"]),
|
||||||
|
DocumentTitle: cast.ToString(data["document_title"]),
|
||||||
|
FaqID: cast.ToInt64(data["faq_id"]),
|
||||||
|
FaqQuestion: cast.ToString(data["faq_question"]),
|
||||||
|
ChunkNo: cast.ToInt(data["chunk_no"]),
|
||||||
|
ChunkType: cast.ToString(data["chunk_type"]),
|
||||||
|
SectionPath: cast.ToString(data["section_path"]),
|
||||||
|
Title: cast.ToString(data["title"]),
|
||||||
|
Content: cast.ToString(data["content"]),
|
||||||
|
Provider: cast.ToString(data["provider"]),
|
||||||
|
}
|
||||||
|
}
|
||||||
88
internal/ai/rag/vectordb/provider.go
Normal file
88
internal/ai/rag/vectordb/provider.go
Normal file
@@ -0,0 +1,88 @@
|
|||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
var defaultProvider Provider
|
||||||
|
|
||||||
|
func Init(cfg *config.VectorDBConfig) error {
|
||||||
|
if cfg == nil || cfg.Type == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var err error
|
||||||
|
switch enums.VectorDBType(cfg.Type) {
|
||||||
|
case enums.VectorDBTypeQdrant:
|
||||||
|
defaultProvider, err = NewQdrantProvider(&cfg.Qdrant)
|
||||||
|
case enums.VectorDBTypeLanceDB:
|
||||||
|
defaultProvider, err = NewLanceDBProvider(&cfg.LanceDB)
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unsupported vectordb type: %s", cfg.Type)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func GetProvider() Provider {
|
||||||
|
return defaultProvider
|
||||||
|
}
|
||||||
|
|
||||||
|
func Close() error {
|
||||||
|
if defaultProvider != nil {
|
||||||
|
return defaultProvider.Close()
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func CreateCollection(ctx context.Context, name string, dimension int) error {
|
||||||
|
if defaultProvider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return defaultProvider.CreateCollection(ctx, name, dimension)
|
||||||
|
}
|
||||||
|
|
||||||
|
func DeleteCollection(ctx context.Context, name string) error {
|
||||||
|
if defaultProvider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return defaultProvider.DeleteCollection(ctx, name)
|
||||||
|
}
|
||||||
|
|
||||||
|
func GetCollection(ctx context.Context, name string) (*CollectionInfo, error) {
|
||||||
|
if defaultProvider == nil {
|
||||||
|
return nil, fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return defaultProvider.GetCollection(ctx, name)
|
||||||
|
}
|
||||||
|
|
||||||
|
func ListCollections(ctx context.Context) ([]string, error) {
|
||||||
|
if defaultProvider == nil {
|
||||||
|
return nil, fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return defaultProvider.ListCollections(ctx)
|
||||||
|
}
|
||||||
|
|
||||||
|
func UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error {
|
||||||
|
if defaultProvider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return defaultProvider.UpsertVectors(ctx, collectionName, vectors)
|
||||||
|
}
|
||||||
|
|
||||||
|
func DeleteVectors(ctx context.Context, collectionName string, ids []string) error {
|
||||||
|
if defaultProvider == nil {
|
||||||
|
return fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return defaultProvider.DeleteVectors(ctx, collectionName, ids)
|
||||||
|
}
|
||||||
|
|
||||||
|
func Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) {
|
||||||
|
if defaultProvider == nil {
|
||||||
|
return nil, fmt.Errorf("vectordb provider not initialized")
|
||||||
|
}
|
||||||
|
return defaultProvider.Search(ctx, req)
|
||||||
|
}
|
||||||
25
internal/ai/rag/vectordb/provider_test.go
Normal file
25
internal/ai/rag/vectordb/provider_test.go
Normal file
@@ -0,0 +1,25 @@
|
|||||||
|
//go:build !lancedb
|
||||||
|
|
||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestInitLanceDBWithoutBuildTagReturnsActionableError(t *testing.T) {
|
||||||
|
err := Init(&config.VectorDBConfig{
|
||||||
|
Type: "lancedb",
|
||||||
|
LanceDB: config.LanceDBVectorDBConfig{
|
||||||
|
Path: "data/lancedb",
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("Init(lancedb) error = nil, want actionable build tag error")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "LanceDB provider is not built") {
|
||||||
|
t.Fatalf("Init(lancedb) error = %q, want build tag guidance", err.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
247
internal/ai/rag/vectordb/qdrant.go
Normal file
247
internal/ai/rag/vectordb/qdrant.go
Normal file
@@ -0,0 +1,247 @@
|
|||||||
|
package vectordb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/qdrant/go-client/qdrant"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
type QdrantProvider struct {
|
||||||
|
client *qdrant.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewQdrantProvider(cfg *config.QdrantVectorDBConfig) (*QdrantProvider, error) {
|
||||||
|
if cfg == nil {
|
||||||
|
return nil, fmt.Errorf("vectordb config is nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
host := cfg.Host
|
||||||
|
if host == "" {
|
||||||
|
host = "localhost"
|
||||||
|
}
|
||||||
|
|
||||||
|
port := cfg.GrpcPort
|
||||||
|
if port <= 0 {
|
||||||
|
port = 6334
|
||||||
|
}
|
||||||
|
|
||||||
|
client, err := qdrant.NewClient(&qdrant.Config{
|
||||||
|
Host: host,
|
||||||
|
Port: port,
|
||||||
|
APIKey: cfg.APIKey,
|
||||||
|
UseTLS: cfg.UseTLS,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create qdrant client: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return &QdrantProvider{client: client}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) Close() error {
|
||||||
|
if p.client != nil {
|
||||||
|
return p.client.Close()
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) CreateCollection(ctx context.Context, name string, dimension int) error {
|
||||||
|
err := p.client.CreateCollection(ctx, &qdrant.CreateCollection{
|
||||||
|
CollectionName: name,
|
||||||
|
VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{
|
||||||
|
Size: uint64(dimension),
|
||||||
|
Distance: qdrant.Distance_Cosine,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to create collection %s: %w", name, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) DeleteCollection(ctx context.Context, name string) error {
|
||||||
|
err := p.client.DeleteCollection(ctx, name)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to delete collection %s: %w", name, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) GetCollection(ctx context.Context, name string) (*CollectionInfo, error) {
|
||||||
|
info, err := p.client.GetCollectionInfo(ctx, name)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to get collection %s: %w", name, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
status := info.GetStatus().String()
|
||||||
|
pointCount := int(info.GetPointsCount())
|
||||||
|
|
||||||
|
dimension := 0
|
||||||
|
if info.Config != nil && info.Config.Params != nil {
|
||||||
|
vectorsConfig := info.Config.Params.VectorsConfig
|
||||||
|
if vectorsConfig != nil {
|
||||||
|
params := vectorsConfig.GetParams()
|
||||||
|
if params != nil {
|
||||||
|
dimension = int(params.Size)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return &CollectionInfo{
|
||||||
|
Name: name,
|
||||||
|
Dimension: dimension,
|
||||||
|
PointCount: pointCount,
|
||||||
|
Status: status,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) ListCollections(ctx context.Context) ([]string, error) {
|
||||||
|
collections, err := p.client.ListCollections(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to list collections: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return collections, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error {
|
||||||
|
if len(vectors) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
points := make([]*qdrant.PointStruct, 0, len(vectors))
|
||||||
|
for _, v := range vectors {
|
||||||
|
points = append(points, &qdrant.PointStruct{
|
||||||
|
Id: qdrant.NewID(v.ID),
|
||||||
|
Vectors: qdrant.NewVectors(v.Vector...),
|
||||||
|
Payload: qdrant.NewValueMap(v.Payload.ToMap()),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := p.client.Upsert(ctx, &qdrant.UpsertPoints{
|
||||||
|
CollectionName: collectionName,
|
||||||
|
Points: points,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to upsert vectors to collection %s: %w", collectionName, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) DeleteVectors(ctx context.Context, collectionName string, ids []string) error {
|
||||||
|
if len(ids) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
pointIDs := make([]*qdrant.PointId, 0, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
pointIDs = append(pointIDs, qdrant.NewID(id))
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := p.client.Delete(ctx, &qdrant.DeletePoints{
|
||||||
|
CollectionName: collectionName,
|
||||||
|
Points: &qdrant.PointsSelector{
|
||||||
|
PointsSelectorOneOf: &qdrant.PointsSelector_Points{
|
||||||
|
Points: &qdrant.PointsIdsList{
|
||||||
|
Ids: pointIDs,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to delete vectors from collection %s: %w", collectionName, err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) {
|
||||||
|
filter := p.buildFilter(req.Filter)
|
||||||
|
|
||||||
|
results, err := p.client.Query(ctx, &qdrant.QueryPoints{
|
||||||
|
CollectionName: req.CollectionName,
|
||||||
|
Query: qdrant.NewQuery(req.Vector...),
|
||||||
|
Limit: qdrant.PtrOf(uint64(req.TopK)),
|
||||||
|
ScoreThreshold: &req.ScoreThreshold,
|
||||||
|
Filter: filter,
|
||||||
|
WithPayload: qdrant.NewWithPayload(true),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to search collection %s: %w", req.CollectionName, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
searchResults := make([]SearchResult, 0, len(results))
|
||||||
|
for _, r := range results {
|
||||||
|
payload := make(map[string]any)
|
||||||
|
if r.Payload != nil {
|
||||||
|
for k, v := range r.Payload {
|
||||||
|
payload[k] = p.extractPayloadValue(v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
id := ""
|
||||||
|
if r.Id != nil {
|
||||||
|
id = r.Id.GetUuid()
|
||||||
|
}
|
||||||
|
|
||||||
|
searchResults = append(searchResults, SearchResult{
|
||||||
|
ID: id,
|
||||||
|
Score: r.Score,
|
||||||
|
Payload: ChunkPayloadFromMap(payload),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return searchResults, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) buildFilter(filter *SearchFilter) *qdrant.Filter {
|
||||||
|
if filter == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
must := make([]*qdrant.Condition, 0, 2)
|
||||||
|
if len(filter.KnowledgeBaseIDs) > 0 {
|
||||||
|
must = append(must, qdrant.NewMatchInts("knowledge_base_id", filter.KnowledgeBaseIDs...))
|
||||||
|
}
|
||||||
|
if len(filter.DocumentIDs) > 0 {
|
||||||
|
must = append(must, qdrant.NewMatchInts("document_id", filter.DocumentIDs...))
|
||||||
|
}
|
||||||
|
if len(must) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return &qdrant.Filter{Must: must}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *QdrantProvider) extractPayloadValue(v *qdrant.Value) interface{} {
|
||||||
|
if v == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
switch val := v.Kind.(type) {
|
||||||
|
case *qdrant.Value_StringValue:
|
||||||
|
return val.StringValue
|
||||||
|
case *qdrant.Value_IntegerValue:
|
||||||
|
return val.IntegerValue
|
||||||
|
case *qdrant.Value_DoubleValue:
|
||||||
|
return val.DoubleValue
|
||||||
|
case *qdrant.Value_BoolValue:
|
||||||
|
return val.BoolValue
|
||||||
|
case *qdrant.Value_ListValue:
|
||||||
|
list := make([]interface{}, 0, len(val.ListValue.Values))
|
||||||
|
for _, item := range val.ListValue.Values {
|
||||||
|
list = append(list, p.extractPayloadValue(item))
|
||||||
|
}
|
||||||
|
return list
|
||||||
|
case *qdrant.Value_StructValue:
|
||||||
|
m := make(map[string]interface{})
|
||||||
|
for k, v := range val.StructValue.Fields {
|
||||||
|
m[k] = p.extractPayloadValue(v)
|
||||||
|
}
|
||||||
|
return m
|
||||||
|
default:
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
48
internal/ai/rag/vectordb/types.go
Normal file
48
internal/ai/rag/vectordb/types.go
Normal file
@@ -0,0 +1,48 @@
|
|||||||
|
package vectordb
|
||||||
|
|
||||||
|
import "context"
|
||||||
|
|
||||||
|
type Vector struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Vector []float32 `json:"vector"`
|
||||||
|
Payload ChunkPayload `json:"payload"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchRequest struct {
|
||||||
|
CollectionName string `json:"collectionName"`
|
||||||
|
Vector []float32 `json:"vector"`
|
||||||
|
TopK int `json:"topK"`
|
||||||
|
ScoreThreshold float32 `json:"scoreThreshold"`
|
||||||
|
Filter *SearchFilter `json:"filter,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchFilter struct {
|
||||||
|
KnowledgeBaseIDs []int64 `json:"knowledgeBaseIds,omitempty"`
|
||||||
|
DocumentIDs []int64 `json:"documentIds,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SearchResult struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Score float32 `json:"score"`
|
||||||
|
Payload ChunkPayload `json:"payload"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type CollectionInfo struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Dimension int `json:"dimension"`
|
||||||
|
PointCount int `json:"pointCount"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type Provider interface {
|
||||||
|
CreateCollection(ctx context.Context, name string, dimension int) error
|
||||||
|
DeleteCollection(ctx context.Context, name string) error
|
||||||
|
GetCollection(ctx context.Context, name string) (*CollectionInfo, error)
|
||||||
|
ListCollections(ctx context.Context) ([]string, error)
|
||||||
|
|
||||||
|
UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error
|
||||||
|
DeleteVectors(ctx context.Context, collectionName string, ids []string) error
|
||||||
|
|
||||||
|
Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error)
|
||||||
|
Close() error
|
||||||
|
}
|
||||||
176
internal/ai/runtime/debug_run.go
Normal file
176
internal/ai/runtime/debug_run.go
Normal file
@@ -0,0 +1,176 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/ai/runtime/graphs"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/dto/request"
|
||||||
|
"agent-desk/internal/pkg/dto/response"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
svc.SkillDebugRunHook = DebugRunSkill
|
||||||
|
svc.SkillDebugResumeHook = DebugResumeSkill
|
||||||
|
}
|
||||||
|
|
||||||
|
func DebugRunSkill(ctx context.Context, req request.SkillDebugRunRequest) (*response.SkillDebugRunResponse, error) {
|
||||||
|
aiAgent := svc.AIAgentService.Get(req.AIAgentID)
|
||||||
|
if aiAgent == nil || aiAgent.Status != enums.StatusOk {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0007")
|
||||||
|
}
|
||||||
|
aiConfig := svc.AIConfigService.Get(aiAgent.AIConfigID)
|
||||||
|
if aiConfig == nil {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0008")
|
||||||
|
}
|
||||||
|
skill := svc.SkillDefinitionService.Get(req.SkillDefinitionID)
|
||||||
|
if skill == nil || skill.Status != enums.StatusOk {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0054")
|
||||||
|
}
|
||||||
|
debugAgent := *aiAgent
|
||||||
|
debugAgent.SkillIDs = fmt.Sprintf("%d", skill.ID)
|
||||||
|
var conversation *models.Conversation
|
||||||
|
if req.ConversationID > 0 {
|
||||||
|
if conversation = svc.ConversationService.Get(req.ConversationID); conversation == nil {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0116")
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
conversation = &models.Conversation{ID: req.ConversationID, AIAgentID: req.AIAgentID}
|
||||||
|
}
|
||||||
|
message := models.Message{
|
||||||
|
ConversationID: req.ConversationID,
|
||||||
|
SenderType: enums.IMSenderTypeCustomer,
|
||||||
|
MessageType: enums.IMMessageTypeText,
|
||||||
|
Content: strings.TrimSpace(req.UserMessage),
|
||||||
|
}
|
||||||
|
summary, err := applicationruntime.DefaultAgentApplicationService.RunPrepared(ctx, applicationruntime.RunInput{
|
||||||
|
Conversation: *conversation,
|
||||||
|
UserMessage: message,
|
||||||
|
AIAgent: debugAgent,
|
||||||
|
AIConfig: *aiConfig,
|
||||||
|
Debug: true,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return buildSkillDebugRunResponse(req, summary, skill), err
|
||||||
|
}
|
||||||
|
return buildSkillDebugRunResponse(req, summary, skill), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func DebugResumeSkill(ctx context.Context, req request.SkillDebugResumeRequest) (*response.SkillDebugRunResponse, error) {
|
||||||
|
aiAgent := svc.AIAgentService.Get(req.AIAgentID)
|
||||||
|
if aiAgent == nil || aiAgent.Status != enums.StatusOk {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0007")
|
||||||
|
}
|
||||||
|
aiConfig := svc.AIConfigService.Get(aiAgent.AIConfigID)
|
||||||
|
if aiConfig == nil {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0008")
|
||||||
|
}
|
||||||
|
pendingInterrupt := svc.ConversationInterruptService.GetByCheckPointID(strings.TrimSpace(req.CheckPointID))
|
||||||
|
if pendingInterrupt == nil {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0014")
|
||||||
|
}
|
||||||
|
if pendingInterrupt.AIAgentID > 0 && pendingInterrupt.AIAgentID != req.AIAgentID {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0015")
|
||||||
|
}
|
||||||
|
conversationID := req.ConversationID
|
||||||
|
if conversationID <= 0 {
|
||||||
|
conversationID = pendingInterrupt.ConversationID
|
||||||
|
}
|
||||||
|
if conversationID <= 0 {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0116")
|
||||||
|
}
|
||||||
|
conversation := svc.ConversationService.Get(conversationID)
|
||||||
|
if conversation == nil {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0116")
|
||||||
|
}
|
||||||
|
if conversation.AIAgentID > 0 && conversation.AIAgentID != req.AIAgentID {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0117")
|
||||||
|
}
|
||||||
|
resumeText := strings.TrimSpace(req.UserMessage)
|
||||||
|
summary, err := applicationruntime.DefaultAgentApplicationService.ResumePrepared(ctx, applicationruntime.ResumeInput{
|
||||||
|
Conversation: *conversation,
|
||||||
|
AIAgent: *aiAgent,
|
||||||
|
AIConfig: *aiConfig,
|
||||||
|
CheckPointID: strings.TrimSpace(req.CheckPointID),
|
||||||
|
ResumeData: map[string]string{
|
||||||
|
strings.TrimSpace(pendingInterrupt.InterruptID): resumeText,
|
||||||
|
},
|
||||||
|
Debug: true,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
if isCheckpointMissingError(err) {
|
||||||
|
summary = &applicationruntime.RunResult{
|
||||||
|
Status: "expired",
|
||||||
|
ReplyText: graphs.ConfirmationExpiredReply,
|
||||||
|
}
|
||||||
|
if pendingInterrupt.ID > 0 {
|
||||||
|
_ = svc.ConversationInterruptService.MarkExpired(pendingInterrupt.ID, 0)
|
||||||
|
}
|
||||||
|
return buildSkillDebugResumeResponse(req, summary, conversationID), nil
|
||||||
|
}
|
||||||
|
return buildSkillDebugResumeResponse(req, summary, conversationID), err
|
||||||
|
}
|
||||||
|
if pendingInterrupt.ID > 0 {
|
||||||
|
if summary != nil && summary.Interrupted {
|
||||||
|
_ = svc.ConversationInterruptService.MarkPendingAgain(pendingInterrupt.ID, firstInterruptID(summary), resolveInterruptPrompt(summary), 0)
|
||||||
|
} else if summary != nil && graphs.IsCancellationReply(summary.ReplyText) {
|
||||||
|
_ = svc.ConversationInterruptService.MarkCancelled(pendingInterrupt.ID, 0)
|
||||||
|
} else {
|
||||||
|
_ = svc.ConversationInterruptService.MarkResolved(pendingInterrupt.ID, 0)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return buildSkillDebugResumeResponse(req, summary, conversationID), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildSkillDebugRunResponse(req request.SkillDebugRunRequest, summary *applicationruntime.RunResult, skill *models.SkillDefinition) *response.SkillDebugRunResponse {
|
||||||
|
resp := &response.SkillDebugRunResponse{
|
||||||
|
ConversationID: req.ConversationID,
|
||||||
|
AIAgentID: req.AIAgentID,
|
||||||
|
}
|
||||||
|
if skill != nil {
|
||||||
|
resp.SkillDefinitionID = skill.ID
|
||||||
|
resp.SkillName = skill.Name
|
||||||
|
}
|
||||||
|
if summary == nil {
|
||||||
|
return resp
|
||||||
|
}
|
||||||
|
if resp.SkillDefinitionID <= 0 {
|
||||||
|
resp.SkillDefinitionID = summary.PlannedSkillID
|
||||||
|
}
|
||||||
|
resp.ReplyText = summary.ReplyText
|
||||||
|
resp.ToolWhitelist = append([]string(nil), summary.SkillAllowedToolCodes...)
|
||||||
|
resp.InvokedToolCodes = append([]string(nil), summary.InvokedToolCodes...)
|
||||||
|
resp.InterruptType = firstInterruptType(summary)
|
||||||
|
resp.CheckPointID = summary.CheckPointID
|
||||||
|
resp.Interrupted = summary.Interrupted
|
||||||
|
resp.TraceData = summary.TraceData
|
||||||
|
resp.ErrorMessage = summary.ErrorMessage
|
||||||
|
return resp
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildSkillDebugResumeResponse(req request.SkillDebugResumeRequest, summary *applicationruntime.RunResult, conversationID int64) *response.SkillDebugRunResponse {
|
||||||
|
resp := &response.SkillDebugRunResponse{
|
||||||
|
ConversationID: conversationID,
|
||||||
|
AIAgentID: req.AIAgentID,
|
||||||
|
}
|
||||||
|
if summary == nil {
|
||||||
|
return resp
|
||||||
|
}
|
||||||
|
resp.SkillDefinitionID = summary.PlannedSkillID
|
||||||
|
resp.SkillName = strings.TrimSpace(summary.PlannedSkillName)
|
||||||
|
resp.ReplyText = summary.ReplyText
|
||||||
|
resp.ToolWhitelist = append([]string(nil), summary.SkillAllowedToolCodes...)
|
||||||
|
resp.InvokedToolCodes = append([]string(nil), summary.InvokedToolCodes...)
|
||||||
|
resp.InterruptType = firstInterruptType(summary)
|
||||||
|
resp.CheckPointID = summary.CheckPointID
|
||||||
|
resp.Interrupted = summary.Interrupted
|
||||||
|
resp.TraceData = summary.TraceData
|
||||||
|
resp.ErrorMessage = summary.ErrorMessage
|
||||||
|
return resp
|
||||||
|
}
|
||||||
44
internal/ai/runtime/evaluation_run.go
Normal file
44
internal/ai/runtime/evaluation_run.go
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/pkg/dto/request"
|
||||||
|
"agent-desk/internal/pkg/dto/response"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/errorsx"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
svc.AgentEvaluationRunHook = RunAgentEvaluation
|
||||||
|
}
|
||||||
|
|
||||||
|
func RunAgentEvaluation(ctx context.Context, req request.RunAgentEvaluationRequest) (*response.AgentEvaluationReportResponse, error) {
|
||||||
|
agent := svc.AIAgentService.Get(req.AIAgentID)
|
||||||
|
if agent == nil || agent.Status != enums.StatusOk {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0007")
|
||||||
|
}
|
||||||
|
config := svc.AIConfigService.Get(agent.AIConfigID)
|
||||||
|
if config == nil {
|
||||||
|
return nil, errorsx.InvalidParamI18n("error.e0008")
|
||||||
|
}
|
||||||
|
cases := make([]applicationruntime.OfflineEvaluationCase, 0, len(req.Cases))
|
||||||
|
for _, item := range req.Cases {
|
||||||
|
cases = append(cases, applicationruntime.OfflineEvaluationCase{ID: item.ID, Category: item.Category, Message: item.Message, History: item.History, Expect: item.Expect})
|
||||||
|
}
|
||||||
|
report, err := applicationruntime.NewService().RunOfflineEvaluation(ctx, *agent, *config, cases)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
csv, err := report.CSV()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
ret := &response.AgentEvaluationReportResponse{Total: report.Total, Passed: report.Passed, CSV: csv, Results: make([]response.AgentEvaluationResultResponse, 0, len(report.Results))}
|
||||||
|
for _, item := range report.Results {
|
||||||
|
ret.Results = append(ret.Results, response.AgentEvaluationResultResponse{CaseID: item.CaseID, Category: item.Category, Passed: item.Passed, ReplyText: item.ReplyText, Interrupted: item.Interrupted, Error: item.Error, Finding: item.Finding})
|
||||||
|
}
|
||||||
|
return ret, nil
|
||||||
|
}
|
||||||
224
internal/ai/runtime/graphs/analyze_conversation_graph.go
Normal file
224
internal/ai/runtime/graphs/analyze_conversation_graph.go
Normal file
@@ -0,0 +1,224 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
type AnalyzeConversationInput struct {
|
||||||
|
Goal string `json:"goal"`
|
||||||
|
ObservedIssue string `json:"observedIssue"`
|
||||||
|
NeedTicket bool `json:"needTicket"`
|
||||||
|
NeedHumanHandoff bool `json:"needHumanHandoff"`
|
||||||
|
NeedQualityCheck bool `json:"needQualityCheck"`
|
||||||
|
AdditionalContext string `json:"additionalContext"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type AnalyzeConversationResult struct {
|
||||||
|
Summary string `json:"summary"`
|
||||||
|
UserIntent string `json:"userIntent"`
|
||||||
|
RiskLevel string `json:"riskLevel"`
|
||||||
|
RiskSignals []string `json:"riskSignals,omitempty"`
|
||||||
|
RecommendedNextAction string `json:"recommendedNextAction"`
|
||||||
|
RecommendedQuestions []string `json:"recommendedQuestions,omitempty"`
|
||||||
|
ConversationFacts []string `json:"conversationFacts,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type AnalyzeConversationGraph struct {
|
||||||
|
conversation models.Conversation
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewAnalyzeConversationGraph(conversation models.Conversation) *AnalyzeConversationGraph {
|
||||||
|
return &AnalyzeConversationGraph{conversation: conversation}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *AnalyzeConversationGraph) Run(_ context.Context, argumentsInJSON string) (string, error) {
|
||||||
|
input, err := g.parseInput(argumentsInJSON)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
messages, _, _ := services.MessageService.FindByConversationIDCursor(g.conversation.ID, 0, 8, "", "")
|
||||||
|
result := buildAnalyzeConversationResult(g.conversation, messages, input)
|
||||||
|
buf, err := json.Marshal(result)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(buf), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *AnalyzeConversationGraph) parseInput(argumentsInJSON string) (AnalyzeConversationInput, error) {
|
||||||
|
var input AnalyzeConversationInput
|
||||||
|
if strings.TrimSpace(argumentsInJSON) == "" {
|
||||||
|
return input, nil
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal([]byte(argumentsInJSON), &input); err != nil {
|
||||||
|
return input, fmt.Errorf("invalid analyze conversation arguments: %w", err)
|
||||||
|
}
|
||||||
|
input.Goal = strings.TrimSpace(input.Goal)
|
||||||
|
input.ObservedIssue = strings.TrimSpace(input.ObservedIssue)
|
||||||
|
input.AdditionalContext = strings.TrimSpace(input.AdditionalContext)
|
||||||
|
return input, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildAnalyzeConversationResult(conversation models.Conversation, messages []models.Message, input AnalyzeConversationInput) AnalyzeConversationResult {
|
||||||
|
joined := strings.ToLower(buildConversationCorpus(conversation, messages, input))
|
||||||
|
signals := collectRiskSignals(joined, input)
|
||||||
|
intent := detectUserIntent(joined, input)
|
||||||
|
recommendedAction := recommendNextAction(intent, signals, input)
|
||||||
|
result := AnalyzeConversationResult{
|
||||||
|
Summary: buildConversationSummary(conversation, messages, input),
|
||||||
|
UserIntent: intent,
|
||||||
|
RiskLevel: deriveRiskLevel(signals),
|
||||||
|
RiskSignals: signals,
|
||||||
|
RecommendedNextAction: recommendedAction,
|
||||||
|
RecommendedQuestions: recommendQuestions(intent, signals, input),
|
||||||
|
ConversationFacts: buildConversationAnalysisFacts(conversation, messages),
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildConversationSummary(conversation models.Conversation, messages []models.Message, input AnalyzeConversationInput) string {
|
||||||
|
parts := make([]string, 0, 4)
|
||||||
|
if input.ObservedIssue != "" {
|
||||||
|
parts = append(parts, "当前问题:"+input.ObservedIssue)
|
||||||
|
} else if strings.TrimSpace(conversation.LastMessageSummary) != "" {
|
||||||
|
parts = append(parts, "当前问题:"+strings.TrimSpace(conversation.LastMessageSummary))
|
||||||
|
}
|
||||||
|
if digest := buildRecentMessageDigest(messages); digest != "" {
|
||||||
|
parts = append(parts, "最近对话:"+digest)
|
||||||
|
}
|
||||||
|
if input.AdditionalContext != "" {
|
||||||
|
parts = append(parts, "补充信息:"+input.AdditionalContext)
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(strings.Join(parts, "\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildConversationCorpus(conversation models.Conversation, messages []models.Message, input AnalyzeConversationInput) string {
|
||||||
|
parts := make([]string, 0, len(messages)+4)
|
||||||
|
parts = append(parts, strings.TrimSpace(conversation.LastMessageSummary))
|
||||||
|
parts = append(parts, input.Goal, input.ObservedIssue, input.AdditionalContext)
|
||||||
|
for i := range messages {
|
||||||
|
parts = append(parts, strings.TrimSpace(messages[i].Content))
|
||||||
|
}
|
||||||
|
return strings.Join(parts, "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func collectRiskSignals(joined string, input AnalyzeConversationInput) []string {
|
||||||
|
signals := make([]string, 0, 6)
|
||||||
|
add := func(signal string) {
|
||||||
|
for _, item := range signals {
|
||||||
|
if item == signal {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
signals = append(signals, signal)
|
||||||
|
}
|
||||||
|
if containsAny(joined, "投诉", "举报", "差评", "曝光", "媒体", "起诉", "律师", "12315") {
|
||||||
|
add("complaint_escalation")
|
||||||
|
}
|
||||||
|
if containsAny(joined, "退款", "赔偿", "损失", "扣款", "重复扣费", "金额") {
|
||||||
|
add("financial_risk")
|
||||||
|
}
|
||||||
|
if containsAny(joined, "生气", "愤怒", "垃圾", "太差", "一直没人", "再不处理") {
|
||||||
|
add("negative_sentiment")
|
||||||
|
}
|
||||||
|
if containsAny(joined, "人工", "转人工", "真人", "客服") || input.NeedHumanHandoff {
|
||||||
|
add("handoff_requested")
|
||||||
|
}
|
||||||
|
if containsAny(joined, "工单", "报障", "售后", "登记", "记录问题") || input.NeedTicket {
|
||||||
|
add("ticket_expected")
|
||||||
|
}
|
||||||
|
if input.NeedQualityCheck {
|
||||||
|
add("quality_review_requested")
|
||||||
|
}
|
||||||
|
return signals
|
||||||
|
}
|
||||||
|
|
||||||
|
func detectUserIntent(joined string, input AnalyzeConversationInput) string {
|
||||||
|
switch {
|
||||||
|
case input.NeedHumanHandoff || containsAny(joined, "人工", "转人工", "真人"):
|
||||||
|
return "handoff_request"
|
||||||
|
case input.NeedTicket || containsAny(joined, "工单", "报障", "售后", "登记问题"):
|
||||||
|
return "ticket_request"
|
||||||
|
case containsAny(joined, "投诉", "举报", "差评", "赔偿"):
|
||||||
|
return "complaint"
|
||||||
|
default:
|
||||||
|
return "general_support"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func deriveRiskLevel(signals []string) string {
|
||||||
|
if len(signals) == 0 {
|
||||||
|
return "low"
|
||||||
|
}
|
||||||
|
if containsSignal(signals, "complaint_escalation") || containsSignal(signals, "financial_risk") {
|
||||||
|
return "high"
|
||||||
|
}
|
||||||
|
if containsSignal(signals, "handoff_requested") || containsSignal(signals, "negative_sentiment") {
|
||||||
|
return "medium"
|
||||||
|
}
|
||||||
|
return "low"
|
||||||
|
}
|
||||||
|
|
||||||
|
func recommendNextAction(intent string, signals []string, input AnalyzeConversationInput) string {
|
||||||
|
switch {
|
||||||
|
case input.NeedQualityCheck:
|
||||||
|
return "quality_review"
|
||||||
|
case containsSignal(signals, "handoff_requested") || intent == "handoff_request":
|
||||||
|
return "handoff_to_human"
|
||||||
|
case containsSignal(signals, "ticket_expected") || intent == "ticket_request":
|
||||||
|
return "prepare_ticket"
|
||||||
|
case containsSignal(signals, "complaint_escalation"):
|
||||||
|
return "handoff_to_human"
|
||||||
|
default:
|
||||||
|
return "continue_answering"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func recommendQuestions(intent string, signals []string, input AnalyzeConversationInput) []string {
|
||||||
|
questions := make([]string, 0, 3)
|
||||||
|
if containsSignal(signals, "ticket_expected") && strings.TrimSpace(input.ObservedIssue) == "" {
|
||||||
|
questions = append(questions, "Please confirm the specific issue, error message, and expected outcome.")
|
||||||
|
}
|
||||||
|
if containsSignal(signals, "handoff_requested") {
|
||||||
|
questions = append(questions, "Please confirm whether the user explicitly requested human support and why the issue needs human handling.")
|
||||||
|
}
|
||||||
|
if intent == "complaint" {
|
||||||
|
questions = append(questions, "Please confirm the complaint, impact, and the user's most important desired outcome.")
|
||||||
|
}
|
||||||
|
return questions
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildConversationAnalysisFacts(conversation models.Conversation, messages []models.Message) []string {
|
||||||
|
facts := make([]string, 0, 4)
|
||||||
|
if strings.TrimSpace(conversation.LastMessageSummary) != "" {
|
||||||
|
facts = append(facts, "最近摘要:"+strings.TrimSpace(conversation.LastMessageSummary))
|
||||||
|
}
|
||||||
|
if digest := buildRecentMessageDigest(messages); digest != "" {
|
||||||
|
facts = append(facts, "最近消息:"+digest)
|
||||||
|
}
|
||||||
|
return facts
|
||||||
|
}
|
||||||
|
|
||||||
|
func containsAny(value string, keywords ...string) bool {
|
||||||
|
for _, keyword := range keywords {
|
||||||
|
if keyword != "" && strings.Contains(value, strings.ToLower(keyword)) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func containsSignal(signals []string, target string) bool {
|
||||||
|
for _, item := range signals {
|
||||||
|
if item == target {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
@@ -0,0 +1,51 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuildAnalyzeConversationResult_RecommendsHandoffForComplaint(t *testing.T) {
|
||||||
|
conversation := models.Conversation{
|
||||||
|
LastMessageSummary: "用户反馈被重复扣费,并要求人工处理",
|
||||||
|
}
|
||||||
|
messages := []models.Message{
|
||||||
|
{SenderType: enums.IMSenderTypeCustomer, Content: "你们重复扣费了,我要投诉并转人工"},
|
||||||
|
}
|
||||||
|
|
||||||
|
got := buildAnalyzeConversationResult(conversation, messages, AnalyzeConversationInput{
|
||||||
|
NeedHumanHandoff: true,
|
||||||
|
})
|
||||||
|
|
||||||
|
if got.UserIntent != "handoff_request" {
|
||||||
|
t.Fatalf("expected handoff_request, got %q", got.UserIntent)
|
||||||
|
}
|
||||||
|
if got.RiskLevel != "high" {
|
||||||
|
t.Fatalf("expected high risk, got %q", got.RiskLevel)
|
||||||
|
}
|
||||||
|
if got.RecommendedNextAction != "handoff_to_human" {
|
||||||
|
t.Fatalf("expected handoff_to_human, got %q", got.RecommendedNextAction)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildAnalyzeConversationResult_RecommendsPrepareTicket(t *testing.T) {
|
||||||
|
conversation := models.Conversation{
|
||||||
|
LastMessageSummary: "用户要求登记问题并尽快处理",
|
||||||
|
}
|
||||||
|
messages := []models.Message{
|
||||||
|
{SenderType: enums.IMSenderTypeCustomer, Content: "麻烦帮我建个工单,订单一直支付失败"},
|
||||||
|
}
|
||||||
|
|
||||||
|
got := buildAnalyzeConversationResult(conversation, messages, AnalyzeConversationInput{
|
||||||
|
NeedTicket: true,
|
||||||
|
})
|
||||||
|
|
||||||
|
if got.UserIntent != "ticket_request" {
|
||||||
|
t.Fatalf("expected ticket_request, got %q", got.UserIntent)
|
||||||
|
}
|
||||||
|
if got.RecommendedNextAction != "prepare_ticket" {
|
||||||
|
t.Fatalf("expected prepare_ticket, got %q", got.RecommendedNextAction)
|
||||||
|
}
|
||||||
|
}
|
||||||
153
internal/ai/runtime/graphs/create_ticket_graph.go
Normal file
153
internal/ai/runtime/graphs/create_ticket_graph.go
Normal file
@@ -0,0 +1,153 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/runtime/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/dto"
|
||||||
|
"agent-desk/internal/pkg/dto/request"
|
||||||
|
"agent-desk/internal/pkg/i18nx"
|
||||||
|
"agent-desk/internal/services"
|
||||||
|
|
||||||
|
componenttool "github.com/cloudwego/eino/components/tool"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
type CreateTicketGraphState struct {
|
||||||
|
Request request.CreateTicketFromConversationRequest
|
||||||
|
}
|
||||||
|
|
||||||
|
type CreateTicketGraphInterruptInfo struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Message string `json:"message"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type createTicketGraphArgs struct {
|
||||||
|
Title string `json:"title"`
|
||||||
|
Description string `json:"description"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
schema.RegisterName[CreateTicketGraphState]("cs_ai_agent_create_ticket_graph_state")
|
||||||
|
schema.RegisterName[CreateTicketGraphInterruptInfo]("cs_ai_agent_create_ticket_graph_interrupt_info")
|
||||||
|
}
|
||||||
|
|
||||||
|
type CreateTicketGraph struct {
|
||||||
|
conversation models.Conversation
|
||||||
|
aiAgent models.AIAgent
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewCreateTicketGraph(conversation models.Conversation, aiAgent models.AIAgent) *CreateTicketGraph {
|
||||||
|
return &CreateTicketGraph{
|
||||||
|
conversation: conversation,
|
||||||
|
aiAgent: aiAgent,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *CreateTicketGraph) Run(ctx context.Context, argumentsInJSON string) (string, error) {
|
||||||
|
locale := services.ConversationHumanDispatchService.ResolveConversationLocale(g.conversation.ID)
|
||||||
|
wasInterrupted, hasState, state := componenttool.GetInterruptState[CreateTicketGraphState](ctx)
|
||||||
|
if !wasInterrupted {
|
||||||
|
req, err := g.buildCreateRequest(argumentsInJSON)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
info := CreateTicketGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeTicketCreationConfirmation,
|
||||||
|
Message: g.buildConfirmationPrompt(req, locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, CreateTicketGraphState{Request: req})
|
||||||
|
}
|
||||||
|
if !hasState {
|
||||||
|
return "", fmt.Errorf("create ticket graph state missing")
|
||||||
|
}
|
||||||
|
isResumeTarget, hasData, resumeText := componenttool.GetResumeContext[string](ctx)
|
||||||
|
if !isResumeTarget {
|
||||||
|
info := CreateTicketGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeTicketCreationConfirmation,
|
||||||
|
Message: g.buildConfirmationPrompt(state.Request, locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, state)
|
||||||
|
}
|
||||||
|
if !hasData {
|
||||||
|
info := CreateTicketGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeTicketCreationConfirmation,
|
||||||
|
Message: ConfirmOrCancelPromptForLocale(locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, state)
|
||||||
|
}
|
||||||
|
decision := ParseConfirmationDecision(resumeText)
|
||||||
|
switch decision {
|
||||||
|
case ConfirmationDecisionConfirm:
|
||||||
|
item, err := services.TicketService.CreateFromConversation(state.Request, g.buildAIPrincipal())
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return tooling.MarshalToolResult(tooling.ToolResult{
|
||||||
|
Handled: true,
|
||||||
|
Terminal: true,
|
||||||
|
Action: "ticket_created",
|
||||||
|
ReplyText: i18nx.Getf(locale, "graph.ticketCreated", strings.TrimSpace(item.TicketNo), strings.TrimSpace(item.Title)),
|
||||||
|
ShouldRetry: false,
|
||||||
|
}), nil
|
||||||
|
case ConfirmationDecisionCancel:
|
||||||
|
return tooling.MarshalToolResult(tooling.ToolResult{
|
||||||
|
Handled: true,
|
||||||
|
Terminal: true,
|
||||||
|
Action: "ticket_cancelled",
|
||||||
|
ReplyText: CancelCreateTicketReplyForLocale(locale),
|
||||||
|
ShouldRetry: false,
|
||||||
|
}), nil
|
||||||
|
default:
|
||||||
|
info := CreateTicketGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeTicketCreationConfirmation,
|
||||||
|
Message: NeedExplicitConfirmationPromptForLocale(locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, state)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *CreateTicketGraph) buildCreateRequest(argumentsInJSON string) (request.CreateTicketFromConversationRequest, error) {
|
||||||
|
req := request.CreateTicketFromConversationRequest{
|
||||||
|
ConversationID: g.conversation.ID,
|
||||||
|
}
|
||||||
|
var args createTicketGraphArgs
|
||||||
|
if strings.TrimSpace(argumentsInJSON) != "" {
|
||||||
|
if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {
|
||||||
|
return req, fmt.Errorf("invalid create ticket arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
req.Title = strings.TrimSpace(args.Title)
|
||||||
|
req.Description = strings.TrimSpace(args.Description)
|
||||||
|
if req.Title == "" {
|
||||||
|
req.Title = strings.TrimSpace(g.conversation.LastMessageSummary)
|
||||||
|
}
|
||||||
|
if req.Description == "" {
|
||||||
|
req.Description = strings.TrimSpace(g.conversation.LastMessageSummary)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(req.Title) == "" {
|
||||||
|
return req, fmt.Errorf("ticket title is required")
|
||||||
|
}
|
||||||
|
return req, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *CreateTicketGraph) buildConfirmationPrompt(req request.CreateTicketFromConversationRequest, locale string) string {
|
||||||
|
return i18nx.Getf(locale, "graph.createTicketConfirmPrompt",
|
||||||
|
strings.TrimSpace(req.Title), strings.TrimSpace(req.Description))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *CreateTicketGraph) buildAIPrincipal() *dto.AuthPrincipal {
|
||||||
|
username := "AI"
|
||||||
|
if strings.TrimSpace(g.aiAgent.Name) != "" {
|
||||||
|
username = strings.TrimSpace(g.aiAgent.Name)
|
||||||
|
}
|
||||||
|
return &dto.AuthPrincipal{
|
||||||
|
UserID: 0,
|
||||||
|
Username: username,
|
||||||
|
Nickname: username,
|
||||||
|
}
|
||||||
|
}
|
||||||
150
internal/ai/runtime/graphs/handoff_graph.go
Normal file
150
internal/ai/runtime/graphs/handoff_graph.go
Normal file
@@ -0,0 +1,150 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/runtime/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/i18nx"
|
||||||
|
"agent-desk/internal/pkg/tracex"
|
||||||
|
"agent-desk/internal/services"
|
||||||
|
|
||||||
|
componenttool "github.com/cloudwego/eino/components/tool"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
type HandoffGraphState struct {
|
||||||
|
Reason string `json:"reason"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type HandoffGraphInterruptInfo struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Message string `json:"message"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type handoffGraphArgs struct {
|
||||||
|
Reason string `json:"reason"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
schema.RegisterName[HandoffGraphState]("cs_ai_agent_handoff_graph_state")
|
||||||
|
schema.RegisterName[HandoffGraphInterruptInfo]("cs_ai_agent_handoff_graph_interrupt_info")
|
||||||
|
}
|
||||||
|
|
||||||
|
type HandoffGraph struct {
|
||||||
|
conversation models.Conversation
|
||||||
|
aiAgent models.AIAgent
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewHandoffGraph(conversation models.Conversation, aiAgent models.AIAgent) *HandoffGraph {
|
||||||
|
return &HandoffGraph{
|
||||||
|
conversation: conversation,
|
||||||
|
aiAgent: aiAgent,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *HandoffGraph) Run(ctx context.Context, argumentsInJSON string) (string, error) {
|
||||||
|
locale := services.ConversationHumanDispatchService.ResolveConversationLocale(g.conversation.ID)
|
||||||
|
wasInterrupted, hasState, state := componenttool.GetInterruptState[HandoffGraphState](ctx)
|
||||||
|
if !wasInterrupted {
|
||||||
|
reason, err := g.buildReason(argumentsInJSON, locale)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
requestID := tracex.RequestIDFromContext(ctx)
|
||||||
|
handled, err := services.ConversationService.TryOffHoursHandoffByAIWithRequestID(g.conversation.ID, g.aiAgent, reason, requestID)
|
||||||
|
if err != nil || handled {
|
||||||
|
if handled && err == nil {
|
||||||
|
offHoursMsg := strings.TrimSpace(g.aiAgent.OffHoursMessage)
|
||||||
|
if offHoursMsg == "" {
|
||||||
|
offHoursMsg = services.HandoffOffHoursMessageForLocale(locale)
|
||||||
|
}
|
||||||
|
return tooling.MarshalToolResult(tooling.ToolResult{
|
||||||
|
Handled: true,
|
||||||
|
Terminal: true,
|
||||||
|
Action: "off_hours_handoff",
|
||||||
|
ReplyText: offHoursMsg,
|
||||||
|
ReplySent: true,
|
||||||
|
ShouldRetry: false,
|
||||||
|
}), nil
|
||||||
|
}
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
info := HandoffGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeHandoffConfirmation,
|
||||||
|
Message: g.buildConfirmationPrompt(reason, locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, HandoffGraphState{Reason: reason})
|
||||||
|
}
|
||||||
|
if !hasState {
|
||||||
|
return "", fmt.Errorf("handoff graph state missing")
|
||||||
|
}
|
||||||
|
isResumeTarget, hasData, resumeText := componenttool.GetResumeContext[string](ctx)
|
||||||
|
if !isResumeTarget {
|
||||||
|
info := HandoffGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeHandoffConfirmation,
|
||||||
|
Message: g.buildConfirmationPrompt(state.Reason, locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, state)
|
||||||
|
}
|
||||||
|
if !hasData {
|
||||||
|
info := HandoffGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeHandoffConfirmation,
|
||||||
|
Message: ConfirmOrCancelPromptForLocale(locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, state)
|
||||||
|
}
|
||||||
|
switch parseHandoffDecision(resumeText) {
|
||||||
|
case ConfirmationDecisionConfirm:
|
||||||
|
if err := services.ConversationService.HandoffByAIWithRequestID(g.conversation.ID, g.aiAgent, state.Reason, tracex.RequestIDFromContext(ctx)); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
// ConversationService sends the customer-visible handoff notice according to the dispatch decision.
|
||||||
|
return tooling.MarshalToolResult(tooling.ToolResult{
|
||||||
|
Handled: true,
|
||||||
|
Terminal: true,
|
||||||
|
Action: "handoff_confirmed",
|
||||||
|
ReplySent: true,
|
||||||
|
ShouldRetry: false,
|
||||||
|
}), nil
|
||||||
|
case ConfirmationDecisionCancel:
|
||||||
|
return tooling.MarshalToolResult(tooling.ToolResult{
|
||||||
|
Handled: true,
|
||||||
|
Terminal: true,
|
||||||
|
Action: "handoff_cancelled",
|
||||||
|
ReplyText: CancelHandoffReplyForLocale(locale),
|
||||||
|
ShouldRetry: false,
|
||||||
|
}), nil
|
||||||
|
default:
|
||||||
|
info := HandoffGraphInterruptInfo{
|
||||||
|
Type: InterruptTypeHandoffConfirmation,
|
||||||
|
Message: NeedExplicitConfirmationPromptForLocale(locale),
|
||||||
|
}
|
||||||
|
return "", componenttool.StatefulInterrupt(ctx, info, state)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *HandoffGraph) buildReason(argumentsInJSON string, locale string) (string, error) {
|
||||||
|
reason := i18nx.Getf(locale, "graph.defaultHandoffReason")
|
||||||
|
var args handoffGraphArgs
|
||||||
|
if strings.TrimSpace(argumentsInJSON) != "" {
|
||||||
|
if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil {
|
||||||
|
return "", fmt.Errorf("invalid handoff arguments: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if parsed := strings.TrimSpace(args.Reason); parsed != "" {
|
||||||
|
reason = parsed
|
||||||
|
}
|
||||||
|
return reason, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *HandoffGraph) buildConfirmationPrompt(reason string, locale string) string {
|
||||||
|
return i18nx.Getf(locale, "graph.handoffConfirmPrompt", strings.TrimSpace(reason))
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseHandoffDecision(value string) ConfirmationDecision {
|
||||||
|
return ParseConfirmationDecision(value)
|
||||||
|
}
|
||||||
247
internal/ai/runtime/graphs/handoff_graph_test.go
Normal file
247
internal/ai/runtime/graphs/handoff_graph_test.go
Normal file
@@ -0,0 +1,247 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/runtime/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/services"
|
||||||
|
|
||||||
|
"github.com/glebarez/sqlite"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
"gorm.io/gorm"
|
||||||
|
"gorm.io/gorm/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestHandoffGraphOffHoursSendsNoticeWithoutConfirmation(t *testing.T) {
|
||||||
|
db := setupHandoffGraphTestDB(t)
|
||||||
|
aiAgent := createHandoffGraphAIAgent(t, db, "1")
|
||||||
|
conversation := createHandoffGraphConversation(t, db, aiAgent.ID)
|
||||||
|
|
||||||
|
reply, err := NewHandoffGraph(conversation, aiAgent).Run(context.Background(), `{"reason":"用户要求转人工"}`)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Run() error = %v", err)
|
||||||
|
}
|
||||||
|
var result tooling.ToolResult
|
||||||
|
if err := json.Unmarshal([]byte(reply), &result); err != nil {
|
||||||
|
t.Fatalf("expected graph tool result JSON, got %q: %v", reply, err)
|
||||||
|
}
|
||||||
|
if !result.Handled || !result.Terminal || result.ShouldRetry {
|
||||||
|
t.Fatalf("unexpected graph result flags: %+v", result)
|
||||||
|
}
|
||||||
|
if !result.ReplySent {
|
||||||
|
t.Fatalf("expected graph result to mark replySent, got %+v", result)
|
||||||
|
}
|
||||||
|
expectedMsg := services.HandoffOffHoursMessageForLocale(services.ResolveConversationLocale(conversation.ID))
|
||||||
|
if result.Action != "off_hours_handoff" || result.ReplyText != expectedMsg {
|
||||||
|
t.Fatalf("unexpected off-hours graph result: %+v, want %q", result, expectedMsg)
|
||||||
|
}
|
||||||
|
|
||||||
|
message := services.MessageService.FindOne(sqls.NewCnd().Eq("conversation_id", conversation.ID).Desc("id"))
|
||||||
|
if message == nil {
|
||||||
|
t.Fatalf("expected off-hours notice message")
|
||||||
|
}
|
||||||
|
if message.Content != expectedMsg {
|
||||||
|
t.Fatalf("expected off-hours notice %q, got %q", expectedMsg, message.Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
current := services.ConversationService.Get(conversation.ID)
|
||||||
|
if current == nil {
|
||||||
|
t.Fatalf("expected conversation")
|
||||||
|
}
|
||||||
|
if current.Status != enums.IMConversationStatusPending {
|
||||||
|
t.Fatalf("expected conversation to be queued in pending status, got %d", current.Status)
|
||||||
|
}
|
||||||
|
if current.HandoffAt == nil {
|
||||||
|
t.Fatalf("expected handoff_at to be set, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandoffGraphOffHoursSendsNoticeInRussianForRussianConversation(t *testing.T) {
|
||||||
|
db := setupHandoffGraphTestDB(t)
|
||||||
|
aiAgent := createHandoffGraphAIAgent(t, db, "1")
|
||||||
|
conversation := createHandoffGraphConversation(t, db, aiAgent.ID)
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
if err := db.Create(&models.Message{
|
||||||
|
ConversationID: conversation.ID,
|
||||||
|
SenderType: enums.IMSenderTypeCustomer,
|
||||||
|
SenderID: 1,
|
||||||
|
Content: "Привет хочу возрат денег",
|
||||||
|
SentAt: &now,
|
||||||
|
}).Error; err != nil {
|
||||||
|
t.Fatalf("create message error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
reply, err := NewHandoffGraph(conversation, aiAgent).Run(context.Background(), `{"reason":"перенаправляю оператору"}`)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Run() error = %v", err)
|
||||||
|
}
|
||||||
|
var result tooling.ToolResult
|
||||||
|
if err := json.Unmarshal([]byte(reply), &result); err != nil {
|
||||||
|
t.Fatalf("expected graph tool result JSON, got %q: %v", reply, err)
|
||||||
|
}
|
||||||
|
expectedRuMsg := services.HandoffOffHoursMessageForLocale("ru-RU")
|
||||||
|
if result.Action != "off_hours_handoff" || result.ReplyText != expectedRuMsg {
|
||||||
|
t.Fatalf("unexpected off-hours graph result for ru-RU: %+v, want %q", result, expectedRuMsg)
|
||||||
|
}
|
||||||
|
|
||||||
|
message := services.MessageService.FindOne(sqls.NewCnd().Eq("conversation_id", conversation.ID).Eq("sender_type", enums.IMSenderTypeAI).Desc("id"))
|
||||||
|
if message == nil {
|
||||||
|
t.Fatalf("expected off-hours notice message")
|
||||||
|
}
|
||||||
|
if message.Content != expectedRuMsg {
|
||||||
|
t.Fatalf("expected Russian off-hours notice %q, got %q", expectedRuMsg, message.Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandoffGraphOffHoursUsesCustomAgentNotice(t *testing.T) {
|
||||||
|
db := setupHandoffGraphTestDB(t)
|
||||||
|
aiAgent := createHandoffGraphAIAgent(t, db, "1")
|
||||||
|
aiAgent.OffHoursMessage = "Операторы сейчас оффлайн. Мы ответим позже!"
|
||||||
|
conversation := createHandoffGraphConversation(t, db, aiAgent.ID)
|
||||||
|
|
||||||
|
reply, err := NewHandoffGraph(conversation, aiAgent).Run(context.Background(), `{"reason":"позови человека"}`)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Run() error = %v", err)
|
||||||
|
}
|
||||||
|
var result tooling.ToolResult
|
||||||
|
if err := json.Unmarshal([]byte(reply), &result); err != nil {
|
||||||
|
t.Fatalf("expected graph tool result JSON, got %q: %v", reply, err)
|
||||||
|
}
|
||||||
|
if result.Action != "off_hours_handoff" || result.ReplyText != aiAgent.OffHoursMessage {
|
||||||
|
t.Fatalf("expected custom notice %q, got %+v", aiAgent.OffHoursMessage, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
current := services.ConversationService.Get(conversation.ID)
|
||||||
|
if current == nil || current.Status != enums.IMConversationStatusPending {
|
||||||
|
t.Fatalf("expected conversation to be queued in pending status, got %v", current)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandoffGraphWithActiveScheduleStillRequestsConfirmation(t *testing.T) {
|
||||||
|
db := setupHandoffGraphTestDB(t)
|
||||||
|
aiAgent := createHandoffGraphAIAgent(t, db, "1")
|
||||||
|
createHandoffGraphTeam(t, db, 1)
|
||||||
|
createHandoffGraphActiveSchedule(t, db, 1)
|
||||||
|
conversation := createHandoffGraphConversation(t, db, aiAgent.ID)
|
||||||
|
|
||||||
|
reply, err := NewHandoffGraph(conversation, aiAgent).Run(context.Background(), `{"reason":"用户要求转人工"}`)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatalf("expected confirmation interrupt")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), InterruptTypeHandoffConfirmation) {
|
||||||
|
t.Fatalf("expected handoff confirmation interrupt, got %v", err)
|
||||||
|
}
|
||||||
|
if reply != "" {
|
||||||
|
t.Fatalf("expected no reply before confirmation, got %q", reply)
|
||||||
|
}
|
||||||
|
|
||||||
|
var count int64
|
||||||
|
if err := db.Model(&models.Message{}).Where("conversation_id = ?", conversation.ID).Count(&count).Error; err != nil {
|
||||||
|
t.Fatalf("count messages error = %v", err)
|
||||||
|
}
|
||||||
|
if count != 0 {
|
||||||
|
t.Fatalf("expected no notice before confirmation, got %d messages", count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func setupHandoffGraphTestDB(t *testing.T) *gorm.DB {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
dbName := strings.NewReplacer("/", "_", " ", "_").Replace(t.Name())
|
||||||
|
db, err := gorm.Open(sqlite.Open("file:"+dbName+"?mode=memory&cache=shared"), &gorm.Config{
|
||||||
|
NamingStrategy: schema.NamingStrategy{
|
||||||
|
TablePrefix: "t_",
|
||||||
|
SingularTable: true,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open sqlite error = %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
sqlDB, err := db.DB()
|
||||||
|
if err == nil {
|
||||||
|
_ = sqlDB.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
if err := db.AutoMigrate(
|
||||||
|
&models.Channel{},
|
||||||
|
&models.CustomerIdentity{},
|
||||||
|
&models.AIAgent{},
|
||||||
|
&models.AgentTeam{},
|
||||||
|
&models.AgentTeamSchedule{},
|
||||||
|
&models.Conversation{},
|
||||||
|
&models.ConversationAssignment{},
|
||||||
|
&models.ConversationEventLog{},
|
||||||
|
&models.ConversationReadState{},
|
||||||
|
&models.Message{},
|
||||||
|
&models.ChannelMessageOutbox{},
|
||||||
|
); err != nil {
|
||||||
|
t.Fatalf("auto migrate error = %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
return db
|
||||||
|
}
|
||||||
|
|
||||||
|
func createHandoffGraphAIAgent(t *testing.T, db *gorm.DB, teamIDs string) models.AIAgent {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
item := models.AIAgent{
|
||||||
|
Name: "测试AI",
|
||||||
|
ServiceMode: enums.IMConversationServiceModeAIFirst,
|
||||||
|
TeamIDs: teamIDs,
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
}
|
||||||
|
if err := db.Create(&item).Error; err != nil {
|
||||||
|
t.Fatalf("create ai agent error = %v", err)
|
||||||
|
}
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
|
||||||
|
func createHandoffGraphTeam(t *testing.T, db *gorm.DB, id int64) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
if err := db.Create(&models.AgentTeam{ID: id, Name: "售后支持组", Status: enums.StatusOk}).Error; err != nil {
|
||||||
|
t.Fatalf("create team error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func createHandoffGraphActiveSchedule(t *testing.T, db *gorm.DB, teamID int64) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
if err := db.Create(&models.AgentTeamSchedule{
|
||||||
|
TeamID: teamID,
|
||||||
|
StartAt: now.Add(-time.Hour),
|
||||||
|
EndAt: now.Add(time.Hour),
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
}).Error; err != nil {
|
||||||
|
t.Fatalf("create schedule error = %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func createHandoffGraphConversation(t *testing.T, db *gorm.DB, aiAgentID int64) models.Conversation {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
item := models.Conversation{
|
||||||
|
AIAgentID: aiAgentID,
|
||||||
|
ChannelID: 1,
|
||||||
|
CustomerID: 1,
|
||||||
|
CustomerName: "测试访客",
|
||||||
|
Status: enums.IMConversationStatusAIServing,
|
||||||
|
ServiceMode: enums.IMConversationServiceModeAIFirst,
|
||||||
|
LastMessageAt: now,
|
||||||
|
LastActiveAt: now,
|
||||||
|
}
|
||||||
|
if err := db.Create(&item).Error; err != nil {
|
||||||
|
t.Fatalf("create conversation error = %v", err)
|
||||||
|
}
|
||||||
|
return item
|
||||||
|
}
|
||||||
101
internal/ai/runtime/graphs/hitl.go
Normal file
101
internal/ai/runtime/graphs/hitl.go
Normal file
@@ -0,0 +1,101 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/i18nx"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
InterruptTypeTicketCreationConfirmation = "ticket_creation_confirmation"
|
||||||
|
InterruptTypeHandoffConfirmation = "handoff_confirmation"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
ConfirmOrCancelPrompt = i18nx.Get("graph.confirmOrCancel")
|
||||||
|
NeedExplicitConfirmationPrompt = i18nx.Get("graph.needExplicitConfirmation")
|
||||||
|
ConfirmationExpiredReply = i18nx.Get("graph.confirmationExpired")
|
||||||
|
CancelCreateTicketReply = i18nx.Get("graph.cancelCreateTicket")
|
||||||
|
CancelHandoffReply = i18nx.Get("graph.cancelHandoff")
|
||||||
|
)
|
||||||
|
|
||||||
|
func ConfirmOrCancelPromptForLocale(locale string) string {
|
||||||
|
if strings.TrimSpace(locale) == "" {
|
||||||
|
locale = i18nx.DefaultLocale
|
||||||
|
}
|
||||||
|
return i18nx.Getf(locale, "graph.confirmOrCancel")
|
||||||
|
}
|
||||||
|
|
||||||
|
func NeedExplicitConfirmationPromptForLocale(locale string) string {
|
||||||
|
if strings.TrimSpace(locale) == "" {
|
||||||
|
locale = i18nx.DefaultLocale
|
||||||
|
}
|
||||||
|
return i18nx.Getf(locale, "graph.needExplicitConfirmation")
|
||||||
|
}
|
||||||
|
|
||||||
|
func ConfirmationExpiredReplyForLocale(locale string) string {
|
||||||
|
if strings.TrimSpace(locale) == "" {
|
||||||
|
locale = i18nx.DefaultLocale
|
||||||
|
}
|
||||||
|
return i18nx.Getf(locale, "graph.confirmationExpired")
|
||||||
|
}
|
||||||
|
|
||||||
|
func CancelCreateTicketReplyForLocale(locale string) string {
|
||||||
|
if strings.TrimSpace(locale) == "" {
|
||||||
|
locale = i18nx.DefaultLocale
|
||||||
|
}
|
||||||
|
return i18nx.Getf(locale, "graph.cancelCreateTicket")
|
||||||
|
}
|
||||||
|
|
||||||
|
func CancelHandoffReplyForLocale(locale string) string {
|
||||||
|
if strings.TrimSpace(locale) == "" {
|
||||||
|
locale = i18nx.DefaultLocale
|
||||||
|
}
|
||||||
|
return i18nx.Getf(locale, "graph.cancelHandoff")
|
||||||
|
}
|
||||||
|
|
||||||
|
type ConfirmationDecision string
|
||||||
|
|
||||||
|
const (
|
||||||
|
ConfirmationDecisionConfirm ConfirmationDecision = "confirm"
|
||||||
|
ConfirmationDecisionCancel ConfirmationDecision = "cancel"
|
||||||
|
)
|
||||||
|
|
||||||
|
func ParseConfirmationDecision(value string) ConfirmationDecision {
|
||||||
|
value = strings.ToLower(strings.TrimSpace(value))
|
||||||
|
if value == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
cancelWords := []string{
|
||||||
|
"取消", "不用", "不需要", "算了", "no", "cancel", "stop",
|
||||||
|
"нет", "отмена", "отменить", "не надо", "не нужно", "стоп", "хватит", "отбой",
|
||||||
|
}
|
||||||
|
for _, item := range cancelWords {
|
||||||
|
if strings.Contains(value, item) {
|
||||||
|
return ConfirmationDecisionCancel
|
||||||
|
}
|
||||||
|
}
|
||||||
|
confirmWords := []string{
|
||||||
|
"确认", "是", "好的", "可以", "ok", "yes", "continue", "confirm", "继续", "同意",
|
||||||
|
"да", "подтверждаю", "подтвердить", "переводи", "переведи", "давай",
|
||||||
|
"согласен", "согласна", "верно", "точно", "конечно",
|
||||||
|
}
|
||||||
|
for _, item := range confirmWords {
|
||||||
|
if strings.Contains(value, item) {
|
||||||
|
return ConfirmationDecisionConfirm
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func IsCancellationReply(replyText string) bool {
|
||||||
|
replyText = strings.TrimSpace(replyText)
|
||||||
|
return strings.Contains(replyText, CancelCreateTicketReply) ||
|
||||||
|
strings.Contains(replyText, CancelHandoffReply) ||
|
||||||
|
strings.Contains(replyText, "已取消本次工单创建。") ||
|
||||||
|
strings.Contains(replyText, "已取消本次转人工。") ||
|
||||||
|
strings.Contains(replyText, "Создание обращения отменено.") ||
|
||||||
|
strings.Contains(replyText, "Передача человеку отменена.") ||
|
||||||
|
strings.Contains(replyText, "Ticket creation cancelled.") ||
|
||||||
|
strings.Contains(replyText, "Human handoff cancelled.")
|
||||||
|
}
|
||||||
186
internal/ai/runtime/graphs/prepare_ticket_draft_graph.go
Normal file
186
internal/ai/runtime/graphs/prepare_ticket_draft_graph.go
Normal file
@@ -0,0 +1,186 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
type PrepareTicketDraftInput struct {
|
||||||
|
Title string `json:"title"`
|
||||||
|
Description string `json:"description"`
|
||||||
|
Issue string `json:"issue"`
|
||||||
|
Impact string `json:"impact"`
|
||||||
|
ExpectedOutcome string `json:"expectedOutcome"`
|
||||||
|
CurrentAttempt string `json:"currentAttempt"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type PrepareTicketDraftResult struct {
|
||||||
|
Ready bool `json:"ready"`
|
||||||
|
Title string `json:"title"`
|
||||||
|
Description string `json:"description"`
|
||||||
|
MissingFields []string `json:"missingFields,omitempty"`
|
||||||
|
FollowUpQuestions []string `json:"followUpQuestions,omitempty"`
|
||||||
|
ConversationFacts []string `json:"conversationFacts,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type PrepareTicketDraftGraph struct {
|
||||||
|
conversation models.Conversation
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewPrepareTicketDraftGraph(conversation models.Conversation) *PrepareTicketDraftGraph {
|
||||||
|
return &PrepareTicketDraftGraph{conversation: conversation}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *PrepareTicketDraftGraph) Run(_ context.Context, argumentsInJSON string) (string, error) {
|
||||||
|
input, err := g.parseInput(argumentsInJSON)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
messages, _, _ := services.MessageService.FindByConversationIDCursor(g.conversation.ID, 0, 6, "", "")
|
||||||
|
result := buildPrepareTicketDraftResult(g.conversation, messages, input)
|
||||||
|
buf, err := json.Marshal(result)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(buf), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *PrepareTicketDraftGraph) parseInput(argumentsInJSON string) (PrepareTicketDraftInput, error) {
|
||||||
|
var input PrepareTicketDraftInput
|
||||||
|
if strings.TrimSpace(argumentsInJSON) == "" {
|
||||||
|
return input, nil
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal([]byte(argumentsInJSON), &input); err != nil {
|
||||||
|
return input, fmt.Errorf("invalid prepare ticket draft arguments: %w", err)
|
||||||
|
}
|
||||||
|
input.Title = strings.TrimSpace(input.Title)
|
||||||
|
input.Description = strings.TrimSpace(input.Description)
|
||||||
|
input.Issue = strings.TrimSpace(input.Issue)
|
||||||
|
input.Impact = strings.TrimSpace(input.Impact)
|
||||||
|
input.ExpectedOutcome = strings.TrimSpace(input.ExpectedOutcome)
|
||||||
|
input.CurrentAttempt = strings.TrimSpace(input.CurrentAttempt)
|
||||||
|
return input, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildPrepareTicketDraftResult(conversation models.Conversation, messages []models.Message, input PrepareTicketDraftInput) PrepareTicketDraftResult {
|
||||||
|
result := PrepareTicketDraftResult{
|
||||||
|
MissingFields: make([]string, 0, 2),
|
||||||
|
FollowUpQuestions: make([]string, 0, 2),
|
||||||
|
ConversationFacts: buildConversationFacts(conversation, messages),
|
||||||
|
}
|
||||||
|
result.Title = buildDraftTitle(conversation, input)
|
||||||
|
result.Description = buildDraftDescription(conversation, messages, input)
|
||||||
|
if strings.TrimSpace(result.Title) == "" {
|
||||||
|
result.MissingFields = append(result.MissingFields, "title")
|
||||||
|
result.FollowUpQuestions = append(result.FollowUpQuestions, "Please provide a concise ticket title that clearly summarizes the issue.")
|
||||||
|
}
|
||||||
|
if !hasSufficientIssueContext(input, result.Description) {
|
||||||
|
result.MissingFields = append(result.MissingFields, "issue")
|
||||||
|
result.FollowUpQuestions = append(result.FollowUpQuestions, "Please provide the specific issue, error message, or request so I can prepare the ticket.")
|
||||||
|
}
|
||||||
|
result.Ready = result.Title != "" && result.Description != "" && len(result.MissingFields) == 0
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildDraftTitle(conversation models.Conversation, input PrepareTicketDraftInput) string {
|
||||||
|
switch {
|
||||||
|
case input.Title != "":
|
||||||
|
return limitText(input.Title, 80)
|
||||||
|
case input.Issue != "":
|
||||||
|
return limitText(input.Issue, 80)
|
||||||
|
case strings.TrimSpace(conversation.LastMessageSummary) != "":
|
||||||
|
return limitText(conversation.LastMessageSummary, 80)
|
||||||
|
default:
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildDraftDescription(conversation models.Conversation, messages []models.Message, input PrepareTicketDraftInput) string {
|
||||||
|
if input.Description != "" {
|
||||||
|
return input.Description
|
||||||
|
}
|
||||||
|
parts := make([]string, 0, 6)
|
||||||
|
if input.Issue != "" {
|
||||||
|
parts = append(parts, "Issue: "+input.Issue)
|
||||||
|
}
|
||||||
|
if input.Impact != "" {
|
||||||
|
parts = append(parts, "Impact: "+input.Impact)
|
||||||
|
}
|
||||||
|
if input.ExpectedOutcome != "" {
|
||||||
|
parts = append(parts, "Requested outcome: "+input.ExpectedOutcome)
|
||||||
|
}
|
||||||
|
if input.CurrentAttempt != "" {
|
||||||
|
parts = append(parts, "Attempts so far: "+input.CurrentAttempt)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(conversation.LastMessageSummary) != "" {
|
||||||
|
parts = append(parts, "Conversation summary: "+strings.TrimSpace(conversation.LastMessageSummary))
|
||||||
|
}
|
||||||
|
if recent := buildRecentMessageDigest(messages); recent != "" {
|
||||||
|
parts = append(parts, "Recent messages: "+recent)
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(strings.Join(parts, "\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func hasSufficientIssueContext(input PrepareTicketDraftInput, description string) bool {
|
||||||
|
if input.Issue != "" || input.Description != "" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return len([]rune(strings.TrimSpace(description))) >= 30
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildConversationFacts(conversation models.Conversation, messages []models.Message) []string {
|
||||||
|
facts := make([]string, 0, 4)
|
||||||
|
if strings.TrimSpace(conversation.LastMessageSummary) != "" {
|
||||||
|
facts = append(facts, "Recent summary: "+strings.TrimSpace(conversation.LastMessageSummary))
|
||||||
|
}
|
||||||
|
if digest := buildRecentMessageDigest(messages); digest != "" {
|
||||||
|
facts = append(facts, "Recent messages: "+digest)
|
||||||
|
}
|
||||||
|
return facts
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildRecentMessageDigest(messages []models.Message) string {
|
||||||
|
if len(messages) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
parts := make([]string, 0, len(messages))
|
||||||
|
for i := range messages {
|
||||||
|
content := strings.TrimSpace(messages[i].Content)
|
||||||
|
if content == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
parts = append(parts, messageSenderLabel(messages[i].SenderType)+":"+limitText(content, 60))
|
||||||
|
}
|
||||||
|
return strings.Join(parts, " | ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func messageSenderLabel(senderType enums.IMSenderType) string {
|
||||||
|
switch senderType {
|
||||||
|
case enums.IMSenderTypeCustomer:
|
||||||
|
return "Customer"
|
||||||
|
case enums.IMSenderTypeAgent:
|
||||||
|
return "Agent"
|
||||||
|
case enums.IMSenderTypeAI:
|
||||||
|
return "AI"
|
||||||
|
default:
|
||||||
|
return "Message"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func limitText(value string, max int) string {
|
||||||
|
value = strings.TrimSpace(value)
|
||||||
|
if max <= 0 {
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
runes := []rune(value)
|
||||||
|
if len(runes) <= max {
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(string(runes[:max])) + "..."
|
||||||
|
}
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestBuildPrepareTicketDraftResult_UsesConversationFallbacks(t *testing.T) {
|
||||||
|
conversation := models.Conversation{
|
||||||
|
LastMessageSummary: "用户反馈企业微信扫码后页面空白,无法进入工作台",
|
||||||
|
}
|
||||||
|
messages := []models.Message{
|
||||||
|
{SenderType: enums.IMSenderTypeCustomer, Content: "扫码登录后一直白屏"},
|
||||||
|
{SenderType: enums.IMSenderTypeAI, Content: "请问是否有报错提示"},
|
||||||
|
}
|
||||||
|
|
||||||
|
got := buildPrepareTicketDraftResult(conversation, messages, PrepareTicketDraftInput{
|
||||||
|
Impact: "无法进入后台处理客户消息",
|
||||||
|
ExpectedOutcome: "恢复正常登录",
|
||||||
|
})
|
||||||
|
|
||||||
|
if got.Title == "" {
|
||||||
|
t.Fatalf("expected draft title to be generated")
|
||||||
|
}
|
||||||
|
if got.Description == "" {
|
||||||
|
t.Fatalf("expected draft description to be generated")
|
||||||
|
}
|
||||||
|
if !got.Ready {
|
||||||
|
t.Fatalf("expected conversation summary and recent messages to be enough, got %#v", got)
|
||||||
|
}
|
||||||
|
if len(got.ConversationFacts) == 0 {
|
||||||
|
t.Fatalf("expected conversation facts to be populated")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildPrepareTicketDraftResult_ReadyWithExplicitIssue(t *testing.T) {
|
||||||
|
conversation := models.Conversation{
|
||||||
|
LastMessageSummary: "用户反馈连续支付失败",
|
||||||
|
}
|
||||||
|
|
||||||
|
got := buildPrepareTicketDraftResult(conversation, nil, PrepareTicketDraftInput{
|
||||||
|
Issue: "用户连续三次支付订单失败,页面提示网络异常。",
|
||||||
|
ExpectedOutcome: "希望尽快恢复支付并完成下单。",
|
||||||
|
CurrentAttempt: "已尝试切换网络和刷新页面,问题仍存在。",
|
||||||
|
})
|
||||||
|
|
||||||
|
if !got.Ready {
|
||||||
|
t.Fatalf("expected draft to be ready, got %#v", got)
|
||||||
|
}
|
||||||
|
if got.Title == "" || got.Description == "" {
|
||||||
|
t.Fatalf("expected title and description to be populated, got %#v", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
80
internal/ai/runtime/graphs/triage_service_request_graph.go
Normal file
80
internal/ai/runtime/graphs/triage_service_request_graph.go
Normal file
@@ -0,0 +1,80 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
type TriageServiceRequestInput struct {
|
||||||
|
Goal string `json:"goal"`
|
||||||
|
ObservedIssue string `json:"observedIssue"`
|
||||||
|
NeedTicket bool `json:"needTicket"`
|
||||||
|
NeedHumanHandoff bool `json:"needHumanHandoff"`
|
||||||
|
AdditionalContext string `json:"additionalContext"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type TriageServiceRequestResult struct {
|
||||||
|
Analysis AnalyzeConversationResult `json:"analysis"`
|
||||||
|
TicketDraft *PrepareTicketDraftResult `json:"ticketDraft,omitempty"`
|
||||||
|
RecommendedAction string `json:"recommendedAction"`
|
||||||
|
Ready bool `json:"ready"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type TriageServiceRequestGraph struct {
|
||||||
|
conversation models.Conversation
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewTriageServiceRequestGraph(conversation models.Conversation) *TriageServiceRequestGraph {
|
||||||
|
return &TriageServiceRequestGraph{conversation: conversation}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *TriageServiceRequestGraph) Run(_ context.Context, argumentsInJSON string) (string, error) {
|
||||||
|
input, err := g.parseInput(argumentsInJSON)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
messages, _, _ := services.MessageService.FindByConversationIDCursor(g.conversation.ID, 0, 8, "", "")
|
||||||
|
analysis := buildAnalyzeConversationResult(g.conversation, messages, AnalyzeConversationInput{
|
||||||
|
Goal: input.Goal,
|
||||||
|
ObservedIssue: input.ObservedIssue,
|
||||||
|
NeedTicket: input.NeedTicket,
|
||||||
|
NeedHumanHandoff: input.NeedHumanHandoff,
|
||||||
|
AdditionalContext: input.AdditionalContext,
|
||||||
|
})
|
||||||
|
result := TriageServiceRequestResult{
|
||||||
|
Analysis: analysis,
|
||||||
|
RecommendedAction: analysis.RecommendedNextAction,
|
||||||
|
Ready: analysis.RecommendedNextAction == "continue_answering" || analysis.RecommendedNextAction == "handoff_to_human",
|
||||||
|
}
|
||||||
|
if analysis.RecommendedNextAction == "prepare_ticket" {
|
||||||
|
draft := buildPrepareTicketDraftResult(g.conversation, messages, PrepareTicketDraftInput{
|
||||||
|
Issue: input.ObservedIssue,
|
||||||
|
})
|
||||||
|
result.TicketDraft = &draft
|
||||||
|
result.Ready = draft.Ready
|
||||||
|
}
|
||||||
|
buf, err := json.Marshal(result)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(buf), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (g *TriageServiceRequestGraph) parseInput(argumentsInJSON string) (TriageServiceRequestInput, error) {
|
||||||
|
var input TriageServiceRequestInput
|
||||||
|
if strings.TrimSpace(argumentsInJSON) == "" {
|
||||||
|
return input, nil
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal([]byte(argumentsInJSON), &input); err != nil {
|
||||||
|
return input, fmt.Errorf("invalid triage service request arguments: %w", err)
|
||||||
|
}
|
||||||
|
input.Goal = strings.TrimSpace(input.Goal)
|
||||||
|
input.ObservedIssue = strings.TrimSpace(input.ObservedIssue)
|
||||||
|
input.AdditionalContext = strings.TrimSpace(input.AdditionalContext)
|
||||||
|
return input, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
package graphs
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestTriageServiceRequestResult_PrepareTicket(t *testing.T) {
|
||||||
|
conversation := models.Conversation{
|
||||||
|
LastMessageSummary: "用户要求建单跟进支付失败问题",
|
||||||
|
}
|
||||||
|
messages := []models.Message{
|
||||||
|
{SenderType: enums.IMSenderTypeCustomer, Content: "帮我建个工单,支付一直失败"},
|
||||||
|
}
|
||||||
|
|
||||||
|
analysis := buildAnalyzeConversationResult(conversation, messages, AnalyzeConversationInput{
|
||||||
|
NeedTicket: true,
|
||||||
|
})
|
||||||
|
if analysis.RecommendedNextAction != "prepare_ticket" {
|
||||||
|
t.Fatalf("expected prepare_ticket, got %q", analysis.RecommendedNextAction)
|
||||||
|
}
|
||||||
|
|
||||||
|
draft := buildPrepareTicketDraftResult(conversation, messages, PrepareTicketDraftInput{})
|
||||||
|
if draft.Title == "" || draft.Description == "" {
|
||||||
|
t.Fatalf("expected draft to be populated, got %#v", draft)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTriageServiceRequestResult_Handoff(t *testing.T) {
|
||||||
|
conversation := models.Conversation{
|
||||||
|
LastMessageSummary: "用户要求人工处理扣费投诉",
|
||||||
|
}
|
||||||
|
messages := []models.Message{
|
||||||
|
{SenderType: enums.IMSenderTypeCustomer, Content: "我要投诉并转人工,你们重复扣费了"},
|
||||||
|
}
|
||||||
|
|
||||||
|
analysis := buildAnalyzeConversationResult(conversation, messages, AnalyzeConversationInput{
|
||||||
|
NeedHumanHandoff: true,
|
||||||
|
})
|
||||||
|
if analysis.RecommendedNextAction != "handoff_to_human" {
|
||||||
|
t.Fatalf("expected handoff_to_human, got %q", analysis.RecommendedNextAction)
|
||||||
|
}
|
||||||
|
}
|
||||||
82
internal/ai/runtime/instruction/assembler.go
Normal file
82
internal/ai/runtime/instruction/assembler.go
Normal file
@@ -0,0 +1,82 @@
|
|||||||
|
package instruction
|
||||||
|
|
||||||
|
import "strings"
|
||||||
|
|
||||||
|
type Assembler struct{}
|
||||||
|
|
||||||
|
type AssemblerInput struct {
|
||||||
|
AgentInstruction string
|
||||||
|
SkillInstruction string
|
||||||
|
ToolAppendices []string
|
||||||
|
}
|
||||||
|
|
||||||
|
type AssemblySummary struct {
|
||||||
|
SectionTitles []string
|
||||||
|
HasAgentRule bool
|
||||||
|
HasSkillRule bool
|
||||||
|
HasToolRule bool
|
||||||
|
}
|
||||||
|
|
||||||
|
type AssemblyResult struct {
|
||||||
|
Text string
|
||||||
|
Summary AssemblySummary
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewAssembler() *Assembler {
|
||||||
|
return &Assembler{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Assembler) Build(input AssemblerInput) string {
|
||||||
|
return a.Assemble(input).Text
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Assembler) Assemble(input AssemblerInput) AssemblyResult {
|
||||||
|
parts := make([]string, 0, 3)
|
||||||
|
summary := AssemblySummary{SectionTitles: make([]string, 0, 3)}
|
||||||
|
if agentInstruction := strings.TrimSpace(input.AgentInstruction); agentInstruction != "" {
|
||||||
|
parts = append(parts, buildInstructionSection("Agent Rules", agentInstruction))
|
||||||
|
summary.HasAgentRule = true
|
||||||
|
summary.SectionTitles = append(summary.SectionTitles, "Agent Rules")
|
||||||
|
}
|
||||||
|
if skillInstruction := strings.TrimSpace(input.SkillInstruction); skillInstruction != "" {
|
||||||
|
parts = append(parts, buildInstructionSection("Current Skill Context", skillInstruction))
|
||||||
|
summary.HasSkillRule = true
|
||||||
|
summary.SectionTitles = append(summary.SectionTitles, "Current Skill Context")
|
||||||
|
}
|
||||||
|
if appendix := buildToolAppendix(input.ToolAppendices); appendix != "" {
|
||||||
|
parts = append(parts, buildInstructionSection("Tool Supplementary Rules", appendix))
|
||||||
|
summary.HasToolRule = true
|
||||||
|
summary.SectionTitles = append(summary.SectionTitles, "Tool Supplementary Rules")
|
||||||
|
}
|
||||||
|
return AssemblyResult{
|
||||||
|
Text: strings.TrimSpace(strings.Join(parts, "\n\n")),
|
||||||
|
Summary: summary,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildInstructionSection(title, body string) string {
|
||||||
|
title = strings.TrimSpace(title)
|
||||||
|
body = strings.TrimSpace(body)
|
||||||
|
if body == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if title == "" {
|
||||||
|
return body
|
||||||
|
}
|
||||||
|
return title + ":\n" + body
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildToolAppendix(input []string) string {
|
||||||
|
if len(input) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
parts := make([]string, 0, len(input))
|
||||||
|
for _, item := range input {
|
||||||
|
item = strings.TrimSpace(item)
|
||||||
|
if item == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
parts = append(parts, item)
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(strings.Join(parts, "\n\n"))
|
||||||
|
}
|
||||||
36
internal/ai/runtime/instruction/assembler_test.go
Normal file
36
internal/ai/runtime/instruction/assembler_test.go
Normal file
@@ -0,0 +1,36 @@
|
|||||||
|
package instruction
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAssemblerRespectsProvidedSources(t *testing.T) {
|
||||||
|
result := NewAssembler().Assemble(AssemblerInput{
|
||||||
|
AgentInstruction: "agent-rule",
|
||||||
|
SkillInstruction: "skill-rule",
|
||||||
|
ToolAppendices: []string{"tool-rule-1", "tool-rule-2"},
|
||||||
|
})
|
||||||
|
if !strings.Contains(result.Text, "Agent Rules:\nagent-rule") {
|
||||||
|
t.Fatalf("missing agent instruction: %s", result.Text)
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.Text, "Current Skill Context:\nskill-rule") {
|
||||||
|
t.Fatalf("missing skill instruction: %s", result.Text)
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.Text, "Tool Supplementary Rules:\ntool-rule-1") {
|
||||||
|
t.Fatalf("missing tool appendix: %s", result.Text)
|
||||||
|
}
|
||||||
|
if !result.Summary.HasAgentRule || !result.Summary.HasSkillRule || !result.Summary.HasToolRule {
|
||||||
|
t.Fatalf("unexpected summary: %#v", result.Summary)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAssemblerReturnsEmptyTextWhenInputIsEmpty(t *testing.T) {
|
||||||
|
result := NewAssembler().Assemble(AssemblerInput{})
|
||||||
|
if result.Text != "" {
|
||||||
|
t.Fatalf("expected empty assembled text, got: %s", result.Text)
|
||||||
|
}
|
||||||
|
if len(result.Summary.SectionTitles) != 0 || result.Summary.HasAgentRule || result.Summary.HasSkillRule || result.Summary.HasToolRule {
|
||||||
|
t.Fatalf("expected empty summary, got %#v", result.Summary)
|
||||||
|
}
|
||||||
|
}
|
||||||
89
internal/ai/runtime/instruction/helpers.go
Normal file
89
internal/ai/runtime/instruction/helpers.go
Normal file
@@ -0,0 +1,89 @@
|
|||||||
|
package instruction
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
runtimetooling "agent-desk/internal/ai/runtime/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
)
|
||||||
|
|
||||||
|
func BuildSelectedSkillActivationInstruction(skill *models.SkillDefinition) string {
|
||||||
|
if skill == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
lines := []string{
|
||||||
|
"Active Specialized Skill:",
|
||||||
|
fmt.Sprintf("- id: %d", skill.ID),
|
||||||
|
fmt.Sprintf("- name: %s", strings.TrimSpace(skill.Name)),
|
||||||
|
}
|
||||||
|
if desc := strings.TrimSpace(skill.Description); desc != "" {
|
||||||
|
lines = append(lines, fmt.Sprintf("- description: %s", desc))
|
||||||
|
}
|
||||||
|
lines = append(lines, "", "Execution Requirements:", "- Prioritize handling inquiries within this skill's scope this turn.", fmt.Sprintf("- When specialized handling details are needed, invoke %s to load skill documentation before proceeding.", toolx.BuiltinSkill.Name), "- If key information is missing, ask the customer for clarification first.", "- Do not invoke tools unauthorized for this skill.")
|
||||||
|
return strings.TrimSpace(strings.Join(lines, "\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func BuildSelectedSkillDocument(skill *models.SkillDefinition, toolDefinitions []runtimetooling.MCPToolDefinition) string {
|
||||||
|
return BuildSkillDocument(skill, toolDefinitions)
|
||||||
|
}
|
||||||
|
|
||||||
|
func BuildSkillDocument(skill *models.SkillDefinition, toolDefinitions []runtimetooling.MCPToolDefinition) string {
|
||||||
|
if skill == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
lines := []string{
|
||||||
|
"Active Specialized Skill:",
|
||||||
|
fmt.Sprintf("- id: %d", skill.ID),
|
||||||
|
fmt.Sprintf("- name: %s", strings.TrimSpace(skill.Name)),
|
||||||
|
}
|
||||||
|
if desc := strings.TrimSpace(skill.Description); desc != "" {
|
||||||
|
lines = append(lines, fmt.Sprintf("- description: %s", desc))
|
||||||
|
}
|
||||||
|
if content := strings.TrimSpace(skill.Instruction); content != "" {
|
||||||
|
lines = append(lines, "", "Skill Documentation:", content)
|
||||||
|
}
|
||||||
|
if examples := parseJSONStringArray(skill.Examples); len(examples) > 0 {
|
||||||
|
lines = append(lines, "", "Typical Example Inquiries:")
|
||||||
|
for _, item := range examples {
|
||||||
|
lines = append(lines, "- "+item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(toolDefinitions) > 0 {
|
||||||
|
lines = append(lines, "", "Authorized Tools for this Skill:")
|
||||||
|
for _, item := range toolDefinitions {
|
||||||
|
if strings.TrimSpace(item.ToolCode) == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
line := "- " + strings.TrimSpace(item.ToolCode)
|
||||||
|
if title := strings.TrimSpace(item.Title); title != "" {
|
||||||
|
line += " | " + title
|
||||||
|
}
|
||||||
|
lines = append(lines, line)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
lines = append(lines, "", "执行要求:", "- 优先遵循该技能说明完成任务。", "- 如果关键信息不足,先向用户追问。", "- 不得调用当前技能未授权的工具。")
|
||||||
|
return strings.TrimSpace(strings.Join(lines, "\n"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseJSONStringArray(raw string) []string {
|
||||||
|
raw = strings.TrimSpace(raw)
|
||||||
|
if raw == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var ret []string
|
||||||
|
if err := json.Unmarshal([]byte(raw), &ret); err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
out := make([]string, 0, len(ret))
|
||||||
|
for _, item := range ret {
|
||||||
|
item = strings.TrimSpace(item)
|
||||||
|
if item == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, item)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
36
internal/ai/runtime/instruction/providers.go
Normal file
36
internal/ai/runtime/instruction/providers.go
Normal file
@@ -0,0 +1,36 @@
|
|||||||
|
package instruction
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/ai/runtime/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
)
|
||||||
|
|
||||||
|
type ToolAppendixProvider struct{}
|
||||||
|
|
||||||
|
func NewToolAppendixProvider() *ToolAppendixProvider {
|
||||||
|
return &ToolAppendixProvider{}
|
||||||
|
}
|
||||||
|
|
||||||
|
type SkillInstructionProvider struct{}
|
||||||
|
|
||||||
|
func NewSkillInstructionProvider() *SkillInstructionProvider {
|
||||||
|
return &SkillInstructionProvider{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *SkillInstructionProvider) Resolve(selectedSkill *models.SkillDefinition) string {
|
||||||
|
return BuildSelectedSkillActivationInstruction(selectedSkill)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *ToolAppendixProvider) Build(toolDefinitions []tooling.MCPToolDefinition, extraToolCodes map[string]string) []string {
|
||||||
|
appendixParts := make([]string, 0, 1)
|
||||||
|
toolCodes := make([]string, 0, len(toolDefinitions)+len(extraToolCodes))
|
||||||
|
for _, item := range toolDefinitions {
|
||||||
|
toolCodes = append(toolCodes, item.ToolCode)
|
||||||
|
}
|
||||||
|
for _, item := range extraToolCodes {
|
||||||
|
toolCodes = append(toolCodes, item)
|
||||||
|
}
|
||||||
|
appendixParts = append(appendixParts, toolx.BuildToolAppendicesForCodes(len(toolDefinitions) > 0, toolCodes)...)
|
||||||
|
return appendixParts
|
||||||
|
}
|
||||||
60
internal/ai/runtime/instruction/service.go
Normal file
60
internal/ai/runtime/instruction/service.go
Normal file
@@ -0,0 +1,60 @@
|
|||||||
|
package instruction
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
runtimetooling "agent-desk/internal/ai/runtime/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Service struct {
|
||||||
|
assembler *Assembler
|
||||||
|
skillInstructionProvider *SkillInstructionProvider
|
||||||
|
toolAppendixProvider *ToolAppendixProvider
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewService(
|
||||||
|
assembler *Assembler,
|
||||||
|
skillProvider *SkillInstructionProvider,
|
||||||
|
toolProvider *ToolAppendixProvider,
|
||||||
|
) *Service {
|
||||||
|
if assembler == nil {
|
||||||
|
assembler = NewAssembler()
|
||||||
|
}
|
||||||
|
if skillProvider == nil {
|
||||||
|
skillProvider = NewSkillInstructionProvider()
|
||||||
|
}
|
||||||
|
if toolProvider == nil {
|
||||||
|
toolProvider = NewToolAppendixProvider()
|
||||||
|
}
|
||||||
|
return &Service{
|
||||||
|
assembler: assembler,
|
||||||
|
skillInstructionProvider: skillProvider,
|
||||||
|
toolAppendixProvider: toolProvider,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Service) Build(
|
||||||
|
aiAgent models.AIAgent,
|
||||||
|
selectedSkill *models.SkillDefinition,
|
||||||
|
toolDefinitions []runtimetooling.MCPToolDefinition,
|
||||||
|
extraToolCodes map[string]string,
|
||||||
|
) AssemblyResult {
|
||||||
|
skillInstruction := ""
|
||||||
|
toolAppendices := make([]string, 0)
|
||||||
|
if s != nil && s.skillInstructionProvider != nil {
|
||||||
|
skillInstruction = s.skillInstructionProvider.Resolve(selectedSkill)
|
||||||
|
}
|
||||||
|
if s != nil && s.toolAppendixProvider != nil {
|
||||||
|
toolAppendices = s.toolAppendixProvider.Build(toolDefinitions, extraToolCodes)
|
||||||
|
}
|
||||||
|
assembler := NewAssembler()
|
||||||
|
if s != nil && s.assembler != nil {
|
||||||
|
assembler = s.assembler
|
||||||
|
}
|
||||||
|
return assembler.Assemble(AssemblerInput{
|
||||||
|
AgentInstruction: strings.TrimSpace(aiAgent.SystemPrompt),
|
||||||
|
SkillInstruction: skillInstruction,
|
||||||
|
ToolAppendices: toolAppendices,
|
||||||
|
})
|
||||||
|
}
|
||||||
79
internal/ai/runtime/readtools/graph_executor.go
Normal file
79
internal/ai/runtime/readtools/graph_executor.go
Normal file
@@ -0,0 +1,79 @@
|
|||||||
|
// Package readtools executes deterministic, read-only graph tools through the
|
||||||
|
// shared Tool Registry boundary.
|
||||||
|
package readtools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/runtime/graphs"
|
||||||
|
"agent-desk/internal/ai/runtime/retrievers"
|
||||||
|
aitooling "agent-desk/internal/ai/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
)
|
||||||
|
|
||||||
|
func ExecuteGraphTool(ctx context.Context, conversation models.Conversation, toolCode string, arguments map[string]any, policy aitooling.Policy) (aitooling.Definition, string, error) {
|
||||||
|
toolCode = toolx.NormalizeToolCodeAlias(strings.TrimSpace(toolCode))
|
||||||
|
if toolCode != toolx.GraphTriageServiceRequest.Code && toolCode != toolx.GraphAnalyzeConversation.Code && toolCode != toolx.GraphPrepareTicketDraft.Code {
|
||||||
|
return aitooling.Definition{}, "", fmt.Errorf("tool is not a graph read tool")
|
||||||
|
}
|
||||||
|
definition, err := aitooling.DefaultRegistry.Resolve(toolCode)
|
||||||
|
if err != nil {
|
||||||
|
return aitooling.Definition{}, "", err
|
||||||
|
}
|
||||||
|
if err := aitooling.DefaultPolicyGuard.Authorize(aitooling.Invocation{
|
||||||
|
Definition: definition,
|
||||||
|
Arguments: arguments,
|
||||||
|
Policy: policy,
|
||||||
|
}); err != nil {
|
||||||
|
return definition, "", err
|
||||||
|
}
|
||||||
|
if definition.TimeoutMS > 0 {
|
||||||
|
var cancel context.CancelFunc
|
||||||
|
ctx, cancel = context.WithTimeout(ctx, time.Duration(definition.TimeoutMS)*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(arguments)
|
||||||
|
if err != nil {
|
||||||
|
return definition, "", err
|
||||||
|
}
|
||||||
|
switch toolCode {
|
||||||
|
case toolx.GraphTriageServiceRequest.Code:
|
||||||
|
result, err := graphs.NewTriageServiceRequestGraph(conversation).Run(ctx, string(data))
|
||||||
|
return definition, result, err
|
||||||
|
case toolx.GraphAnalyzeConversation.Code:
|
||||||
|
result, err := graphs.NewAnalyzeConversationGraph(conversation).Run(ctx, string(data))
|
||||||
|
return definition, result, err
|
||||||
|
default:
|
||||||
|
result, err := graphs.NewPrepareTicketDraftGraph(conversation).Run(ctx, string(data))
|
||||||
|
return definition, result, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// RetrieveKnowledge executes the built-in knowledge tool after the same
|
||||||
|
// registry policy and timeout checks used by graph tools.
|
||||||
|
func RetrieveKnowledge(ctx context.Context, agent models.AIAgent, knowledgeBaseIDs []int64, query string, policy aitooling.Policy) (aitooling.Definition, *retrievers.KnowledgeRetrieveResult, error) {
|
||||||
|
definition, err := aitooling.DefaultRegistry.Resolve(toolx.BuiltinKnowledgeRetrieve.Code)
|
||||||
|
if err != nil {
|
||||||
|
return aitooling.Definition{}, nil, err
|
||||||
|
}
|
||||||
|
arguments := map[string]any{"query": strings.TrimSpace(query), "knowledgeBaseIds": knowledgeBaseIDs}
|
||||||
|
if err := aitooling.DefaultPolicyGuard.Authorize(aitooling.Invocation{
|
||||||
|
Definition: definition,
|
||||||
|
Arguments: arguments,
|
||||||
|
Policy: policy,
|
||||||
|
}); err != nil {
|
||||||
|
return definition, nil, err
|
||||||
|
}
|
||||||
|
if definition.TimeoutMS > 0 {
|
||||||
|
var cancel context.CancelFunc
|
||||||
|
ctx, cancel = context.WithTimeout(ctx, time.Duration(definition.TimeoutMS)*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
}
|
||||||
|
result, err := retrievers.NewKnowledgeRetriever(agent, knowledgeBaseIDs).RetrieveContext(ctx, strings.TrimSpace(query))
|
||||||
|
return definition, result, err
|
||||||
|
}
|
||||||
26
internal/ai/runtime/readtools/graph_executor_test.go
Normal file
26
internal/ai/runtime/readtools/graph_executor_test.go
Normal file
@@ -0,0 +1,26 @@
|
|||||||
|
package readtools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
aitooling "agent-desk/internal/ai/tooling"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestExecuteGraphToolRejectsDisallowedToolBeforeGraphExecution(t *testing.T) {
|
||||||
|
definition, _, err := ExecuteGraphTool(context.Background(), models.Conversation{}, toolx.GraphAnalyzeConversation.Code, map[string]any{
|
||||||
|
"observedIssue": "需要分析的问题",
|
||||||
|
}, aitooling.Policy{
|
||||||
|
AllowedToolCodes: []string{toolx.GraphPrepareTicketDraft.Code},
|
||||||
|
AllowedRiskLevels: []string{aitooling.RiskLevelRead},
|
||||||
|
Confirmed: true,
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected policy guard to reject the graph tool")
|
||||||
|
}
|
||||||
|
if definition.Code != toolx.GraphAnalyzeConversation.Code {
|
||||||
|
t.Fatalf("definition code = %q, want %q", definition.Code, toolx.GraphAnalyzeConversation.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
103
internal/ai/runtime/registry/registry.go
Normal file
103
internal/ai/runtime/registry/registry.go
Normal file
@@ -0,0 +1,103 @@
|
|||||||
|
package registry
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
|
||||||
|
einotool "github.com/cloudwego/eino/components/tool"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Registry struct {
|
||||||
|
tools []Tool
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewRegistry(tools ...Tool) *Registry {
|
||||||
|
return &Registry{
|
||||||
|
tools: tools,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Registry) Resolve(ctx Context) (*ToolSet, error) {
|
||||||
|
ret := &ToolSet{
|
||||||
|
StaticTools: make([]einotool.BaseTool, 0, len(r.tools)),
|
||||||
|
StaticToolCodes: make(map[string]string),
|
||||||
|
StaticToolMetadata: make(map[string]ToolMetadata),
|
||||||
|
}
|
||||||
|
allowedToolCodes := makeAllowedToolCodeSet(ctx.AllowedToolCodes)
|
||||||
|
for _, toolDef := range r.tools {
|
||||||
|
if toolDef == nil || !toolDef.Enabled(ctx) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
spec := toolDef.Spec()
|
||||||
|
toolCode := strings.TrimSpace(spec.Code)
|
||||||
|
if toolCode == "" {
|
||||||
|
toolCode = strings.TrimSpace(toolDef.Code())
|
||||||
|
}
|
||||||
|
if len(allowedToolCodes) > 0 && !isAllowedToolCode(toolCode, allowedToolCodes) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
tool, err := toolDef.Build(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if tool == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
toolName := strings.TrimSpace(spec.Name)
|
||||||
|
if toolName == "" {
|
||||||
|
toolName = strings.TrimSpace(toolDef.Name())
|
||||||
|
}
|
||||||
|
if toolName == "" || toolCode == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
ret.StaticTools = append(ret.StaticTools, tool)
|
||||||
|
ret.StaticToolCodes[toolName] = toolCode
|
||||||
|
resolvedMetadata := toolx.ResolveToolMetadata(toolCode, toolName)
|
||||||
|
if strings.TrimSpace(resolvedMetadata.ServerCode) == "" {
|
||||||
|
resolvedMetadata.ServerCode = strings.TrimSpace(spec.ServerCode)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(resolvedMetadata.ToolName) == "" {
|
||||||
|
resolvedMetadata.ToolName = toolName
|
||||||
|
}
|
||||||
|
if resolvedMetadata.SourceType == "" {
|
||||||
|
resolvedMetadata.SourceType = spec.SourceType
|
||||||
|
}
|
||||||
|
ret.StaticToolMetadata[toolName] = ToolMetadata{
|
||||||
|
ToolCode: resolvedMetadata.ToolCode,
|
||||||
|
ServerCode: resolvedMetadata.ServerCode,
|
||||||
|
ToolName: resolvedMetadata.ToolName,
|
||||||
|
SourceType: resolvedMetadata.SourceType,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ret, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func isAllowedToolCode(toolCode string, allowedToolCodes map[string]struct{}) bool {
|
||||||
|
if len(allowedToolCodes) == 0 {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
toolCode = toolx.NormalizeToolCodeAlias(strings.TrimSpace(toolCode))
|
||||||
|
if _, ok := allowedToolCodes[toolCode]; ok {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if toolx.IsAlwaysAllowedToolCode(toolCode) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return toolx.IsImpliedAllowedToolCode(toolCode, allowedToolCodes)
|
||||||
|
}
|
||||||
|
|
||||||
|
func makeAllowedToolCodeSet(input []string) map[string]struct{} {
|
||||||
|
if len(input) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
ret := make(map[string]struct{}, len(input))
|
||||||
|
for _, item := range input {
|
||||||
|
item = toolx.NormalizeToolCodeAlias(strings.TrimSpace(item))
|
||||||
|
if item == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
ret[item] = struct{}{}
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
80
internal/ai/runtime/registry/registry_test.go
Normal file
80
internal/ai/runtime/registry/registry_test.go
Normal file
@@ -0,0 +1,80 @@
|
|||||||
|
package registry_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/runtime/registry"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
|
||||||
|
einotool "github.com/cloudwego/eino/components/tool"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
type stubTool struct {
|
||||||
|
name string
|
||||||
|
code string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t stubTool) Spec() toolx.ToolSpec {
|
||||||
|
return toolx.ToolSpec{
|
||||||
|
Code: t.code,
|
||||||
|
Name: t.name,
|
||||||
|
ServerCode: toolx.GraphCreateTicketConfirm.ServerCode,
|
||||||
|
SourceType: toolx.GraphCreateTicketConfirm.SourceType,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t stubTool) Name() string { return t.name }
|
||||||
|
func (t stubTool) Code() string { return t.code }
|
||||||
|
func (t stubTool) Enabled(registry.Context) bool {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
func (t stubTool) Build(registry.Context) (einotool.BaseTool, error) {
|
||||||
|
return stubBaseTool{name: t.name}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type stubBaseTool struct {
|
||||||
|
name string
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t stubBaseTool) Info(context.Context) (*schema.ToolInfo, error) {
|
||||||
|
return &schema.ToolInfo{Name: t.name}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveBuildsStaticToolMetadata(t *testing.T) {
|
||||||
|
r := registry.NewRegistry(stubTool{
|
||||||
|
name: toolx.GraphCreateTicketConfirm.Name,
|
||||||
|
code: toolx.GraphCreateTicketConfirm.Code,
|
||||||
|
})
|
||||||
|
toolSet, err := r.Resolve(registry.Context{
|
||||||
|
Conversation: models.Conversation{ID: 1},
|
||||||
|
AIAgent: models.AIAgent{ID: 1},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("resolve returned error: %v", err)
|
||||||
|
}
|
||||||
|
if toolSet == nil {
|
||||||
|
t.Fatalf("expected tool set")
|
||||||
|
}
|
||||||
|
if len(toolSet.StaticToolMetadata) != 1 {
|
||||||
|
t.Fatalf("expected 1 metadata item, got %d", len(toolSet.StaticToolMetadata))
|
||||||
|
}
|
||||||
|
item, ok := toolSet.StaticToolMetadata[toolx.GraphCreateTicketConfirm.Name]
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("missing metadata for %s", toolx.GraphCreateTicketConfirm.Name)
|
||||||
|
}
|
||||||
|
if item.ToolCode != toolx.GraphCreateTicketConfirm.Code {
|
||||||
|
t.Fatalf("unexpected tool code: %s", item.ToolCode)
|
||||||
|
}
|
||||||
|
if item.ServerCode != toolx.GraphCreateTicketConfirm.ServerCode {
|
||||||
|
t.Fatalf("unexpected server code: %s", item.ServerCode)
|
||||||
|
}
|
||||||
|
if item.ToolName != toolx.GraphCreateTicketConfirm.Name {
|
||||||
|
t.Fatalf("unexpected tool name: %s", item.ToolName)
|
||||||
|
}
|
||||||
|
if item.SourceType != toolx.GraphCreateTicketConfirm.SourceType {
|
||||||
|
t.Fatalf("unexpected source type: %s", item.SourceType)
|
||||||
|
}
|
||||||
|
}
|
||||||
44
internal/ai/runtime/registry/types.go
Normal file
44
internal/ai/runtime/registry/types.go
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
package registry
|
||||||
|
|
||||||
|
import (
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/toolx"
|
||||||
|
|
||||||
|
einotool "github.com/cloudwego/eino/components/tool"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Context struct {
|
||||||
|
Conversation models.Conversation
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
AIConfig models.AIConfig
|
||||||
|
UserMessage models.Message
|
||||||
|
AllowedToolCodes []string
|
||||||
|
}
|
||||||
|
|
||||||
|
type ToolMetadata struct {
|
||||||
|
ToolCode string
|
||||||
|
ServerCode string
|
||||||
|
ToolName string
|
||||||
|
SourceType enums.ToolSourceType
|
||||||
|
}
|
||||||
|
|
||||||
|
// ToolSet 描述当前运行时可直接挂载到 ToolsNode 上的固定工具集合。
|
||||||
|
//
|
||||||
|
// 这里不承载通过 tool_search 动态暴露的 MCP 工具;动态工具仍由 engine 层单独装配。
|
||||||
|
type ToolSet struct {
|
||||||
|
// StaticTools 为固定可见工具实例,例如 Graph Tool。
|
||||||
|
StaticTools []einotool.BaseTool
|
||||||
|
// StaticToolCodes 为固定工具的 modelName -> toolCode 映射,用于 trace 和运行日志归因。
|
||||||
|
StaticToolCodes map[string]string
|
||||||
|
// StaticToolMetadata 为固定工具的 modelName -> metadata 映射,用于 trace、运行日志和后续装配。
|
||||||
|
StaticToolMetadata map[string]ToolMetadata
|
||||||
|
}
|
||||||
|
|
||||||
|
type Tool interface {
|
||||||
|
Spec() toolx.ToolSpec
|
||||||
|
Name() string
|
||||||
|
Code() string
|
||||||
|
Enabled(ctx Context) bool
|
||||||
|
Build(ctx Context) (einotool.BaseTool, error)
|
||||||
|
}
|
||||||
84
internal/ai/runtime/reply_commit_service.go
Normal file
84
internal/ai/runtime/reply_commit_service.go
Normal file
@@ -0,0 +1,84 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
aitooling "agent-desk/internal/ai/tooling"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/dto"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
type replyCommitService struct{}
|
||||||
|
|
||||||
|
type replyCommitInput struct {
|
||||||
|
Conversation models.Conversation
|
||||||
|
Message models.Message
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
ReplyText string
|
||||||
|
ClientPrefix string
|
||||||
|
WorkflowRunID int64
|
||||||
|
IncrementRound bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func newReplyCommitService() *replyCommitService {
|
||||||
|
return &replyCommitService{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *replyCommitService) SendAIReply(input replyCommitInput) (*models.Message, error) {
|
||||||
|
replyText, err := aitooling.NormalizeCustomerReply(input.ReplyText)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
replyMessage, err := svc.MessageService.SendAIMessageWithRequestIDAndWorkflowRunID(
|
||||||
|
input.Conversation.ID,
|
||||||
|
input.AIAgent.ID,
|
||||||
|
fmt.Sprintf("%s_%d", strings.TrimSpace(input.ClientPrefix), input.Message.ID),
|
||||||
|
enums.IMMessageTypeText,
|
||||||
|
replyText,
|
||||||
|
"",
|
||||||
|
s.buildAIPrincipal(input.AIAgent),
|
||||||
|
input.Message.RequestID,
|
||||||
|
input.WorkflowRunID,
|
||||||
|
)
|
||||||
|
if err != nil || !input.IncrementRound {
|
||||||
|
return replyMessage, err
|
||||||
|
}
|
||||||
|
if err := s.IncrementAIReplyRounds(input.Conversation.ID, input.Conversation.AIReplyRounds+1, input.AIAgent.Name); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return replyMessage, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *replyCommitService) CommitAIReply(input replyCommitInput) (*models.Message, error) {
|
||||||
|
input.IncrementRound = true
|
||||||
|
return s.SendAIReply(input)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *replyCommitService) IncrementAIReplyRounds(conversationID int64, nextRounds int, aiAgentName string) error {
|
||||||
|
return repositories.ConversationRepository.Updates(sqls.DB(), conversationID, map[string]any{
|
||||||
|
"ai_reply_rounds": nextRounds,
|
||||||
|
"update_user_id": 0,
|
||||||
|
"update_user_name": strings.TrimSpace(aiAgentName),
|
||||||
|
"updated_at": time.Now(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *replyCommitService) buildAIPrincipal(aiAgent models.AIAgent) *dto.AuthPrincipal {
|
||||||
|
username := "AI"
|
||||||
|
if strings.TrimSpace(aiAgent.Name) != "" {
|
||||||
|
username = aiAgent.Name
|
||||||
|
}
|
||||||
|
return &dto.AuthPrincipal{
|
||||||
|
UserID: 0,
|
||||||
|
Username: username,
|
||||||
|
Nickname: username,
|
||||||
|
}
|
||||||
|
}
|
||||||
140
internal/ai/runtime/reply_commit_service_test.go
Normal file
140
internal/ai/runtime/reply_commit_service_test.go
Normal file
@@ -0,0 +1,140 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
|
||||||
|
"github.com/glebarez/sqlite"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
"gorm.io/gorm"
|
||||||
|
"gorm.io/gorm/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestReplyCommitStoresWorkflowRunIDOnAIMessage(t *testing.T) {
|
||||||
|
db := setupReplyCommitTestDB(t)
|
||||||
|
aiAgent := createReplyCommitTestAIAgent(t, db)
|
||||||
|
conversation := createReplyCommitTestConversation(t, db, aiAgent.ID)
|
||||||
|
|
||||||
|
replyMessage, err := newReplyCommitService().CommitAIReply(replyCommitInput{
|
||||||
|
Conversation: *conversation,
|
||||||
|
Message: models.Message{ID: 101, RequestID: "trace-101"},
|
||||||
|
AIAgent: *aiAgent,
|
||||||
|
ReplyText: "AI reply",
|
||||||
|
ClientPrefix: "ai_reply",
|
||||||
|
WorkflowRunID: 9988,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CommitAIReply() error = %v", err)
|
||||||
|
}
|
||||||
|
if replyMessage == nil {
|
||||||
|
t.Fatalf("expected reply message")
|
||||||
|
}
|
||||||
|
if replyMessage.WorkflowRunID != 9988 {
|
||||||
|
t.Fatalf("replyMessage.WorkflowRunID=%d want 9988", replyMessage.WorkflowRunID)
|
||||||
|
}
|
||||||
|
|
||||||
|
var stored models.Message
|
||||||
|
if err := db.First(&stored, replyMessage.ID).Error; err != nil {
|
||||||
|
t.Fatalf("find reply message: %v", err)
|
||||||
|
}
|
||||||
|
if stored.WorkflowRunID != 9988 {
|
||||||
|
t.Fatalf("stored.WorkflowRunID=%d want 9988", stored.WorkflowRunID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReplyCommitRejectsSensitiveModelOutput(t *testing.T) {
|
||||||
|
db := setupReplyCommitTestDB(t)
|
||||||
|
aiAgent := createReplyCommitTestAIAgent(t, db)
|
||||||
|
conversation := createReplyCommitTestConversation(t, db, aiAgent.ID)
|
||||||
|
_, err := newReplyCommitService().CommitAIReply(replyCommitInput{
|
||||||
|
Conversation: *conversation, Message: models.Message{ID: 102, RequestID: "trace-102"}, AIAgent: *aiAgent,
|
||||||
|
ReplyText: "authorization=Bearer-secret", ClientPrefix: "ai_reply",
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected sensitive model output to be rejected")
|
||||||
|
}
|
||||||
|
var count int64
|
||||||
|
if err := db.Model(&models.Message{}).Where("conversation_id = ?", conversation.ID).Count(&count).Error; err != nil {
|
||||||
|
t.Fatalf("count messages: %v", err)
|
||||||
|
}
|
||||||
|
if count != 0 {
|
||||||
|
t.Fatalf("unexpected message written for rejected output: %d", count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func setupReplyCommitTestDB(t *testing.T) *gorm.DB {
|
||||||
|
t.Helper()
|
||||||
|
dbName := "reply_commit_test_" + strings.NewReplacer("/", "_").Replace(t.Name())
|
||||||
|
db, err := gorm.Open(sqlite.Open("file:"+dbName+"?mode=memory&cache=shared"), &gorm.Config{
|
||||||
|
NamingStrategy: schema.NamingStrategy{
|
||||||
|
TablePrefix: "t_",
|
||||||
|
SingularTable: true,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open sqlite db: %v", err)
|
||||||
|
}
|
||||||
|
sqlDB, err := db.DB()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get sqlite db: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
if err := sqlDB.Close(); err != nil {
|
||||||
|
t.Fatalf("close sqlite db: %v", err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
if err := db.AutoMigrate(
|
||||||
|
&models.AIAgent{},
|
||||||
|
&models.Channel{},
|
||||||
|
&models.ChannelMessageOutbox{},
|
||||||
|
&models.Conversation{},
|
||||||
|
&models.ConversationReadState{},
|
||||||
|
&models.ConversationEventLog{},
|
||||||
|
&models.Message{},
|
||||||
|
); err != nil {
|
||||||
|
t.Fatalf("auto migrate: %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
return db
|
||||||
|
}
|
||||||
|
|
||||||
|
func createReplyCommitTestAIAgent(t *testing.T, db *gorm.DB) *models.AIAgent {
|
||||||
|
t.Helper()
|
||||||
|
now := time.Now()
|
||||||
|
item := &models.AIAgent{
|
||||||
|
Name: "reply-agent",
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
AuditFields: models.AuditFields{
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := db.Create(item).Error; err != nil {
|
||||||
|
t.Fatalf("create ai agent: %v", err)
|
||||||
|
}
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
|
||||||
|
func createReplyCommitTestConversation(t *testing.T, db *gorm.DB, aiAgentID int64) *models.Conversation {
|
||||||
|
t.Helper()
|
||||||
|
now := time.Now()
|
||||||
|
item := &models.Conversation{
|
||||||
|
CustomerID: 1,
|
||||||
|
ChannelID: 11,
|
||||||
|
AIAgentID: aiAgentID,
|
||||||
|
Status: enums.IMConversationStatusAIServing,
|
||||||
|
LastActiveAt: now,
|
||||||
|
AuditFields: models.AuditFields{
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := db.Create(item).Error; err != nil {
|
||||||
|
t.Fatalf("create conversation: %v", err)
|
||||||
|
}
|
||||||
|
return item
|
||||||
|
}
|
||||||
20
internal/ai/runtime/reply_context.go
Normal file
20
internal/ai/runtime/reply_context.go
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
type aiReplyContext struct {
|
||||||
|
Conversation models.Conversation
|
||||||
|
Message models.Message
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
SummaryRef **applicationruntime.RunResult
|
||||||
|
PendingInterrupt *models.ConversationInterrupt
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c aiReplyContext) setSummary(summary *applicationruntime.RunResult) {
|
||||||
|
if c.SummaryRef != nil {
|
||||||
|
*c.SummaryRef = summary
|
||||||
|
}
|
||||||
|
}
|
||||||
64
internal/ai/runtime/reply_eligibility.go
Normal file
64
internal/ai/runtime/reply_eligibility.go
Normal file
@@ -0,0 +1,64 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/binary"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/common/strs"
|
||||||
|
)
|
||||||
|
|
||||||
|
type replyEligibility struct{}
|
||||||
|
|
||||||
|
func newReplyEligibility() *replyEligibility {
|
||||||
|
return &replyEligibility{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsAIAgentRolloutEligible uses a stable conversation bucket so one customer
|
||||||
|
// remains consistently inside or outside a gray release throughout a session.
|
||||||
|
// Missing legacy values are treated as 100 to preserve existing behavior.
|
||||||
|
func IsAIAgentRolloutEligible(conversation models.Conversation, aiAgent models.AIAgent, channel *models.Channel) bool {
|
||||||
|
percent := normalizedRolloutPercent(aiAgent.RolloutPercent)
|
||||||
|
if channel != nil {
|
||||||
|
channelPercent := normalizedRolloutPercent(channel.AIAgentRolloutPercent)
|
||||||
|
if channelPercent < percent {
|
||||||
|
percent = channelPercent
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if percent >= 100 {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if conversation.ID <= 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
seed := fmt.Sprintf("channel=%d;conversation=%d;agent=%d", conversation.ChannelID, conversation.ID, aiAgent.ID)
|
||||||
|
sum := sha256.Sum256([]byte(seed))
|
||||||
|
bucket := int(binary.BigEndian.Uint64(sum[:8]) % 100)
|
||||||
|
return bucket < percent
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizedRolloutPercent(percent int) int {
|
||||||
|
if percent <= 0 || percent > 100 {
|
||||||
|
return 100
|
||||||
|
}
|
||||||
|
return percent
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *replyEligibility) CanReply(conversation models.Conversation, message models.Message, aiAgent models.AIAgent) bool {
|
||||||
|
if message.SenderType != enums.IMSenderTypeCustomer {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if conversation.CurrentAssigneeID > 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if aiAgent.ServiceMode == enums.IMConversationServiceModeHumanOnly {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if strs.IsBlank(message.Content) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
64
internal/ai/runtime/reply_helpers_test.go
Normal file
64
internal/ai/runtime/reply_helpers_test.go
Normal file
@@ -0,0 +1,64 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestExtractInterruptMessageAndCheckpointError(t *testing.T) {
|
||||||
|
if got := extractInterruptMessage(`{"message":"请补充订单号"}`); got != "请补充订单号" {
|
||||||
|
t.Fatalf("unexpected interrupt message: %q", got)
|
||||||
|
}
|
||||||
|
if got := extractInterruptMessage("not-json"); got != "" {
|
||||||
|
t.Fatalf("expected empty message for invalid json, got %q", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
err := fakeErr("Failed to load from checkpoint: record does not exist")
|
||||||
|
if !isCheckpointMissingError(err) {
|
||||||
|
t.Fatalf("expected checkpoint missing error to be detected")
|
||||||
|
}
|
||||||
|
if isCheckpointMissingError(fakeErr("other error")) {
|
||||||
|
t.Fatalf("expected unrelated error to be ignored")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildConversationInterruptStoresWorkflowCheckpointData(t *testing.T) {
|
||||||
|
item := buildConversationInterrupt(testConversation(1), testMessage(2), testAIAgent(3), &applicationruntime.RunResult{
|
||||||
|
CheckPointData: `{"confirmNodeId":"confirm_1"}`,
|
||||||
|
Interrupted: true,
|
||||||
|
WorkflowRunID: 99,
|
||||||
|
AgentRunID: 88,
|
||||||
|
Interrupts: []applicationruntime.InterruptContextSummary{
|
||||||
|
{Type: "human_confirm", ID: "confirm_1", InfoPreview: `{"message":"请确认"}`},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if item == nil {
|
||||||
|
t.Fatalf("expected interrupt item")
|
||||||
|
}
|
||||||
|
if item.RequestData != `{"confirmNodeId":"confirm_1"}` {
|
||||||
|
t.Fatalf("unexpected request data: %q", item.RequestData)
|
||||||
|
}
|
||||||
|
if item.WorkflowRunID != 99 || item.AgentRunID != 88 || item.WorkflowNodeID != "confirm_1" {
|
||||||
|
t.Fatalf("unexpected workflow interrupt identity: run=%d node=%q", item.WorkflowRunID, item.WorkflowNodeID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeErr string
|
||||||
|
|
||||||
|
func (e fakeErr) Error() string {
|
||||||
|
return string(e)
|
||||||
|
}
|
||||||
|
|
||||||
|
func testConversation(id int64) models.Conversation {
|
||||||
|
return models.Conversation{ID: id}
|
||||||
|
}
|
||||||
|
|
||||||
|
func testMessage(id int64) models.Message {
|
||||||
|
return models.Message{ID: id}
|
||||||
|
}
|
||||||
|
|
||||||
|
func testAIAgent(id int64) models.AIAgent {
|
||||||
|
return models.AIAgent{ID: id}
|
||||||
|
}
|
||||||
114
internal/ai/runtime/reply_interrupt_helpers.go
Normal file
114
internal/ai/runtime/reply_interrupt_helpers.go
Normal file
@@ -0,0 +1,114 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/i18nx"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
type interruptMessagePreview struct {
|
||||||
|
Message string `json:"message"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildConversationInterrupt(conversation models.Conversation, message models.Message, aiAgent models.AIAgent, summary *applicationruntime.RunResult) *models.ConversationInterrupt {
|
||||||
|
if summary == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
item := svc.ConversationInterruptService.GetByCheckPointID(summary.CheckPointID)
|
||||||
|
if item == nil {
|
||||||
|
item = &models.ConversationInterrupt{
|
||||||
|
CheckPointID: summary.CheckPointID,
|
||||||
|
CreatedAt: now,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
item.ConversationID = conversation.ID
|
||||||
|
item.AIAgentID = aiAgent.ID
|
||||||
|
item.AgentRunID = summary.AgentRunID
|
||||||
|
item.SourceMessageID = message.ID
|
||||||
|
item.InterruptID = firstInterruptID(summary)
|
||||||
|
item.InterruptType = firstInterruptType(summary)
|
||||||
|
item.WorkflowRunID = summary.WorkflowRunID
|
||||||
|
item.WorkflowNodeID = firstInterruptID(summary)
|
||||||
|
item.Status = "pending"
|
||||||
|
locale := svc.ConversationHumanDispatchService.ResolveConversationLocale(conversation.ID)
|
||||||
|
item.PromptText = resolveInterruptPrompt(summary, locale)
|
||||||
|
item.RequestData = strings.TrimSpace(summary.CheckPointData)
|
||||||
|
item.UpdatedAt = now
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveInterruptPrompt(summary *applicationruntime.RunResult, locale ...string) string {
|
||||||
|
targetLocale := i18nx.DefaultLocale
|
||||||
|
if len(locale) > 0 && strings.TrimSpace(locale[0]) != "" {
|
||||||
|
targetLocale = locale[0]
|
||||||
|
}
|
||||||
|
if summary == nil || len(summary.Interrupts) == 0 {
|
||||||
|
return i18nx.Getf(targetLocale, "conversation.interrupt.defaultPrompt")
|
||||||
|
}
|
||||||
|
interrupt := summary.Interrupts[0]
|
||||||
|
if prompt := strings.TrimSpace(interrupt.PromptText); prompt != "" {
|
||||||
|
return prompt
|
||||||
|
}
|
||||||
|
if prompt := extractInterruptMessage(interrupt.InfoPreview); prompt != "" {
|
||||||
|
return prompt
|
||||||
|
}
|
||||||
|
if prompt := strings.TrimSpace(summary.ReplyText); prompt != "" {
|
||||||
|
return prompt
|
||||||
|
}
|
||||||
|
if interrupt.Type != "tool_confirmation" {
|
||||||
|
if prompt := strings.TrimSpace(interrupt.InfoPreview); prompt != "" {
|
||||||
|
return prompt
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if displayName := strings.TrimSpace(interrupt.DisplayName); displayName != "" {
|
||||||
|
switch targetLocale {
|
||||||
|
case "ru-RU":
|
||||||
|
return "Готовится выполнение «" + displayName + "». Подтверждаете продолжение?"
|
||||||
|
case "en-US":
|
||||||
|
return "About to execute \"" + displayName + "\". Confirm to continue?"
|
||||||
|
default:
|
||||||
|
return "即将执行“" + displayName + "”,是否确认继续?"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return i18nx.Getf(targetLocale, "conversation.interrupt.defaultPrompt")
|
||||||
|
}
|
||||||
|
|
||||||
|
func extractInterruptMessage(infoPreview string) string {
|
||||||
|
infoPreview = strings.TrimSpace(infoPreview)
|
||||||
|
if infoPreview == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
var payload interruptMessagePreview
|
||||||
|
if err := json.Unmarshal([]byte(infoPreview), &payload); err != nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(payload.Message)
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstInterruptID(summary *applicationruntime.RunResult) string {
|
||||||
|
if summary == nil || len(summary.Interrupts) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(summary.Interrupts[0].ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstInterruptType(summary *applicationruntime.RunResult) string {
|
||||||
|
if summary == nil || len(summary.Interrupts) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(summary.Interrupts[0].Type)
|
||||||
|
}
|
||||||
|
|
||||||
|
func isCheckpointMissingError(err error) bool {
|
||||||
|
if err == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
message := strings.ToLower(strings.TrimSpace(err.Error()))
|
||||||
|
return strings.Contains(message, "failed to load from checkpoint") && strings.Contains(message, "not exist")
|
||||||
|
}
|
||||||
132
internal/ai/runtime/reply_interrupt_service.go
Normal file
132
internal/ai/runtime/reply_interrupt_service.go
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/ai/runtime/graphs"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
type replyInterruptService struct{}
|
||||||
|
|
||||||
|
func newReplyInterruptService() *replyInterruptService {
|
||||||
|
return &replyInterruptService{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *replyInterruptService) ResumePendingInterrupt(ctx context.Context, owner *aiReplyService, replyCtx aiReplyContext) error {
|
||||||
|
if replyCtx.PendingInterrupt == nil {
|
||||||
|
return fmt.Errorf("pending interrupt is required")
|
||||||
|
}
|
||||||
|
summary, err := owner.executor.ResumePendingInterrupt(ctx, runtimeReplyResumeInput{
|
||||||
|
Conversation: replyCtx.Conversation,
|
||||||
|
Message: replyCtx.Message,
|
||||||
|
AIAgent: replyCtx.AIAgent,
|
||||||
|
PendingInterrupt: replyCtx.PendingInterrupt,
|
||||||
|
})
|
||||||
|
replyCtx.setSummary(summary)
|
||||||
|
if err != nil {
|
||||||
|
if isCheckpointMissingError(err) {
|
||||||
|
summary = expiredInterruptSummary()
|
||||||
|
replyCtx.setSummary(summary)
|
||||||
|
replyMessage, expireErr := owner.commit.CommitAIReply(replyCommitInput{
|
||||||
|
Conversation: replyCtx.Conversation,
|
||||||
|
Message: replyCtx.Message,
|
||||||
|
AIAgent: replyCtx.AIAgent,
|
||||||
|
ReplyText: summary.ReplyText,
|
||||||
|
ClientPrefix: "ai_interrupt_expired",
|
||||||
|
WorkflowRunID: summary.WorkflowRunID,
|
||||||
|
})
|
||||||
|
if expireErr != nil {
|
||||||
|
return expireErr
|
||||||
|
}
|
||||||
|
lastResumeMessageID := int64(0)
|
||||||
|
if replyMessage != nil {
|
||||||
|
lastResumeMessageID = replyMessage.ID
|
||||||
|
}
|
||||||
|
if expireMarkErr := svc.ConversationInterruptService.MarkExpired(replyCtx.PendingInterrupt.ID, lastResumeMessageID); expireMarkErr != nil {
|
||||||
|
return expireMarkErr
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if summary != nil && summary.Interrupted {
|
||||||
|
return s.HandleInterruptedResume(owner, replyCtx, summary)
|
||||||
|
}
|
||||||
|
if summary != nil && strings.TrimSpace(summary.ReplyText) != "" {
|
||||||
|
replyMessage, err := owner.commit.CommitAIReply(replyCommitInput{
|
||||||
|
Conversation: replyCtx.Conversation,
|
||||||
|
Message: replyCtx.Message,
|
||||||
|
AIAgent: replyCtx.AIAgent,
|
||||||
|
ReplyText: summary.ReplyText,
|
||||||
|
ClientPrefix: "ai_resume",
|
||||||
|
WorkflowRunID: summary.WorkflowRunID,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
replyMessageID := int64(0)
|
||||||
|
if replyMessage != nil {
|
||||||
|
replyMessageID = replyMessage.ID
|
||||||
|
}
|
||||||
|
if graphs.IsCancellationReply(summary.ReplyText) {
|
||||||
|
return svc.ConversationInterruptService.MarkCancelled(replyCtx.PendingInterrupt.ID, replyMessageID)
|
||||||
|
}
|
||||||
|
return svc.ConversationInterruptService.MarkResolved(replyCtx.PendingInterrupt.ID, replyMessageID)
|
||||||
|
}
|
||||||
|
return svc.ConversationInterruptService.MarkResolved(replyCtx.PendingInterrupt.ID, 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *replyInterruptService) HandleInterruptedSummary(owner *aiReplyService, replyCtx aiReplyContext, summary *applicationruntime.RunResult) error {
|
||||||
|
pending := buildConversationInterrupt(replyCtx.Conversation, replyCtx.Message, replyCtx.AIAgent, summary)
|
||||||
|
if pending != nil && pending.AgentRunID > 0 {
|
||||||
|
pending.AgentStepID = svc.AgentRunService.GetLatestStepID(pending.AgentRunID)
|
||||||
|
}
|
||||||
|
if err := svc.ConversationInterruptService.CreateOrUpdatePending(pending); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
pending = svc.ConversationInterruptService.GetByCheckPointID(summary.CheckPointID)
|
||||||
|
locale := svc.ConversationHumanDispatchService.ResolveConversationLocale(replyCtx.Conversation.ID)
|
||||||
|
replyText := resolveInterruptPrompt(summary, locale)
|
||||||
|
replyMessage, err := owner.commit.CommitAIReply(replyCommitInput{
|
||||||
|
Conversation: replyCtx.Conversation,
|
||||||
|
Message: replyCtx.Message,
|
||||||
|
AIAgent: replyCtx.AIAgent,
|
||||||
|
ReplyText: replyText,
|
||||||
|
ClientPrefix: "ai_interrupt",
|
||||||
|
WorkflowRunID: summary.WorkflowRunID,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if replyMessage != nil && pending != nil {
|
||||||
|
return svc.ConversationInterruptService.MarkPendingAgain(pending.ID, pending.InterruptID, replyText, replyMessage.ID)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *replyInterruptService) HandleInterruptedResume(owner *aiReplyService, replyCtx aiReplyContext, summary *applicationruntime.RunResult) error {
|
||||||
|
if replyCtx.PendingInterrupt == nil {
|
||||||
|
return fmt.Errorf("pending interrupt is required")
|
||||||
|
}
|
||||||
|
locale := svc.ConversationHumanDispatchService.ResolveConversationLocale(replyCtx.Conversation.ID)
|
||||||
|
replyText := resolveInterruptPrompt(summary, locale)
|
||||||
|
replyMessage, err := owner.commit.CommitAIReply(replyCommitInput{
|
||||||
|
Conversation: replyCtx.Conversation,
|
||||||
|
Message: replyCtx.Message,
|
||||||
|
AIAgent: replyCtx.AIAgent,
|
||||||
|
ReplyText: replyText,
|
||||||
|
ClientPrefix: "ai_interrupt_resume",
|
||||||
|
WorkflowRunID: summary.WorkflowRunID,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if replyMessage != nil {
|
||||||
|
return svc.ConversationInterruptService.MarkPendingAgain(replyCtx.PendingInterrupt.ID, firstInterruptID(summary), replyText, replyMessage.ID)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
50
internal/ai/runtime/reply_service.go
Normal file
50
internal/ai/runtime/reply_service.go
Normal file
@@ -0,0 +1,50 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
var AIReplyService = newAIReplyService()
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
svc.TriggerAIReplyAsyncHook = AIReplyService.TriggerReplyAsync
|
||||||
|
}
|
||||||
|
|
||||||
|
type replyExecutor interface {
|
||||||
|
Run(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error)
|
||||||
|
ResumePendingInterrupt(ctx context.Context, input runtimeReplyResumeInput) (*applicationruntime.RunResult, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
func newAIReplyService() *aiReplyService {
|
||||||
|
return newAIReplyServiceWithExecutor(newRuntimeReplyExecutor())
|
||||||
|
}
|
||||||
|
|
||||||
|
func newAIReplyServiceWithExecutor(executor replyExecutor) *aiReplyService {
|
||||||
|
return &aiReplyService{
|
||||||
|
eligibility: newReplyEligibility(),
|
||||||
|
executor: executor,
|
||||||
|
interrupts: newReplyInterruptService(),
|
||||||
|
commit: newReplyCommitService(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type aiReplyService struct {
|
||||||
|
eligibility *replyEligibility
|
||||||
|
executor replyExecutor
|
||||||
|
interrupts *replyInterruptService
|
||||||
|
commit *replyCommitService
|
||||||
|
}
|
||||||
|
|
||||||
|
func firstInvokedToolCode(summary *applicationruntime.RunResult) string {
|
||||||
|
if summary == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if len(summary.InvokedToolCodes) > 0 {
|
||||||
|
return strings.TrimSpace(summary.InvokedToolCodes[0])
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
145
internal/ai/runtime/reply_service_test.go
Normal file
145
internal/ai/runtime/reply_service_test.go
Normal file
@@ -0,0 +1,145 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestReplyEligibilityCanReply(t *testing.T) {
|
||||||
|
eligibility := newReplyEligibility()
|
||||||
|
conversation := newConversationFixture()
|
||||||
|
message := newCustomerMessageFixture("hello")
|
||||||
|
aiAgent := newAIAgentFixture()
|
||||||
|
|
||||||
|
if !eligibility.CanReply(conversation, message, aiAgent) {
|
||||||
|
t.Fatalf("expected customer message to be replyable")
|
||||||
|
}
|
||||||
|
|
||||||
|
message.SenderType = enums.IMSenderTypeAgent
|
||||||
|
if eligibility.CanReply(conversation, message, aiAgent) {
|
||||||
|
t.Fatalf("expected non-customer message to be rejected")
|
||||||
|
}
|
||||||
|
|
||||||
|
message = newCustomerMessageFixture("hello")
|
||||||
|
conversation.HandoffAt = ptrTime(time.Now())
|
||||||
|
conversation.CurrentAssigneeID = 0
|
||||||
|
if !eligibility.CanReply(conversation, message, aiAgent) {
|
||||||
|
t.Fatalf("expected unassigned handed-off conversation to still be replyable by AI")
|
||||||
|
}
|
||||||
|
|
||||||
|
conversation = newConversationFixture()
|
||||||
|
conversation.CurrentAssigneeID = 1
|
||||||
|
if eligibility.CanReply(conversation, message, aiAgent) {
|
||||||
|
t.Fatalf("expected assigned conversation to be rejected")
|
||||||
|
}
|
||||||
|
|
||||||
|
conversation = newConversationFixture()
|
||||||
|
aiAgent.ServiceMode = enums.IMConversationServiceModeHumanOnly
|
||||||
|
if eligibility.CanReply(conversation, message, aiAgent) {
|
||||||
|
t.Fatalf("expected human-only agent to be rejected")
|
||||||
|
}
|
||||||
|
|
||||||
|
aiAgent = newAIAgentFixture()
|
||||||
|
message.Content = " "
|
||||||
|
if eligibility.CanReply(conversation, message, aiAgent) {
|
||||||
|
t.Fatalf("expected blank message to be rejected")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAIAgentRolloutUsesStableConversationBucket(t *testing.T) {
|
||||||
|
conversation := models.Conversation{ID: 101, ChannelID: 7}
|
||||||
|
agent := models.AIAgent{ID: 9, RolloutPercent: 50}
|
||||||
|
first := IsAIAgentRolloutEligible(conversation, agent, &models.Channel{AIAgentRolloutPercent: 100})
|
||||||
|
for range 20 {
|
||||||
|
if got := IsAIAgentRolloutEligible(conversation, agent, &models.Channel{AIAgentRolloutPercent: 100}); got != first {
|
||||||
|
t.Fatalf("rollout bucket changed within one conversation: first=%t got=%t", first, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if normalizedRolloutPercent(0) != 100 || normalizedRolloutPercent(101) != 100 || normalizedRolloutPercent(25) != 25 {
|
||||||
|
t.Fatal("unexpected rollout percent normalization")
|
||||||
|
}
|
||||||
|
if !IsAIAgentRolloutEligible(conversation, models.AIAgent{ID: 9, RolloutPercent: 0}, &models.Channel{}) {
|
||||||
|
t.Fatal("legacy zero rollout values must preserve full rollout")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveReplyTimeout(t *testing.T) {
|
||||||
|
service := newAIReplyService()
|
||||||
|
aiAgent := newAIAgentFixture()
|
||||||
|
|
||||||
|
if got := service.resolveReplyTimeout(aiAgent); got != 180*time.Second {
|
||||||
|
t.Fatalf("expected default timeout, got %v", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
aiAgent.ReplyTimeoutSeconds = 30
|
||||||
|
if got := service.resolveReplyTimeout(aiAgent); got != 30*time.Second {
|
||||||
|
t.Fatalf("expected exact timeout, got %v", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
aiAgent.ReplyTimeoutSeconds = 999
|
||||||
|
if got := service.resolveReplyTimeout(aiAgent); got != 600*time.Second {
|
||||||
|
t.Fatalf("expected clamped timeout, got %v", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveInterruptPrompt(t *testing.T) {
|
||||||
|
summary := &applicationruntime.RunResult{
|
||||||
|
Interrupts: []applicationruntime.InterruptContextSummary{
|
||||||
|
{
|
||||||
|
ID: "interrupt-1",
|
||||||
|
Type: "question",
|
||||||
|
InfoPreview: `{"message":"请补充订单号"}`,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if got := resolveInterruptPrompt(summary); got != "请补充订单号" {
|
||||||
|
t.Fatalf("unexpected interrupt prompt: %q", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
summary.Interrupts[0].InfoPreview = "直接补充手机号"
|
||||||
|
if got := resolveInterruptPrompt(summary); got != "直接补充手机号" {
|
||||||
|
t.Fatalf("unexpected raw interrupt prompt: %q", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
summary.ReplyText = ""
|
||||||
|
summary.Interrupts[0] = applicationruntime.InterruptContextSummary{
|
||||||
|
ID: "system/server_time",
|
||||||
|
Type: "tool_confirmation",
|
||||||
|
DisplayName: "获取当前时间",
|
||||||
|
InfoPreview: "system/server_time",
|
||||||
|
}
|
||||||
|
if got := resolveInterruptPrompt(summary, "zh-CN"); got != "即将执行“获取当前时间”,是否确认继续?" {
|
||||||
|
t.Fatalf("tool code leaked into customer prompt: %q", got)
|
||||||
|
}
|
||||||
|
if got := resolveInterruptPrompt(summary, "en-US"); got != "About to execute \"获取当前时间\". Confirm to continue?" {
|
||||||
|
t.Fatalf("expected en-US prompt, got: %q", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
summary.Interrupts[0].PromptText = "请确认是否更新客户资料。"
|
||||||
|
if got := resolveInterruptPrompt(summary); got != "请确认是否更新客户资料。" {
|
||||||
|
t.Fatalf("explicit customer prompt was not preferred: %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func newConversationFixture() models.Conversation {
|
||||||
|
return models.Conversation{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func newCustomerMessageFixture(content string) models.Message {
|
||||||
|
return models.Message{
|
||||||
|
SenderType: enums.IMSenderTypeCustomer,
|
||||||
|
Content: content,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func newAIAgentFixture() models.AIAgent {
|
||||||
|
return models.AIAgent{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func ptrTime(v time.Time) *time.Time {
|
||||||
|
return &v
|
||||||
|
}
|
||||||
144
internal/ai/runtime/reply_trigger_service.go
Normal file
144
internal/ai/runtime/reply_trigger_service.go
Normal file
@@ -0,0 +1,144 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"log/slog"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/pkg/tracex"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *aiReplyService) resolveReplyTimeout(aiAgent models.AIAgent) time.Duration {
|
||||||
|
if aiAgent.ReplyTimeoutSeconds <= 0 {
|
||||||
|
return time.Duration(defaultAIReplyAsyncTimeoutSeconds) * time.Second
|
||||||
|
}
|
||||||
|
if aiAgent.ReplyTimeoutSeconds > maxAIReplyAsyncTimeoutSeconds {
|
||||||
|
return time.Duration(maxAIReplyAsyncTimeoutSeconds) * time.Second
|
||||||
|
}
|
||||||
|
return time.Duration(aiAgent.ReplyTimeoutSeconds) * time.Second
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *aiReplyService) TriggerReplyAsync(conversation models.Conversation, message models.Message) {
|
||||||
|
go func() {
|
||||||
|
aiAgent := svc.AIAgentService.Get(conversation.AIAgentID)
|
||||||
|
if aiAgent == nil || aiAgent.Status != enums.StatusOk {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
startedAt := time.Now()
|
||||||
|
timeout := s.resolveReplyTimeout(*aiAgent)
|
||||||
|
ctx, cancel := context.WithTimeout(tracex.ContextWithRequestID(context.Background(), message.RequestID), timeout)
|
||||||
|
defer cancel()
|
||||||
|
if err := s.TriggerReply(ctx, conversation, message, *aiAgent); err != nil {
|
||||||
|
slog.Error("failed to trigger ai reply",
|
||||||
|
"requestId", message.RequestID,
|
||||||
|
"message_id", message.ID,
|
||||||
|
"timeout_ms", timeout.Milliseconds(),
|
||||||
|
"elapsed_ms", time.Since(startedAt).Milliseconds(),
|
||||||
|
"error", err)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *aiReplyService) TriggerReply(ctx context.Context, conversation models.Conversation, message models.Message, aiAgent models.AIAgent) (retErr error) {
|
||||||
|
var summary *applicationruntime.RunResult
|
||||||
|
replyCtx := aiReplyContext{
|
||||||
|
Conversation: conversation,
|
||||||
|
Message: message,
|
||||||
|
AIAgent: aiAgent,
|
||||||
|
SummaryRef: &summary,
|
||||||
|
}
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if s.eligibility != nil && !s.eligibility.CanReply(conversation, message, aiAgent) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if !IsAIAgentRolloutEligible(conversation, aiAgent, svc.ChannelService.Get(conversation.ChannelID)) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if pendingInterrupt := svc.ConversationInterruptService.FindLatestPendingByConversationID(conversation.ID); pendingInterrupt != nil {
|
||||||
|
replyCtx.PendingInterrupt = pendingInterrupt
|
||||||
|
return s.resumePendingInterrupt(ctx, replyCtx)
|
||||||
|
}
|
||||||
|
return s.executeReply(ctx, replyCtx)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *aiReplyService) resumePendingInterrupt(ctx context.Context, replyCtx aiReplyContext) error {
|
||||||
|
return s.interrupts.ResumePendingInterrupt(ctx, s, replyCtx)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *aiReplyService) executeReply(ctx context.Context, replyCtx aiReplyContext) error {
|
||||||
|
summary, err := s.executor.Run(ctx, runtimeReplyRunInput{
|
||||||
|
Conversation: replyCtx.Conversation,
|
||||||
|
Message: replyCtx.Message,
|
||||||
|
AIAgent: replyCtx.AIAgent,
|
||||||
|
})
|
||||||
|
replyCtx.setSummary(summary)
|
||||||
|
if err != nil {
|
||||||
|
if replyCtx.AIAgent.AutoHandoffOnFailure {
|
||||||
|
if _, handoffErr := svc.ConversationHumanDispatchService.HandoffByAIWithRequestID(
|
||||||
|
replyCtx.Conversation.ID,
|
||||||
|
replyCtx.AIAgent,
|
||||||
|
"AI execution failure: "+err.Error(),
|
||||||
|
replyCtx.Message.RequestID,
|
||||||
|
); handoffErr != nil {
|
||||||
|
slog.Error("failed to auto-handoff on AI execution failure",
|
||||||
|
"conversation_id", replyCtx.Conversation.ID,
|
||||||
|
"ai_agent_id", replyCtx.AIAgent.ID,
|
||||||
|
"error", handoffErr,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if summary != nil && summary.Interrupted {
|
||||||
|
return s.interrupts.HandleInterruptedSummary(s, replyCtx, summary)
|
||||||
|
}
|
||||||
|
if summary != nil && summary.HandoffRequested {
|
||||||
|
if _, err := svc.ConversationHumanDispatchService.HandoffByAIWithRequestID(
|
||||||
|
replyCtx.Conversation.ID,
|
||||||
|
replyCtx.AIAgent,
|
||||||
|
summary.HandoffReason,
|
||||||
|
replyCtx.Message.RequestID,
|
||||||
|
); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if summary != nil && strings.TrimSpace(summary.ReplyText) != "" {
|
||||||
|
_, err := s.commit.CommitAIReply(replyCommitInput{
|
||||||
|
Conversation: replyCtx.Conversation,
|
||||||
|
Message: replyCtx.Message,
|
||||||
|
AIAgent: replyCtx.AIAgent,
|
||||||
|
ReplyText: summary.ReplyText,
|
||||||
|
ClientPrefix: "ai_reply",
|
||||||
|
WorkflowRunID: summary.WorkflowRunID,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if replyCtx.AIAgent.AutoHandoffOnFailure {
|
||||||
|
if _, handoffErr := svc.ConversationHumanDispatchService.HandoffByAIWithRequestID(
|
||||||
|
replyCtx.Conversation.ID,
|
||||||
|
replyCtx.AIAgent,
|
||||||
|
"AI generated empty reply or fallback unavailable",
|
||||||
|
replyCtx.Message.RequestID,
|
||||||
|
); handoffErr != nil {
|
||||||
|
slog.Error("failed to auto-handoff on empty AI reply",
|
||||||
|
"conversation_id", replyCtx.Conversation.ID,
|
||||||
|
"ai_agent_id", replyCtx.AIAgent.ID,
|
||||||
|
"error", handoffErr,
|
||||||
|
)
|
||||||
|
return handoffErr
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
342
internal/ai/runtime/reply_trigger_service_test.go
Normal file
342
internal/ai/runtime/reply_trigger_service_test.go
Normal file
@@ -0,0 +1,342 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
svc "agent-desk/internal/services"
|
||||||
|
|
||||||
|
"github.com/glebarez/sqlite"
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
"gorm.io/gorm"
|
||||||
|
"gorm.io/gorm/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
type mockReplyExecutor struct {
|
||||||
|
runFunc func(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error)
|
||||||
|
resumeFunc func(ctx context.Context, input runtimeReplyResumeInput) (*applicationruntime.RunResult, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *mockReplyExecutor) Run(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
if m.runFunc != nil {
|
||||||
|
return m.runFunc(ctx, input)
|
||||||
|
}
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *mockReplyExecutor) ResumePendingInterrupt(ctx context.Context, input runtimeReplyResumeInput) (*applicationruntime.RunResult, error) {
|
||||||
|
if m.resumeFunc != nil {
|
||||||
|
return m.resumeFunc(ctx, input)
|
||||||
|
}
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func setupReplyTriggerTestDB(t *testing.T) *gorm.DB {
|
||||||
|
t.Helper()
|
||||||
|
dbName := "reply_trigger_test_" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name())
|
||||||
|
db, err := gorm.Open(sqlite.Open("file:"+dbName+"?mode=memory&cache=shared"), &gorm.Config{
|
||||||
|
NamingStrategy: schema.NamingStrategy{
|
||||||
|
TablePrefix: "t_",
|
||||||
|
SingularTable: true,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("open sqlite db: %v", err)
|
||||||
|
}
|
||||||
|
sqlDB, err := db.DB()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get sqlite db: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
_ = sqlDB.Close()
|
||||||
|
})
|
||||||
|
if err := db.AutoMigrate(
|
||||||
|
&models.User{},
|
||||||
|
&models.Customer{},
|
||||||
|
&models.CustomerIdentity{},
|
||||||
|
&models.AIAgent{},
|
||||||
|
&models.AgentTeam{},
|
||||||
|
&models.AgentTeamSchedule{},
|
||||||
|
&models.AgentProfile{},
|
||||||
|
&models.Conversation{},
|
||||||
|
&models.ConversationParticipant{},
|
||||||
|
&models.ConversationAssignment{},
|
||||||
|
&models.ConversationEventLog{},
|
||||||
|
&models.ConversationReadState{},
|
||||||
|
&models.ConversationInterrupt{},
|
||||||
|
&models.Message{},
|
||||||
|
&models.Channel{},
|
||||||
|
&models.ChannelMessageOutbox{},
|
||||||
|
); err != nil {
|
||||||
|
t.Fatalf("auto migrate: %v", err)
|
||||||
|
}
|
||||||
|
sqls.SetDB(db)
|
||||||
|
return db
|
||||||
|
}
|
||||||
|
|
||||||
|
func createReplyTriggerTestAIAgent(t *testing.T, db *gorm.DB, autoHandoff bool, handoffMsg string) models.AIAgent {
|
||||||
|
t.Helper()
|
||||||
|
now := time.Now()
|
||||||
|
item := models.AIAgent{
|
||||||
|
Name: "TriggerAI",
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
ServiceMode: enums.IMConversationServiceModeAIFirst,
|
||||||
|
AutoHandoffOnFailure: autoHandoff,
|
||||||
|
HandoffMessage: handoffMsg,
|
||||||
|
TeamIDs: "1",
|
||||||
|
RolloutPercent: 100,
|
||||||
|
AuditFields: models.AuditFields{
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := db.Create(&item).Error; err != nil {
|
||||||
|
t.Fatalf("create ai agent: %v", err)
|
||||||
|
}
|
||||||
|
if !autoHandoff {
|
||||||
|
if err := db.Model(&item).Update("auto_handoff_on_failure", false).Error; err != nil {
|
||||||
|
t.Fatalf("update auto_handoff_on_failure: %v", err)
|
||||||
|
}
|
||||||
|
item.AutoHandoffOnFailure = false
|
||||||
|
}
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
|
||||||
|
func createReplyTriggerTestTeam(t *testing.T, db *gorm.DB, teamID int64) {
|
||||||
|
t.Helper()
|
||||||
|
if err := db.Create(&models.AgentTeam{ID: teamID, Name: "支持组", Status: enums.StatusOk}).Error; err != nil {
|
||||||
|
t.Fatalf("create team: %v", err)
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
if err := db.Create(&models.AgentTeamSchedule{
|
||||||
|
TeamID: teamID,
|
||||||
|
StartAt: now.Add(-time.Hour),
|
||||||
|
EndAt: now.Add(time.Hour),
|
||||||
|
Status: enums.StatusOk,
|
||||||
|
}).Error; err != nil {
|
||||||
|
t.Fatalf("create schedule: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func createReplyTriggerTestConversation(t *testing.T, db *gorm.DB, aiAgentID int64) models.Conversation {
|
||||||
|
t.Helper()
|
||||||
|
now := time.Now()
|
||||||
|
item := models.Conversation{
|
||||||
|
CustomerID: 1,
|
||||||
|
ChannelID: 1,
|
||||||
|
AIAgentID: aiAgentID,
|
||||||
|
Status: enums.IMConversationStatusAIServing,
|
||||||
|
ServiceMode: enums.IMConversationServiceModeAIFirst,
|
||||||
|
LastActiveAt: now,
|
||||||
|
AuditFields: models.AuditFields{
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := db.Create(&item).Error; err != nil {
|
||||||
|
t.Fatalf("create conversation: %v", err)
|
||||||
|
}
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
|
||||||
|
func createCustomerTestMessage(t *testing.T, db *gorm.DB, conversationID int64, content string) models.Message {
|
||||||
|
t.Helper()
|
||||||
|
now := time.Now()
|
||||||
|
item := models.Message{
|
||||||
|
ConversationID: conversationID,
|
||||||
|
SenderType: enums.IMSenderTypeCustomer,
|
||||||
|
SenderID: 1,
|
||||||
|
Content: content,
|
||||||
|
RequestID: "req-trigger-test",
|
||||||
|
SendStatus: enums.IMMessageStatusSent,
|
||||||
|
AuditFields: models.AuditFields{
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if err := db.Create(&item).Error; err != nil {
|
||||||
|
t.Fatalf("create message: %v", err)
|
||||||
|
}
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAIReplyService_AutoHandoffOnExecutionFailure(t *testing.T) {
|
||||||
|
db := setupReplyTriggerTestDB(t)
|
||||||
|
createReplyTriggerTestTeam(t, db, 1)
|
||||||
|
aiAgent := createReplyTriggerTestAIAgent(t, db, true, "AI故障,正在为您转接人工客服...")
|
||||||
|
conversation := createReplyTriggerTestConversation(t, db, aiAgent.ID)
|
||||||
|
message := createCustomerTestMessage(t, db, conversation.ID, "帮助我")
|
||||||
|
|
||||||
|
mock := &mockReplyExecutor{
|
||||||
|
runFunc: func(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
return nil, fmt.Errorf("llm connection timeout")
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
service := newAIReplyServiceWithExecutor(mock)
|
||||||
|
err := service.TriggerReply(context.Background(), conversation, message, aiAgent)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "llm connection timeout") {
|
||||||
|
t.Fatalf("expected llm timeout error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
currentConv := svc.ConversationService.Get(conversation.ID)
|
||||||
|
if currentConv.Status != enums.IMConversationStatusPending {
|
||||||
|
t.Fatalf("expected conversation status pending after handoff, got %v", currentConv.Status)
|
||||||
|
}
|
||||||
|
if currentConv.HandoffAt == nil {
|
||||||
|
t.Fatalf("expected handoffAt to be set")
|
||||||
|
}
|
||||||
|
if !strings.Contains(currentConv.HandoffReason, "AI execution failure: llm connection timeout") {
|
||||||
|
t.Fatalf("unexpected handoff reason: %q", currentConv.HandoffReason)
|
||||||
|
}
|
||||||
|
|
||||||
|
sentMsg := svc.MessageService.FindOne(sqls.NewCnd().Eq("conversation_id", conversation.ID).Desc("id"))
|
||||||
|
if sentMsg == nil {
|
||||||
|
t.Fatalf("expected customer notice message to be sent")
|
||||||
|
}
|
||||||
|
if sentMsg.Content != "AI故障,正在为您转接人工客服..." {
|
||||||
|
t.Fatalf("unexpected notice message content: %q", sentMsg.Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAIReplyService_AutoHandoffDisabledOnExecutionFailure(t *testing.T) {
|
||||||
|
db := setupReplyTriggerTestDB(t)
|
||||||
|
createReplyTriggerTestTeam(t, db, 1)
|
||||||
|
aiAgent := createReplyTriggerTestAIAgent(t, db, false, "")
|
||||||
|
conversation := createReplyTriggerTestConversation(t, db, aiAgent.ID)
|
||||||
|
message := createCustomerTestMessage(t, db, conversation.ID, "帮助我")
|
||||||
|
|
||||||
|
mock := &mockReplyExecutor{
|
||||||
|
runFunc: func(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
return nil, fmt.Errorf("llm connection timeout")
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
service := newAIReplyServiceWithExecutor(mock)
|
||||||
|
err := service.TriggerReply(context.Background(), conversation, message, aiAgent)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "llm connection timeout") {
|
||||||
|
t.Fatalf("expected llm timeout error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
currentConv := svc.ConversationService.Get(conversation.ID)
|
||||||
|
if currentConv.Status != enums.IMConversationStatusAIServing {
|
||||||
|
t.Fatalf("expected conversation to stay AI serving, got %v", currentConv.Status)
|
||||||
|
}
|
||||||
|
if currentConv.HandoffAt != nil {
|
||||||
|
t.Fatalf("expected handoffAt to stay nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAIReplyService_AutoHandoffOnEmptyReply(t *testing.T) {
|
||||||
|
db := setupReplyTriggerTestDB(t)
|
||||||
|
createReplyTriggerTestTeam(t, db, 1)
|
||||||
|
aiAgent := createReplyTriggerTestAIAgent(t, db, true, "AI无法解答,已为您转接人工客服")
|
||||||
|
conversation := createReplyTriggerTestConversation(t, db, aiAgent.ID)
|
||||||
|
message := createCustomerTestMessage(t, db, conversation.ID, "复杂未知问题")
|
||||||
|
|
||||||
|
mock := &mockReplyExecutor{
|
||||||
|
runFunc: func(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
return &applicationruntime.RunResult{
|
||||||
|
ReplyText: " ",
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
service := newAIReplyServiceWithExecutor(mock)
|
||||||
|
err := service.TriggerReply(context.Background(), conversation, message, aiAgent)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected nil error on successful empty reply auto-handoff, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
currentConv := svc.ConversationService.Get(conversation.ID)
|
||||||
|
if currentConv.Status != enums.IMConversationStatusPending {
|
||||||
|
t.Fatalf("expected conversation status pending after handoff, got %v", currentConv.Status)
|
||||||
|
}
|
||||||
|
if currentConv.HandoffAt == nil {
|
||||||
|
t.Fatalf("expected handoffAt to be set")
|
||||||
|
}
|
||||||
|
if !strings.Contains(currentConv.HandoffReason, "AI generated empty reply or fallback unavailable") {
|
||||||
|
t.Fatalf("unexpected handoff reason: %q", currentConv.HandoffReason)
|
||||||
|
}
|
||||||
|
|
||||||
|
sentMsg := svc.MessageService.FindOne(sqls.NewCnd().Eq("conversation_id", conversation.ID).Desc("id"))
|
||||||
|
if sentMsg == nil {
|
||||||
|
t.Fatalf("expected customer notice message to be sent")
|
||||||
|
}
|
||||||
|
if sentMsg.Content != "AI无法解答,已为您转接人工客服" {
|
||||||
|
t.Fatalf("unexpected notice message content: %q", sentMsg.Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAIReplyService_AutoHandoffDisabledOnEmptyReply(t *testing.T) {
|
||||||
|
db := setupReplyTriggerTestDB(t)
|
||||||
|
createReplyTriggerTestTeam(t, db, 1)
|
||||||
|
aiAgent := createReplyTriggerTestAIAgent(t, db, false, "")
|
||||||
|
conversation := createReplyTriggerTestConversation(t, db, aiAgent.ID)
|
||||||
|
message := createCustomerTestMessage(t, db, conversation.ID, "复杂未知问题")
|
||||||
|
|
||||||
|
mock := &mockReplyExecutor{
|
||||||
|
runFunc: func(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
return &applicationruntime.RunResult{
|
||||||
|
ReplyText: "",
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
service := newAIReplyServiceWithExecutor(mock)
|
||||||
|
err := service.TriggerReply(context.Background(), conversation, message, aiAgent)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected nil error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
currentConv := svc.ConversationService.Get(conversation.ID)
|
||||||
|
if currentConv.Status != enums.IMConversationStatusAIServing {
|
||||||
|
t.Fatalf("expected conversation to stay AI serving, got %v", currentConv.Status)
|
||||||
|
}
|
||||||
|
if currentConv.HandoffAt != nil {
|
||||||
|
t.Fatalf("expected handoffAt to stay nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAIReplyService_NormalReplyCommitted(t *testing.T) {
|
||||||
|
db := setupReplyTriggerTestDB(t)
|
||||||
|
aiAgent := createReplyTriggerTestAIAgent(t, db, true, "")
|
||||||
|
conversation := createReplyTriggerTestConversation(t, db, aiAgent.ID)
|
||||||
|
message := createCustomerTestMessage(t, db, conversation.ID, "你好")
|
||||||
|
|
||||||
|
mock := &mockReplyExecutor{
|
||||||
|
runFunc: func(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
return &applicationruntime.RunResult{
|
||||||
|
ReplyText: "您好,有什么可以帮您?",
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
service := newAIReplyServiceWithExecutor(mock)
|
||||||
|
err := service.TriggerReply(context.Background(), conversation, message, aiAgent)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected nil error on normal reply, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
currentConv := svc.ConversationService.Get(conversation.ID)
|
||||||
|
if currentConv.Status != enums.IMConversationStatusAIServing {
|
||||||
|
t.Fatalf("expected conversation to stay AI serving, got %v", currentConv.Status)
|
||||||
|
}
|
||||||
|
if currentConv.HandoffAt != nil {
|
||||||
|
t.Fatalf("expected handoffAt to stay nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
sentMsg := svc.MessageService.FindOne(sqls.NewCnd().Eq("conversation_id", conversation.ID).Desc("id"))
|
||||||
|
if sentMsg == nil {
|
||||||
|
t.Fatalf("expected reply message to be sent")
|
||||||
|
}
|
||||||
|
if sentMsg.Content != "您好,有什么可以帮您?" {
|
||||||
|
t.Fatalf("unexpected reply content: %q", sentMsg.Content)
|
||||||
|
}
|
||||||
|
}
|
||||||
6
internal/ai/runtime/reply_types.go
Normal file
6
internal/ai/runtime/reply_types.go
Normal file
@@ -0,0 +1,6 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultAIReplyAsyncTimeoutSeconds = 180
|
||||||
|
maxAIReplyAsyncTimeoutSeconds = 600
|
||||||
|
)
|
||||||
296
internal/ai/runtime/retrievers/knowledge_retriever.go
Normal file
296
internal/ai/runtime/retrievers/knowledge_retriever.go
Normal file
@@ -0,0 +1,296 @@
|
|||||||
|
package retrievers
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/rag"
|
||||||
|
"agent-desk/internal/ai/runtime/traces"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
"agent-desk/internal/pkg/enums"
|
||||||
|
"agent-desk/internal/repositories"
|
||||||
|
|
||||||
|
"github.com/mlogclub/simple/sqls"
|
||||||
|
)
|
||||||
|
|
||||||
|
const defaultRuntimeKnowledgeContextTokens = 4000
|
||||||
|
const defaultRuntimeKnowledgeTopK = 8
|
||||||
|
const defaultRuntimeKnowledgeScoreThreshold = 0.3
|
||||||
|
const defaultRuntimeKnowledgeMaxContextItems = 5
|
||||||
|
|
||||||
|
type KnowledgeRetriever struct {
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
knowledgeBaseIDs []int64
|
||||||
|
}
|
||||||
|
|
||||||
|
type KnowledgeRetrieveOptions struct {
|
||||||
|
ContextMaxTokens int
|
||||||
|
MaxContextItems int
|
||||||
|
TopK int
|
||||||
|
ScoreThreshold float64
|
||||||
|
QueryPreview string
|
||||||
|
}
|
||||||
|
|
||||||
|
type KnowledgeBaseRetrievePolicy struct {
|
||||||
|
KnowledgeBaseID int64
|
||||||
|
TopK int
|
||||||
|
ScoreThreshold float64
|
||||||
|
}
|
||||||
|
|
||||||
|
type KnowledgeRetrieveResult struct {
|
||||||
|
KnowledgeBaseIDs []int64
|
||||||
|
Query string
|
||||||
|
Options KnowledgeRetrieveOptions
|
||||||
|
Hits []rag.RetrieveResult
|
||||||
|
ContextResults []rag.RetrieveResult
|
||||||
|
ContextText string
|
||||||
|
TopScore float64
|
||||||
|
AnswerMode enums.KnowledgeAnswerMode
|
||||||
|
Trace *rag.RetrieveTrace
|
||||||
|
TraceItems []traces.RetrieverTraceItem
|
||||||
|
TraceSummary traces.RetrieverTraceSummary
|
||||||
|
Policies []KnowledgeBaseRetrievePolicy
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewKnowledgeRetriever(aiAgent models.AIAgent, knowledgeBaseIDs []int64) *KnowledgeRetriever {
|
||||||
|
return &KnowledgeRetriever{
|
||||||
|
AIAgent: aiAgent,
|
||||||
|
knowledgeBaseIDs: append([]int64(nil), knowledgeBaseIDs...),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func DefaultKnowledgeRetrieveOptions() KnowledgeRetrieveOptions {
|
||||||
|
return KnowledgeRetrieveOptions{
|
||||||
|
ContextMaxTokens: defaultRuntimeKnowledgeContextTokens,
|
||||||
|
MaxContextItems: defaultRuntimeKnowledgeMaxContextItems,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *KnowledgeRetriever) ConfiguredKnowledgeBaseIDs() []int64 {
|
||||||
|
return append([]int64(nil), r.knowledgeBaseIDs...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *KnowledgeRetriever) Retrieve(ctx context.Context, query string) ([]rag.RetrieveResult, *rag.RetrieveTrace, error) {
|
||||||
|
return r.RetrieveByOptions(ctx, DefaultKnowledgeRetrieveOptions(), query)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *KnowledgeRetriever) RetrieveByOptions(ctx context.Context, opts KnowledgeRetrieveOptions, query string) ([]rag.RetrieveResult, *rag.RetrieveTrace, error) {
|
||||||
|
ids := r.ConfiguredKnowledgeBaseIDs()
|
||||||
|
return rag.Retrieve.RetrieveWithTrace(ctx, rag.RetrieveRequest{
|
||||||
|
Query: query,
|
||||||
|
KnowledgeBaseIDs: ids,
|
||||||
|
TopK: opts.TopK,
|
||||||
|
ScoreThreshold: opts.ScoreThreshold,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *KnowledgeRetriever) RetrieveContext(ctx context.Context, query string) (*KnowledgeRetrieveResult, error) {
|
||||||
|
return r.RetrieveContextByOptions(ctx, DefaultKnowledgeRetrieveOptions(), query)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *KnowledgeRetriever) RetrieveContextByOptions(ctx context.Context, opts KnowledgeRetrieveOptions, query string) (*KnowledgeRetrieveResult, error) {
|
||||||
|
query = strings.TrimSpace(query)
|
||||||
|
knowledgeBaseIDs := r.ConfiguredKnowledgeBaseIDs()
|
||||||
|
policies := r.resolvePolicies(knowledgeBaseIDs, opts)
|
||||||
|
contextMaxTokens := opts.ContextMaxTokens
|
||||||
|
if contextMaxTokens <= 0 {
|
||||||
|
contextMaxTokens = defaultRuntimeKnowledgeContextTokens
|
||||||
|
}
|
||||||
|
maxContextItems := opts.MaxContextItems
|
||||||
|
if maxContextItems <= 0 {
|
||||||
|
maxContextItems = defaultRuntimeKnowledgeMaxContextItems
|
||||||
|
}
|
||||||
|
queryPreview := strings.TrimSpace(opts.QueryPreview)
|
||||||
|
if queryPreview == "" {
|
||||||
|
queryPreview = query
|
||||||
|
}
|
||||||
|
ret := &KnowledgeRetrieveResult{
|
||||||
|
KnowledgeBaseIDs: append([]int64(nil), knowledgeBaseIDs...),
|
||||||
|
Query: query,
|
||||||
|
Options: KnowledgeRetrieveOptions{
|
||||||
|
ContextMaxTokens: contextMaxTokens,
|
||||||
|
MaxContextItems: maxContextItems,
|
||||||
|
TopK: opts.TopK,
|
||||||
|
ScoreThreshold: opts.ScoreThreshold,
|
||||||
|
QueryPreview: queryPreview,
|
||||||
|
},
|
||||||
|
Policies: append([]KnowledgeBaseRetrievePolicy(nil), policies...),
|
||||||
|
}
|
||||||
|
if query == "" || len(knowledgeBaseIDs) == 0 {
|
||||||
|
return ret, nil
|
||||||
|
}
|
||||||
|
results, trace, err := r.RetrieveByOptions(ctx, opts, query)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
ret.Hits = append([]rag.RetrieveResult(nil), results...)
|
||||||
|
ret.Trace = trace
|
||||||
|
ret.ContextResults = rag.Retrieve.SelectContextResults(results, contextMaxTokens)
|
||||||
|
ret.ContextResults = limitContextResults(ret.ContextResults, maxContextItems)
|
||||||
|
ret.ContextText = strings.TrimSpace(buildContextText(ret.ContextResults))
|
||||||
|
ret.TopScore = resolveTopScore(results)
|
||||||
|
ret.AnswerMode = resolveRuntimeAnswerMode(knowledgeBaseIDs, results)
|
||||||
|
ret.TraceItems = buildRetrieverTraceItems(queryPreview, results, trace)
|
||||||
|
ret.TraceSummary = buildRetrieverTraceSummary(ret.Options, ret.Policies, ret.ContextResults, results, trace)
|
||||||
|
return ret, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func limitContextResults(results []rag.RetrieveResult, maxItems int) []rag.RetrieveResult {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if maxItems <= 0 || len(results) <= maxItems {
|
||||||
|
return append([]rag.RetrieveResult(nil), results...)
|
||||||
|
}
|
||||||
|
return append([]rag.RetrieveResult(nil), results[:maxItems]...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildContextText(results []rag.RetrieveResult) string {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(rag.Retrieve.BuildContext(context.Background(), results, 1<<30))
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveTopScore(results []rag.RetrieveResult) float64 {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return float64(results[0].Score)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *KnowledgeRetriever) resolvePolicies(knowledgeBaseIDs []int64, opts KnowledgeRetrieveOptions) []KnowledgeBaseRetrievePolicy {
|
||||||
|
if len(knowledgeBaseIDs) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
knowledgeBases := loadRuntimeKnowledgeBases(knowledgeBaseIDs)
|
||||||
|
ret := make([]KnowledgeBaseRetrievePolicy, 0, len(knowledgeBaseIDs))
|
||||||
|
for _, knowledgeBaseID := range knowledgeBaseIDs {
|
||||||
|
policy := KnowledgeBaseRetrievePolicy{
|
||||||
|
KnowledgeBaseID: knowledgeBaseID,
|
||||||
|
TopK: defaultRuntimeKnowledgeTopK,
|
||||||
|
ScoreThreshold: defaultRuntimeKnowledgeScoreThreshold,
|
||||||
|
}
|
||||||
|
if knowledgeBase, ok := knowledgeBases[knowledgeBaseID]; ok {
|
||||||
|
if knowledgeBase.DefaultTopK > 0 {
|
||||||
|
policy.TopK = knowledgeBase.DefaultTopK
|
||||||
|
}
|
||||||
|
if knowledgeBase.DefaultScoreThreshold > 0 {
|
||||||
|
policy.ScoreThreshold = knowledgeBase.DefaultScoreThreshold
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if opts.TopK > 0 {
|
||||||
|
policy.TopK = opts.TopK
|
||||||
|
}
|
||||||
|
if opts.ScoreThreshold > 0 {
|
||||||
|
policy.ScoreThreshold = opts.ScoreThreshold
|
||||||
|
}
|
||||||
|
ret = append(ret, policy)
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveRuntimeAnswerMode(knowledgeBaseIDs []int64, results []rag.RetrieveResult) enums.KnowledgeAnswerMode {
|
||||||
|
knowledgeBases := loadRuntimeKnowledgeBases(knowledgeBaseIDs)
|
||||||
|
if len(knowledgeBases) == 0 {
|
||||||
|
return enums.KnowledgeAnswerModeStrict
|
||||||
|
}
|
||||||
|
if len(results) > 0 {
|
||||||
|
if knowledgeBase, ok := knowledgeBases[results[0].KnowledgeBaseID]; ok {
|
||||||
|
return normalizeRuntimeAnswerMode(knowledgeBase)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, knowledgeBaseID := range knowledgeBaseIDs {
|
||||||
|
if knowledgeBase, ok := knowledgeBases[knowledgeBaseID]; ok {
|
||||||
|
return normalizeRuntimeAnswerMode(knowledgeBase)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return enums.KnowledgeAnswerModeStrict
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeRuntimeAnswerMode(knowledgeBase models.KnowledgeBase) enums.KnowledgeAnswerMode {
|
||||||
|
answerMode := enums.KnowledgeAnswerMode(knowledgeBase.AnswerMode)
|
||||||
|
if answerMode == 0 {
|
||||||
|
answerMode = enums.KnowledgeAnswerModeStrict
|
||||||
|
}
|
||||||
|
return answerMode
|
||||||
|
}
|
||||||
|
|
||||||
|
func loadRuntimeKnowledgeBases(ids []int64) map[int64]models.KnowledgeBase {
|
||||||
|
if len(ids) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
items := repositories.KnowledgeBaseRepository.Find(sqls.DB(), sqls.NewCnd().In("id", ids))
|
||||||
|
if len(items) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
ret := make(map[int64]models.KnowledgeBase, len(items))
|
||||||
|
for _, item := range items {
|
||||||
|
if item.Status != enums.StatusOk {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
ret[item.ID] = item
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildRetrieverTraceItems(queryPreview string, results []rag.RetrieveResult, trace *rag.RetrieveTrace) []traces.RetrieverTraceItem {
|
||||||
|
if len(results) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
latencyMs := int64(0)
|
||||||
|
if trace != nil {
|
||||||
|
latencyMs = trace.EmbeddingMs + trace.VectorSearchMs + trace.HydrateMs
|
||||||
|
}
|
||||||
|
ret := make([]traces.RetrieverTraceItem, 0, len(results))
|
||||||
|
for _, item := range results {
|
||||||
|
ret = append(ret, traces.RetrieverTraceItem{
|
||||||
|
Query: queryPreview,
|
||||||
|
KnowledgeBaseID: item.KnowledgeBaseID,
|
||||||
|
DocumentID: item.DocumentID,
|
||||||
|
DocumentTitle: item.DocumentTitle,
|
||||||
|
Score: float64(item.Score),
|
||||||
|
LatencyMs: latencyMs,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildRetrieverTraceSummary(opts KnowledgeRetrieveOptions, policies []KnowledgeBaseRetrievePolicy, contextResults []rag.RetrieveResult, results []rag.RetrieveResult, trace *rag.RetrieveTrace) traces.RetrieverTraceSummary {
|
||||||
|
ret := traces.RetrieverTraceSummary{
|
||||||
|
TopK: opts.TopK,
|
||||||
|
ScoreThreshold: opts.ScoreThreshold,
|
||||||
|
ContextMaxTokens: opts.ContextMaxTokens,
|
||||||
|
MaxContextItems: opts.MaxContextItems,
|
||||||
|
HitCount: len(results),
|
||||||
|
ContextCount: len(contextResults),
|
||||||
|
Policies: buildRetrieverPolicyTraceItems(policies),
|
||||||
|
}
|
||||||
|
if ret.TopK <= 0 && len(policies) == 1 {
|
||||||
|
ret.TopK = policies[0].TopK
|
||||||
|
}
|
||||||
|
if ret.ScoreThreshold <= 0 && len(policies) == 1 {
|
||||||
|
ret.ScoreThreshold = policies[0].ScoreThreshold
|
||||||
|
}
|
||||||
|
if trace != nil {
|
||||||
|
ret.EmbeddingMs = trace.EmbeddingMs
|
||||||
|
ret.VectorSearchMs = trace.VectorSearchMs
|
||||||
|
ret.HydrateMs = trace.HydrateMs
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildRetrieverPolicyTraceItems(policies []KnowledgeBaseRetrievePolicy) []traces.RetrieverPolicyTraceItem {
|
||||||
|
if len(policies) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
ret := make([]traces.RetrieverPolicyTraceItem, 0, len(policies))
|
||||||
|
for _, item := range policies {
|
||||||
|
ret = append(ret, traces.RetrieverPolicyTraceItem{
|
||||||
|
KnowledgeBaseID: item.KnowledgeBaseID,
|
||||||
|
TopK: item.TopK,
|
||||||
|
ScoreThreshold: item.ScoreThreshold,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return ret
|
||||||
|
}
|
||||||
64
internal/ai/runtime/runtime_reply_executor.go
Normal file
64
internal/ai/runtime/runtime_reply_executor.go
Normal file
@@ -0,0 +1,64 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
"agent-desk/internal/ai/runtime/graphs"
|
||||||
|
"agent-desk/internal/models"
|
||||||
|
)
|
||||||
|
|
||||||
|
type runtimeReplyExecutor struct{}
|
||||||
|
|
||||||
|
type runtimeReplyRunInput struct {
|
||||||
|
Conversation models.Conversation
|
||||||
|
Message models.Message
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
}
|
||||||
|
|
||||||
|
type runtimeReplyResumeInput struct {
|
||||||
|
Conversation models.Conversation
|
||||||
|
Message models.Message
|
||||||
|
AIAgent models.AIAgent
|
||||||
|
PendingInterrupt *models.ConversationInterrupt
|
||||||
|
}
|
||||||
|
|
||||||
|
func newRuntimeReplyExecutor() *runtimeReplyExecutor {
|
||||||
|
return &runtimeReplyExecutor{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *runtimeReplyExecutor) Run(ctx context.Context, input runtimeReplyRunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
summary, err := applicationruntime.DefaultAgentApplicationService.Run(ctx, applicationruntime.ApplicationRunInput{
|
||||||
|
ConversationID: input.Conversation.ID,
|
||||||
|
MessageID: input.Message.ID,
|
||||||
|
AIAgentID: input.AIAgent.ID,
|
||||||
|
})
|
||||||
|
return summary, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *runtimeReplyExecutor) ResumePendingInterrupt(ctx context.Context, input runtimeReplyResumeInput) (*applicationruntime.RunResult, error) {
|
||||||
|
if input.PendingInterrupt == nil {
|
||||||
|
return nil, fmt.Errorf("pending interrupt is required")
|
||||||
|
}
|
||||||
|
summary, err := applicationruntime.DefaultAgentApplicationService.Resume(ctx, applicationruntime.ApplicationResumeInput{
|
||||||
|
ApplicationRunInput: applicationruntime.ApplicationRunInput{
|
||||||
|
ConversationID: input.Conversation.ID,
|
||||||
|
MessageID: input.Message.ID,
|
||||||
|
AIAgentID: input.AIAgent.ID,
|
||||||
|
},
|
||||||
|
CheckPointID: strings.TrimSpace(input.PendingInterrupt.CheckPointID),
|
||||||
|
ResumeData: map[string]string{
|
||||||
|
strings.TrimSpace(input.PendingInterrupt.InterruptID): strings.TrimSpace(input.Message.Content),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
return summary, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func expiredInterruptSummary() *applicationruntime.RunResult {
|
||||||
|
return &applicationruntime.RunResult{
|
||||||
|
Status: "expired",
|
||||||
|
ReplyText: graphs.ConfirmationExpiredReply,
|
||||||
|
}
|
||||||
|
}
|
||||||
27
internal/ai/runtime/service.go
Normal file
27
internal/ai/runtime/service.go
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
package runtime
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
|
||||||
|
applicationruntime "agent-desk/internal/ai/application/runtime"
|
||||||
|
)
|
||||||
|
|
||||||
|
var Service = newService()
|
||||||
|
|
||||||
|
func newService() *service {
|
||||||
|
return &service{
|
||||||
|
app: applicationruntime.NewService(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type service struct {
|
||||||
|
app *applicationruntime.Service
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *service) Run(ctx context.Context, req applicationruntime.RunInput) (*applicationruntime.RunResult, error) {
|
||||||
|
return s.app.Run(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *service) Resume(ctx context.Context, req applicationruntime.ResumeInput) (*applicationruntime.RunResult, error) {
|
||||||
|
return s.app.Resume(ctx, req)
|
||||||
|
}
|
||||||
34
internal/ai/runtime/tooling/mcp_tool_definition.go
Normal file
34
internal/ai/runtime/tooling/mcp_tool_definition.go
Normal file
@@ -0,0 +1,34 @@
|
|||||||
|
package tooling
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"hash/crc32"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
var toolNameSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_]`)
|
||||||
|
|
||||||
|
type MCPToolDefinition struct {
|
||||||
|
ToolCode string
|
||||||
|
ServerCode string
|
||||||
|
ToolName string
|
||||||
|
ModelName string
|
||||||
|
Title string
|
||||||
|
Description string
|
||||||
|
FixedArgs map[string]string
|
||||||
|
}
|
||||||
|
|
||||||
|
func BuildModelToolName(definition MCPToolDefinition) string {
|
||||||
|
if strings.TrimSpace(definition.ModelName) != "" {
|
||||||
|
return strings.TrimSpace(definition.ModelName)
|
||||||
|
}
|
||||||
|
base := "mcp_" + strings.TrimSpace(definition.ServerCode) + "_" + strings.TrimSpace(definition.ToolName)
|
||||||
|
base = toolNameSanitizer.ReplaceAllString(base, "_")
|
||||||
|
base = strings.Trim(base, "_")
|
||||||
|
if base == "" {
|
||||||
|
base = "mcp_tool"
|
||||||
|
}
|
||||||
|
checksum := crc32.ChecksumIEEE([]byte(definition.ToolCode))
|
||||||
|
return fmt.Sprintf("%s_%08x", base, checksum)
|
||||||
|
}
|
||||||
42
internal/ai/runtime/tooling/tool_result.go
Normal file
42
internal/ai/runtime/tooling/tool_result.go
Normal file
@@ -0,0 +1,42 @@
|
|||||||
|
package tooling
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
type ToolResult struct {
|
||||||
|
Handled bool `json:"handled"`
|
||||||
|
Terminal bool `json:"terminal"`
|
||||||
|
Action string `json:"action"`
|
||||||
|
ReplyText string `json:"replyText,omitempty"`
|
||||||
|
ReplySent bool `json:"replySent,omitempty"`
|
||||||
|
ShouldRetry bool `json:"shouldRetry"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func MarshalToolResult(result ToolResult) string {
|
||||||
|
result.Action = strings.TrimSpace(result.Action)
|
||||||
|
result.ReplyText = strings.TrimSpace(result.ReplyText)
|
||||||
|
buf, err := json.Marshal(result)
|
||||||
|
if err != nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return string(buf)
|
||||||
|
}
|
||||||
|
|
||||||
|
func ParseToolResult(raw string) (ToolResult, bool) {
|
||||||
|
raw = strings.TrimSpace(raw)
|
||||||
|
if raw == "" {
|
||||||
|
return ToolResult{}, false
|
||||||
|
}
|
||||||
|
var result ToolResult
|
||||||
|
if err := json.Unmarshal([]byte(raw), &result); err != nil {
|
||||||
|
return ToolResult{}, false
|
||||||
|
}
|
||||||
|
result.Action = strings.TrimSpace(result.Action)
|
||||||
|
result.ReplyText = strings.TrimSpace(result.ReplyText)
|
||||||
|
if result.Action == "" && result.ReplyText == "" && !result.Handled && !result.Terminal {
|
||||||
|
return ToolResult{}, false
|
||||||
|
}
|
||||||
|
return result, true
|
||||||
|
}
|
||||||
124
internal/ai/runtime/tooling/tool_result_reducer.go
Normal file
124
internal/ai/runtime/tooling/tool_result_reducer.go
Normal file
@@ -0,0 +1,124 @@
|
|||||||
|
package tooling
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"regexp"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"agent-desk/internal/ai/mcps"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
maxToolResultSummaryChars = 2500
|
||||||
|
maxToolResultSegments = 12
|
||||||
|
)
|
||||||
|
|
||||||
|
var reductionInfoPattern = regexp.MustCompile(`\[tool result reduced: original_length=(\d+), kept_length=(\d+)\]`)
|
||||||
|
|
||||||
|
type ReductionInfo struct {
|
||||||
|
Reduced bool
|
||||||
|
OriginalChars int
|
||||||
|
KeptChars int
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildReducedToolResultSummary returns a bounded text summary for MCP tool results.
|
||||||
|
// It keeps the main payload visible to the model while preventing a single large tool
|
||||||
|
// response from exhausting too much context.
|
||||||
|
func BuildReducedToolResultSummary(result *mcps.ToolCallResult) string {
|
||||||
|
if result == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
segments := collectToolResultSegments(result)
|
||||||
|
if len(segments) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
text := strings.TrimSpace(strings.Join(segments, "\n"))
|
||||||
|
if text == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
runes := []rune(text)
|
||||||
|
if len(runes) <= maxToolResultSummaryChars {
|
||||||
|
return text
|
||||||
|
}
|
||||||
|
truncated := strings.TrimSpace(string(runes[:maxToolResultSummaryChars]))
|
||||||
|
return fmt.Sprintf("%s\n\n[tool result reduced: original_length=%d, kept_length=%d]", truncated, len(runes), maxToolResultSummaryChars)
|
||||||
|
}
|
||||||
|
|
||||||
|
func ParseReductionInfo(summary string) ReductionInfo {
|
||||||
|
matches := reductionInfoPattern.FindStringSubmatch(strings.TrimSpace(summary))
|
||||||
|
if len(matches) != 3 {
|
||||||
|
return ReductionInfo{}
|
||||||
|
}
|
||||||
|
originalChars, err1 := strconv.Atoi(matches[1])
|
||||||
|
keptChars, err2 := strconv.Atoi(matches[2])
|
||||||
|
if err1 != nil || err2 != nil {
|
||||||
|
return ReductionInfo{}
|
||||||
|
}
|
||||||
|
return ReductionInfo{
|
||||||
|
Reduced: true,
|
||||||
|
OriginalChars: originalChars,
|
||||||
|
KeptChars: keptChars,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func collectToolResultSegments(result *mcps.ToolCallResult) []string {
|
||||||
|
segments := make([]string, 0, len(result.Content)+2)
|
||||||
|
if result.IsError {
|
||||||
|
segments = append(segments, "tool returned an error")
|
||||||
|
}
|
||||||
|
if result.StructuredContent != nil {
|
||||||
|
if data, err := json.Marshal(result.StructuredContent); err == nil {
|
||||||
|
segments = appendNonBlankSegment(segments, string(data))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, item := range result.Content {
|
||||||
|
if len(segments) >= maxToolResultSegments {
|
||||||
|
segments = append(segments, "[tool result reduced: remaining segments omitted]")
|
||||||
|
break
|
||||||
|
}
|
||||||
|
switch item.Type {
|
||||||
|
case "text":
|
||||||
|
segments = appendNonBlankSegment(segments, item.Text)
|
||||||
|
default:
|
||||||
|
if item.Data == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if data, err := json.Marshal(item.Data); err == nil {
|
||||||
|
segments = appendNonBlankSegment(segments, string(data))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return segments
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendNonBlankSegment(input []string, value string) []string {
|
||||||
|
value = strings.TrimSpace(value)
|
||||||
|
if value == "" {
|
||||||
|
return input
|
||||||
|
}
|
||||||
|
key := canonicalToolResultSegment(value)
|
||||||
|
for _, existing := range input {
|
||||||
|
if canonicalToolResultSegment(existing) == key {
|
||||||
|
return input
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return append(input, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
func canonicalToolResultSegment(value string) string {
|
||||||
|
value = strings.TrimSpace(value)
|
||||||
|
if value == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
var payload any
|
||||||
|
if err := json.Unmarshal([]byte(value), &payload); err != nil {
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
return string(data)
|
||||||
|
}
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user