feat(observability): instrument agent ReAct loop and tool calls in Langfuse
The existing Langfuse integration covered Chat / Embedding / Rerank / VLM /
ASR generations plus the HTTP + asynq spans, but the agent's own execution
tree was invisible: tool calls never appeared, multi-round ReAct iterations
were flat under the HTTP trace, and there was no single node representing
"one agent run".
This change adds three levels of agent-side spans:
- agent.execute — wraps AgentEngine.Execute, records query preview,
knowledge bases, allowed tools, final-answer length
and totals on finish.
- agent.round.<N> — wraps each ReAct iteration; records finish_reason,
tool-call count, token usage and duration.
- agent.tool.<name> — wraps each tool invocation; records arguments,
success, duration, output preview (rune-safe, 4KB
cap), error, data keys and image count.
To keep the loop's many exit paths (natural stop, stuck loop, empty-content
retry, final_answer, context cancellation) span-safe, the iteration body was
extracted into runReActIteration with a single defer span.Finish() and an
iterOutcome sentinel driving the outer loop. database_query arguments are
redacted (keys only) to avoid leaking raw SQL into the observability
backend, mirroring the existing UI hint policy.
Adds unit tests for the new helpers (truncateForLangfuse, argKeys, dataKeys,
finishToolSpan nil-safety, iterOutcome.String).
This commit is contained in:
+4
-2
@@ -228,8 +228,8 @@ Dev 相关容器都带 `-dev` 后缀、用独立网络 `WeKnora-network-dev`,
|
||||
| Langfuse 概念 | WeKnora 对应 | 备注 |
|
||||
| --- | --- | --- |
|
||||
| Trace | 一次 HTTP 请求(含其触发的所有 asynq 任务) | 对于 `knowledge-chat`、`agent-chat`、`knowledge-search`、`generate_title`、`evaluation`、模型连通性测试等在线请求;以及文件上传/URL 入库/manual/reparse/move/copy、FAQ 导入、知识修改、wiki auto-fix、数据源手工触发等入库请求,HTTP 层都会开启 trace,并把 `trace_id` / `parent_observation_id` 注入 asynq payload。 |
|
||||
| Span(type=SPAN) | 每个 asynq 任务的执行窗口 | 由 `internal/tracing/langfuse/AsynqMiddleware` 在 `mux.Use` 注册;对每个 handler 自动创建 `asynq.<task_type>` 的 SPAN,并记录 `task_id` / `queue` / `retry` / `payload_bytes`。定时任务(无上游 trace)会退化为 `asynq.<task_type>` 独立 trace。 |
|
||||
| Generation(type=GENERATION) | 每次 chat / embedding / rerank / VLM / ASR 调用 | 若位于 span 下会自动设置 `parentObservationId`,所以 Langfuse UI 呈现 trace → asynq-span → generation 的树状结构。 |
|
||||
| Span(type=SPAN) | 每个 asynq 任务的执行窗口 / 每次 Agent 执行及其每一轮 / 每次工具调用 | 由 `internal/tracing/langfuse/AsynqMiddleware` 在 `mux.Use` 注册;对每个 handler 自动创建 `asynq.<task_type>` 的 SPAN,并记录 `task_id` / `queue` / `retry` / `payload_bytes`。定时任务(无上游 trace)会退化为 `asynq.<task_type>` 独立 trace。**Agent 相关**:`AgentEngine.Execute` 会开 `agent.execute` 顶层 SPAN,其下每一轮 ReAct 循环开 `agent.round.N` SPAN,每次工具调用开 `agent.tool.<tool_name>` SPAN(参数、输出、耗时、成败、错误都会写入)。 |
|
||||
| Generation(type=GENERATION) | 每次 chat / embedding / rerank / VLM / ASR 调用 | 若位于 span 下会自动设置 `parentObservationId`,所以 Langfuse UI 呈现 trace → asynq-span → generation 的树状结构;Agent 模式下是 trace → agent.execute → agent.round.N → (chat.completion.stream + agent.tool.X → rerank/embedding...) 的完整树。 |
|
||||
| Input Tokens | `TokenUsage.PromptTokens` | 来自模型返回的 usage 字段。 |
|
||||
| Output Tokens | `TokenUsage.CompletionTokens` | 来自模型返回的 usage 字段。 |
|
||||
| Total Tokens | `TokenUsage.TotalTokens` | 大多数厂商返回;未返回时自动求和。 |
|
||||
@@ -300,6 +300,8 @@ Dev 相关容器都带 `-dev` 后缀、用独立网络 `WeKnora-network-dev`,
|
||||
- `internal/models/rerank/langfuse_wrapper.go` — Rerank 调用装饰器。
|
||||
- `internal/models/vlm/langfuse_wrapper.go` — VLM(视觉语言模型)调用装饰器。
|
||||
- `internal/models/asr/langfuse_wrapper.go` — ASR(语音识别)调用装饰器。
|
||||
- `internal/agent/engine.go` — `agent.execute` 顶层 SPAN 和 `agent.round.<N>` 每轮 SPAN。
|
||||
- `internal/agent/act.go` — `agent.tool.<tool_name>` 工具调用 SPAN(包含参数、输出、耗时、成败)。
|
||||
- `internal/router/router.go` — 注册 `langfuse.GinMiddleware()`。
|
||||
- `internal/router/task.go` — 在 asynq mux 上 `mux.Use(langfuse.AsynqMiddleware())`,使所有 handler 自动被 trace。
|
||||
- `internal/container/container.go` — 初始化 + 资源清理。
|
||||
|
||||
+134
-2
@@ -3,7 +3,9 @@ package agent
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -11,10 +13,106 @@ import (
|
||||
"github.com/Tencent/WeKnora/internal/common"
|
||||
"github.com/Tencent/WeKnora/internal/event"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
// langfuseToolOutputPreview caps the Output field we send to Langfuse for a
|
||||
// tool call. Tool outputs are already truncated by the registry to
|
||||
// DefaultMaxToolOutput (16KB) before this point, but rendering 16KB in the
|
||||
// Langfuse UI for every tool call is noisy. We keep a generous slice so the
|
||||
// gist is preserved, and include the original length in metadata.
|
||||
const langfuseToolOutputPreview = 4000
|
||||
|
||||
// truncateForLangfuse returns s truncated to at most n runes, with a "…"
|
||||
// marker appended when truncated. Runes (not bytes) are used so multi-byte
|
||||
// CJK content is never split mid-character.
|
||||
func truncateForLangfuse(s string, n int) string {
|
||||
if n <= 0 || len(s) == 0 {
|
||||
return s
|
||||
}
|
||||
r := []rune(s)
|
||||
if len(r) <= n {
|
||||
return s
|
||||
}
|
||||
return string(r[:n]) + "…"
|
||||
}
|
||||
|
||||
// argKeys returns the sorted list of top-level keys in a tool's argument
|
||||
// map. Used when we choose not to send the raw arguments to Langfuse
|
||||
// (e.g. database_query's SQL) but still want to signal what was passed in.
|
||||
func argKeys(args map[string]any) []string {
|
||||
keys := make([]string, 0, len(args))
|
||||
for k := range args {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
|
||||
// finishToolSpan serialises a completed tool call into a Langfuse span
|
||||
// update. Extracted from runToolCall so the tool-call pipeline keeps
|
||||
// a single assignment per line and the observability-specific logic
|
||||
// (payload shaping, error classification) lives in one place.
|
||||
func finishToolSpan(span *langfuse.Span, tc types.ToolCall, execErr error, durationMs int64) {
|
||||
if span == nil {
|
||||
return
|
||||
}
|
||||
success := tc.Result != nil && tc.Result.Success
|
||||
output := map[string]interface{}{
|
||||
"success": success,
|
||||
"duration_ms": durationMs,
|
||||
}
|
||||
if tc.Result != nil {
|
||||
if tc.Result.Output != "" {
|
||||
output["output"] = truncateForLangfuse(tc.Result.Output, langfuseToolOutputPreview)
|
||||
output["output_len"] = len(tc.Result.Output)
|
||||
}
|
||||
if tc.Result.Error != "" {
|
||||
output["error"] = tc.Result.Error
|
||||
}
|
||||
if len(tc.Result.Data) > 0 {
|
||||
// Data is structured but can be arbitrarily large (e.g. full
|
||||
// search-result payloads). Only report key shape so Langfuse
|
||||
// users see what was surfaced without blowing up trace size.
|
||||
output["data_keys"] = dataKeys(tc.Result.Data)
|
||||
}
|
||||
if len(tc.Result.Images) > 0 {
|
||||
output["image_count"] = len(tc.Result.Images)
|
||||
}
|
||||
}
|
||||
// Classify the span's outcome: a non-nil execErr is always an error, and
|
||||
// a result with Success=false is treated as an error too (matches the
|
||||
// user-visible behaviour — the LLM would see this as a failed tool call
|
||||
// and try a different approach).
|
||||
var spanErr error
|
||||
switch {
|
||||
case execErr != nil:
|
||||
spanErr = execErr
|
||||
case tc.Result != nil && !tc.Result.Success:
|
||||
msg := tc.Result.Error
|
||||
if msg == "" {
|
||||
msg = "tool returned success=false"
|
||||
}
|
||||
spanErr = errors.New(msg)
|
||||
}
|
||||
span.Finish(output, map[string]interface{}{
|
||||
"success": success,
|
||||
"duration_ms": durationMs,
|
||||
}, spanErr)
|
||||
}
|
||||
|
||||
// dataKeys returns the sorted top-level keys of a tool's Data map.
|
||||
func dataKeys(data map[string]interface{}) []string {
|
||||
keys := make([]string, 0, len(data))
|
||||
for k := range data {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
|
||||
// toolDisplayNames maps internal tool names to user-friendly display labels.
|
||||
var toolDisplayNames = map[string]string{
|
||||
agenttools.ToolThinking: "深度思考",
|
||||
@@ -262,9 +360,41 @@ func (e *AgentEngine) runToolCall(
|
||||
"tool_index": fmt.Sprintf("%d/%s", i+1, total),
|
||||
})
|
||||
|
||||
toolCtx, toolCancel := context.WithTimeout(ctx, defaultToolExecTimeout)
|
||||
// Open a Langfuse span for the tool invocation so the Langfuse UI shows
|
||||
// trace → agent.execute → agent.round.N → agent.tool.<name>, alongside
|
||||
// any nested generations (embedding/rerank/VLM) that the tool itself
|
||||
// triggers. No-op when Langfuse is disabled.
|
||||
mgr := langfuse.GetManager()
|
||||
toolSpanInput := map[string]interface{}{
|
||||
"arguments": args,
|
||||
"tool_call_id": tc.ID,
|
||||
}
|
||||
// database_query's SQL is treated as sensitive by the UI hint layer
|
||||
// (toolHintSensitiveArgs) because it exposes implementation details.
|
||||
// Mirror that policy for Langfuse: redact raw arguments to avoid
|
||||
// leaking raw SQL into the observability backend.
|
||||
if toolHintSensitiveArgs[tc.Function.Name] {
|
||||
toolSpanInput = map[string]interface{}{
|
||||
"tool_call_id": tc.ID,
|
||||
"arg_keys": argKeys(args),
|
||||
"args_redacted": true,
|
||||
}
|
||||
}
|
||||
toolCtx, toolSpan := mgr.StartSpan(ctx, langfuse.SpanOptions{
|
||||
Name: "agent.tool." + tc.Function.Name,
|
||||
Input: toolSpanInput,
|
||||
Metadata: map[string]interface{}{
|
||||
"iteration": iteration,
|
||||
"round": round,
|
||||
"tool_index": i + 1,
|
||||
"tool_call_id": tc.ID,
|
||||
"session_id": sessionID,
|
||||
},
|
||||
})
|
||||
|
||||
execCtx, toolCancel := context.WithTimeout(toolCtx, defaultToolExecTimeout)
|
||||
result, err := e.toolRegistry.ExecuteTool(
|
||||
toolCtx, tc.Function.Name,
|
||||
execCtx, tc.Function.Name,
|
||||
json.RawMessage(tc.Function.Arguments),
|
||||
)
|
||||
toolCancel()
|
||||
@@ -294,6 +424,8 @@ func (e *AgentEngine) runToolCall(
|
||||
toolTag, duration, success, outputLen)
|
||||
}
|
||||
|
||||
finishToolSpan(toolSpan, toolCall, err, duration)
|
||||
|
||||
// Pipeline event for monitoring
|
||||
toolSuccess := toolCall.Result != nil && toolCall.Result.Success
|
||||
pipelineFields := map[string]interface{}{
|
||||
|
||||
+311
-111
@@ -16,10 +16,15 @@ import (
|
||||
"github.com/Tencent/WeKnora/internal/event"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/models/chat"
|
||||
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// langfuseQueryPreview caps the query length we ship as the agent.execute
|
||||
// span input — long quoted-context queries can be many KB of prose.
|
||||
const langfuseQueryPreview = 2000
|
||||
|
||||
// AgentEngine is the core engine for running ReAct agents
|
||||
type AgentEngine struct {
|
||||
config *types.AgentConfig
|
||||
@@ -170,6 +175,41 @@ func (e *AgentEngine) Execute(
|
||||
"context_msgs": len(llmContext),
|
||||
})
|
||||
|
||||
// Open a top-level Langfuse span so the agent run — including every
|
||||
// round's LLM call and every tool execution — groups under a single
|
||||
// node in the Langfuse UI instead of being flat children of the HTTP
|
||||
// trace. No-op when Langfuse is disabled.
|
||||
imgCount := 0
|
||||
if len(imageURLs) > 0 {
|
||||
imgCount = len(imageURLs[0])
|
||||
}
|
||||
kbIDs := make([]string, 0, len(e.knowledgeBasesInfo))
|
||||
for _, kb := range e.knowledgeBasesInfo {
|
||||
if kb != nil {
|
||||
kbIDs = append(kbIDs, kb.ID)
|
||||
}
|
||||
}
|
||||
spanCtx, agentSpan := langfuse.GetManager().StartSpan(ctx, langfuse.SpanOptions{
|
||||
Name: "agent.execute",
|
||||
Input: map[string]interface{}{
|
||||
"query": truncateRunes(query, langfuseQueryPreview),
|
||||
"query_len": len(query),
|
||||
"context_msgs": len(llmContext),
|
||||
"image_count": imgCount,
|
||||
},
|
||||
Metadata: map[string]interface{}{
|
||||
"session_id": sessionID,
|
||||
"message_id": messageID,
|
||||
"max_iterations": e.config.MaxIterations,
|
||||
"parallel_tool_calls": e.config.ParallelToolCalls,
|
||||
"web_search": e.config.WebSearchEnabled,
|
||||
"multi_turn": e.config.MultiTurnEnabled,
|
||||
"knowledge_base_ids": kbIDs,
|
||||
"allowed_tools": e.config.AllowedTools,
|
||||
},
|
||||
})
|
||||
ctx = spanCtx
|
||||
|
||||
// Initialize state
|
||||
state := &types.AgentState{
|
||||
RoundSteps: []types.AgentStep{},
|
||||
@@ -241,6 +281,7 @@ func (e *AgentEngine) Execute(
|
||||
SessionID: sessionID,
|
||||
},
|
||||
})
|
||||
finishAgentSpan(agentSpan, state, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -252,9 +293,51 @@ func (e *AgentEngine) Execute(
|
||||
"steps": len(state.RoundSteps),
|
||||
"complete": state.IsComplete,
|
||||
})
|
||||
finishAgentSpan(agentSpan, state, nil)
|
||||
return state, nil
|
||||
}
|
||||
|
||||
// finishAgentSpan records the final outcome of an agent execution onto the
|
||||
// top-level Langfuse span. Extracted so the same payload is used for both
|
||||
// success and error return paths in Execute().
|
||||
func finishAgentSpan(span *langfuse.Span, state *types.AgentState, err error) {
|
||||
if span == nil {
|
||||
return
|
||||
}
|
||||
totalToolCalls := 0
|
||||
for _, step := range state.RoundSteps {
|
||||
totalToolCalls += len(step.ToolCalls)
|
||||
}
|
||||
output := map[string]interface{}{
|
||||
"rounds": state.CurrentRound,
|
||||
"steps": len(state.RoundSteps),
|
||||
"tool_calls": totalToolCalls,
|
||||
"complete": state.IsComplete,
|
||||
"final_answer_len": len(state.FinalAnswer),
|
||||
"final_answer": truncateRunes(state.FinalAnswer, langfuseQueryPreview),
|
||||
}
|
||||
span.Finish(output, map[string]interface{}{
|
||||
"rounds": state.CurrentRound,
|
||||
"steps": len(state.RoundSteps),
|
||||
"tool_calls": totalToolCalls,
|
||||
"complete": state.IsComplete,
|
||||
}, err)
|
||||
}
|
||||
|
||||
// truncateRunes caps s to n runes and appends "…" when truncated. Identical
|
||||
// in spirit to the helper in act.go, but kept locally so both files stay
|
||||
// independent and the truncation budget can diverge if needed.
|
||||
func truncateRunes(s string, n int) string {
|
||||
if n <= 0 || s == "" {
|
||||
return s
|
||||
}
|
||||
r := []rune(s)
|
||||
if len(r) <= n {
|
||||
return s
|
||||
}
|
||||
return string(r[:n]) + "…"
|
||||
}
|
||||
|
||||
// executeLoop executes the main ReAct loop
|
||||
// All events are emitted through EventBus with the given sessionID
|
||||
func (e *AgentEngine) executeLoop(
|
||||
@@ -273,6 +356,7 @@ func (e *AgentEngine) executeLoop(
|
||||
emptyRetries := 0
|
||||
consecutiveSameContent := 0
|
||||
lastResponseContent := ""
|
||||
loop:
|
||||
for state.CurrentRound < e.config.MaxIterations {
|
||||
// Check for context cancellation (request timeout, user cancel, etc.)
|
||||
select {
|
||||
@@ -290,119 +374,23 @@ func (e *AgentEngine) executeLoop(
|
||||
default:
|
||||
}
|
||||
|
||||
roundStart := time.Now()
|
||||
|
||||
// Context window management: estimate current token count using
|
||||
// the API-reported usage from the previous round plus a BPE delta
|
||||
// for newly appended messages (assistant reply + tool results).
|
||||
currentTokens := e.estimateCurrentTokens(messages)
|
||||
beforeLen := len(messages)
|
||||
messages = e.manageContextWindow(ctx, messages, state.CurrentRound+1, currentTokens)
|
||||
if len(messages) < beforeLen {
|
||||
currentTokens = e.tokenEstimator.EstimateMessages(messages)
|
||||
// Each iteration runs inside an "agent.round.<N>" Langfuse span.
|
||||
// We execute the body in a closure so `defer span.Finish()` fires at
|
||||
// every exit path (break/continue/next) without having to sprinkle
|
||||
// manual finish calls throughout the many branches below.
|
||||
outcome, iterErr := e.runReActIteration(ctx, state, &messages, tools,
|
||||
sessionID, query, &emptyRetries, &consecutiveSameContent, &lastResponseContent)
|
||||
if iterErr != nil {
|
||||
return state, iterErr
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[Agent][Round-%d/%d] Starting: %d messages, %d tools, est_tokens=%d",
|
||||
state.CurrentRound+1, e.config.MaxIterations, len(messages), len(tools), currentTokens)
|
||||
common.PipelineInfo(ctx, "Agent", "round_start", map[string]interface{}{
|
||||
"iteration": state.CurrentRound,
|
||||
"round": state.CurrentRound + 1,
|
||||
"message_count": len(messages),
|
||||
"pending_tools": len(tools),
|
||||
"max_iterations": e.config.MaxIterations,
|
||||
})
|
||||
|
||||
// 1. Think: Call LLM with function calling (includes retry + graceful degradation)
|
||||
e.lastSentMsgCount = len(messages)
|
||||
response, err := e.callLLMWithRetry(ctx, messages, tools, state, query, state.CurrentRound, sessionID)
|
||||
if err != nil {
|
||||
return state, err
|
||||
switch outcome {
|
||||
case iterOutcomeContinue:
|
||||
continue loop
|
||||
case iterOutcomeBreak:
|
||||
break loop
|
||||
case iterOutcomeNext:
|
||||
state.CurrentRound++
|
||||
}
|
||||
if response == nil {
|
||||
break
|
||||
}
|
||||
if response.Usage.TotalTokens > 0 {
|
||||
e.lastUsage = response.Usage
|
||||
logger.Debugf(ctx, "[Agent][Round-%d] Usage: prompt=%d, completion=%d, total=%d",
|
||||
state.CurrentRound+1, response.Usage.PromptTokens,
|
||||
response.Usage.CompletionTokens, response.Usage.TotalTokens)
|
||||
}
|
||||
|
||||
// Detect stuck loops: if the LLM keeps returning the same content
|
||||
// without tool calls (e.g., an unhandled finish reason), break early.
|
||||
if len(response.ToolCalls) == 0 && response.Content != "" {
|
||||
if response.Content == lastResponseContent {
|
||||
consecutiveSameContent++
|
||||
} else {
|
||||
consecutiveSameContent = 0
|
||||
}
|
||||
lastResponseContent = response.Content
|
||||
if consecutiveSameContent >= maxRepeatedResponseRounds {
|
||||
logger.Warnf(ctx, "[Agent][Round-%d] Detected stuck loop: same content repeated %d times (finish=%s), stopping",
|
||||
state.CurrentRound+1, consecutiveSameContent+1, response.FinishReason)
|
||||
state.FinalAnswer = response.Content
|
||||
state.IsComplete = true
|
||||
break
|
||||
}
|
||||
} else {
|
||||
consecutiveSameContent = 0
|
||||
lastResponseContent = ""
|
||||
}
|
||||
|
||||
// Create agent step
|
||||
step := types.AgentStep{
|
||||
Iteration: state.CurrentRound,
|
||||
Thought: response.Content,
|
||||
ToolCalls: make([]types.ToolCall, 0),
|
||||
Timestamp: time.Now(),
|
||||
}
|
||||
|
||||
// 2. Analyze: Check for stop conditions (natural stop or final_answer tool)
|
||||
verdict := e.analyzeResponse(ctx, response, step, state.CurrentRound, sessionID, roundStart)
|
||||
if verdict.isDone {
|
||||
// Guard against empty content: when the LLM stops naturally with no
|
||||
// content and no tool calls (e.g., thinking-only loop without KB),
|
||||
// retry with a nudge message instead of accepting an empty answer.
|
||||
if verdict.emptyContent {
|
||||
emptyRetries++
|
||||
if emptyRetries <= maxEmptyResponseRetries {
|
||||
logger.Warnf(ctx, "[Agent][Round-%d] Empty content with stop - retrying (%d/%d)",
|
||||
state.CurrentRound+1, emptyRetries, maxEmptyResponseRetries)
|
||||
messages = append(messages, chat.Message{
|
||||
Role: "user",
|
||||
Content: "Please provide your answer by calling the final_answer tool.",
|
||||
})
|
||||
continue
|
||||
}
|
||||
// Retries exhausted — use fallback message rather than empty answer
|
||||
logger.Warnf(ctx, "[Agent][Round-%d] Empty content after %d retries - using fallback",
|
||||
state.CurrentRound+1, maxEmptyResponseRetries)
|
||||
state.FinalAnswer = "I'm sorry, I was unable to generate a response. Please try again."
|
||||
state.IsComplete = true
|
||||
state.RoundSteps = append(state.RoundSteps, verdict.step)
|
||||
break
|
||||
}
|
||||
state.FinalAnswer = verdict.finalAnswer
|
||||
state.IsComplete = true
|
||||
state.RoundSteps = append(state.RoundSteps, verdict.step)
|
||||
break
|
||||
}
|
||||
|
||||
// 3. Act: Execute tool calls
|
||||
e.executeToolCalls(ctx, response, &step, state.CurrentRound, sessionID)
|
||||
|
||||
// 4. Observe: Add tool results to messages and write to context
|
||||
state.RoundSteps = append(state.RoundSteps, step)
|
||||
messages = e.appendToolResults(ctx, messages, step)
|
||||
common.PipelineInfo(ctx, "Agent", "round_end", map[string]interface{}{
|
||||
"iteration": state.CurrentRound,
|
||||
"round": state.CurrentRound + 1,
|
||||
"tool_calls": len(step.ToolCalls),
|
||||
"thought_len": len(step.Thought),
|
||||
})
|
||||
|
||||
// 5. Advance to next round
|
||||
state.CurrentRound++
|
||||
}
|
||||
|
||||
// If loop finished without final answer, generate one
|
||||
@@ -416,6 +404,218 @@ func (e *AgentEngine) executeLoop(
|
||||
return state, nil
|
||||
}
|
||||
|
||||
// iterOutcome directs executeLoop's control flow after one ReAct iteration.
|
||||
// Using a sentinel (rather than bare return values from runReActIteration)
|
||||
// keeps the loop's break/continue/next branches explicit in one place.
|
||||
type iterOutcome int
|
||||
|
||||
const (
|
||||
// iterOutcomeNext advances state.CurrentRound and loops again.
|
||||
iterOutcomeNext iterOutcome = iota
|
||||
// iterOutcomeContinue re-runs the loop without advancing the round
|
||||
// counter. Used by the empty-content retry path.
|
||||
iterOutcomeContinue
|
||||
// iterOutcomeBreak exits the loop (final answer, stuck loop, or end).
|
||||
iterOutcomeBreak
|
||||
)
|
||||
|
||||
// runReActIteration executes one ReAct step: think → analyze → act → observe.
|
||||
// Extracted from executeLoop so the whole iteration body can live inside a
|
||||
// single `defer span.Finish()` scope — otherwise we'd need to sprinkle
|
||||
// manual finish calls across every break/continue/return branch.
|
||||
//
|
||||
// The mutable loop state (messages, empty-retry counter, stuck-loop detector)
|
||||
// is passed by pointer so iterations share progress.
|
||||
func (e *AgentEngine) runReActIteration(
|
||||
parentCtx context.Context,
|
||||
state *types.AgentState,
|
||||
messagesPtr *[]chat.Message,
|
||||
tools []chat.Tool,
|
||||
sessionID, query string,
|
||||
emptyRetries, consecutiveSameContent *int,
|
||||
lastResponseContent *string,
|
||||
) (outcome iterOutcome, retErr error) {
|
||||
roundStart := time.Now()
|
||||
round := state.CurrentRound + 1
|
||||
|
||||
// Open the round-level Langfuse span. Any chat/tool calls made inside
|
||||
// this iteration will attach under it via ctx, giving the UI a clean
|
||||
// trace → agent.execute → agent.round.N → (chat + tools) structure.
|
||||
ctx, roundSpan := langfuse.GetManager().StartSpan(parentCtx, langfuse.SpanOptions{
|
||||
Name: fmt.Sprintf("agent.round.%d", round),
|
||||
Input: map[string]interface{}{
|
||||
"round": round,
|
||||
"message_count": len(*messagesPtr),
|
||||
"max_iterations": e.config.MaxIterations,
|
||||
},
|
||||
Metadata: map[string]interface{}{
|
||||
"iteration": state.CurrentRound,
|
||||
"round": round,
|
||||
"session_id": sessionID,
|
||||
},
|
||||
})
|
||||
|
||||
var (
|
||||
response *types.ChatResponse
|
||||
toolCallCount int
|
||||
)
|
||||
defer func() {
|
||||
if roundSpan == nil {
|
||||
return
|
||||
}
|
||||
out := map[string]interface{}{
|
||||
"round": round,
|
||||
"outcome": outcome.String(),
|
||||
"tool_calls": toolCallCount,
|
||||
}
|
||||
if response != nil {
|
||||
out["has_tool_calls"] = len(response.ToolCalls) > 0
|
||||
out["finish_reason"] = response.FinishReason
|
||||
out["content_len"] = len(response.Content)
|
||||
if response.Usage.TotalTokens > 0 {
|
||||
out["prompt_tokens"] = response.Usage.PromptTokens
|
||||
out["completion_tokens"] = response.Usage.CompletionTokens
|
||||
out["total_tokens"] = response.Usage.TotalTokens
|
||||
}
|
||||
}
|
||||
out["duration_ms"] = time.Since(roundStart).Milliseconds()
|
||||
roundSpan.Finish(out, map[string]interface{}{
|
||||
"round": round,
|
||||
"tool_calls": toolCallCount,
|
||||
"outcome": outcome.String(),
|
||||
"duration_ms": time.Since(roundStart).Milliseconds(),
|
||||
}, retErr)
|
||||
}()
|
||||
|
||||
// Context window management: estimate current token count using
|
||||
// the API-reported usage from the previous round plus a BPE delta
|
||||
// for newly appended messages (assistant reply + tool results).
|
||||
currentTokens := e.estimateCurrentTokens(*messagesPtr)
|
||||
beforeLen := len(*messagesPtr)
|
||||
*messagesPtr = e.manageContextWindow(ctx, *messagesPtr, round, currentTokens)
|
||||
if len(*messagesPtr) < beforeLen {
|
||||
currentTokens = e.tokenEstimator.EstimateMessages(*messagesPtr)
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[Agent][Round-%d/%d] Starting: %d messages, %d tools, est_tokens=%d",
|
||||
round, e.config.MaxIterations, len(*messagesPtr), len(tools), currentTokens)
|
||||
common.PipelineInfo(ctx, "Agent", "round_start", map[string]interface{}{
|
||||
"iteration": state.CurrentRound,
|
||||
"round": round,
|
||||
"message_count": len(*messagesPtr),
|
||||
"pending_tools": len(tools),
|
||||
"max_iterations": e.config.MaxIterations,
|
||||
})
|
||||
|
||||
// 1. Think: Call LLM with function calling (includes retry + graceful degradation)
|
||||
e.lastSentMsgCount = len(*messagesPtr)
|
||||
resp, err := e.callLLMWithRetry(ctx, *messagesPtr, tools, state, query, state.CurrentRound, sessionID)
|
||||
if err != nil {
|
||||
retErr = err
|
||||
return iterOutcomeNext, err
|
||||
}
|
||||
if resp == nil {
|
||||
return iterOutcomeBreak, nil
|
||||
}
|
||||
response = resp
|
||||
if response.Usage.TotalTokens > 0 {
|
||||
e.lastUsage = response.Usage
|
||||
logger.Debugf(ctx, "[Agent][Round-%d] Usage: prompt=%d, completion=%d, total=%d",
|
||||
round, response.Usage.PromptTokens,
|
||||
response.Usage.CompletionTokens, response.Usage.TotalTokens)
|
||||
}
|
||||
|
||||
// Detect stuck loops: if the LLM keeps returning the same content
|
||||
// without tool calls (e.g., an unhandled finish reason), break early.
|
||||
if len(response.ToolCalls) == 0 && response.Content != "" {
|
||||
if response.Content == *lastResponseContent {
|
||||
*consecutiveSameContent++
|
||||
} else {
|
||||
*consecutiveSameContent = 0
|
||||
}
|
||||
*lastResponseContent = response.Content
|
||||
if *consecutiveSameContent >= maxRepeatedResponseRounds {
|
||||
logger.Warnf(ctx, "[Agent][Round-%d] Detected stuck loop: same content repeated %d times (finish=%s), stopping",
|
||||
round, *consecutiveSameContent+1, response.FinishReason)
|
||||
state.FinalAnswer = response.Content
|
||||
state.IsComplete = true
|
||||
return iterOutcomeBreak, nil
|
||||
}
|
||||
} else {
|
||||
*consecutiveSameContent = 0
|
||||
*lastResponseContent = ""
|
||||
}
|
||||
|
||||
// Create agent step
|
||||
step := types.AgentStep{
|
||||
Iteration: state.CurrentRound,
|
||||
Thought: response.Content,
|
||||
ToolCalls: make([]types.ToolCall, 0),
|
||||
Timestamp: time.Now(),
|
||||
}
|
||||
|
||||
// 2. Analyze: Check for stop conditions (natural stop or final_answer tool)
|
||||
verdict := e.analyzeResponse(ctx, response, step, state.CurrentRound, sessionID, roundStart)
|
||||
if verdict.isDone {
|
||||
// Guard against empty content: when the LLM stops naturally with no
|
||||
// content and no tool calls (e.g., thinking-only loop without KB),
|
||||
// retry with a nudge message instead of accepting an empty answer.
|
||||
if verdict.emptyContent {
|
||||
*emptyRetries++
|
||||
if *emptyRetries <= maxEmptyResponseRetries {
|
||||
logger.Warnf(ctx, "[Agent][Round-%d] Empty content with stop - retrying (%d/%d)",
|
||||
round, *emptyRetries, maxEmptyResponseRetries)
|
||||
*messagesPtr = append(*messagesPtr, chat.Message{
|
||||
Role: "user",
|
||||
Content: "Please provide your answer by calling the final_answer tool.",
|
||||
})
|
||||
return iterOutcomeContinue, nil
|
||||
}
|
||||
// Retries exhausted — use fallback message rather than empty answer
|
||||
logger.Warnf(ctx, "[Agent][Round-%d] Empty content after %d retries - using fallback",
|
||||
round, maxEmptyResponseRetries)
|
||||
state.FinalAnswer = "I'm sorry, I was unable to generate a response. Please try again."
|
||||
state.IsComplete = true
|
||||
state.RoundSteps = append(state.RoundSteps, verdict.step)
|
||||
return iterOutcomeBreak, nil
|
||||
}
|
||||
state.FinalAnswer = verdict.finalAnswer
|
||||
state.IsComplete = true
|
||||
state.RoundSteps = append(state.RoundSteps, verdict.step)
|
||||
return iterOutcomeBreak, nil
|
||||
}
|
||||
|
||||
// 3. Act: Execute tool calls
|
||||
e.executeToolCalls(ctx, response, &step, state.CurrentRound, sessionID)
|
||||
toolCallCount = len(step.ToolCalls)
|
||||
|
||||
// 4. Observe: Add tool results to messages and write to context
|
||||
state.RoundSteps = append(state.RoundSteps, step)
|
||||
*messagesPtr = e.appendToolResults(ctx, *messagesPtr, step)
|
||||
common.PipelineInfo(ctx, "Agent", "round_end", map[string]interface{}{
|
||||
"iteration": state.CurrentRound,
|
||||
"round": round,
|
||||
"tool_calls": toolCallCount,
|
||||
"thought_len": len(step.Thought),
|
||||
})
|
||||
|
||||
return iterOutcomeNext, nil
|
||||
}
|
||||
|
||||
// String returns a stable label for Langfuse output payloads.
|
||||
func (o iterOutcome) String() string {
|
||||
switch o {
|
||||
case iterOutcomeNext:
|
||||
return "next"
|
||||
case iterOutcomeContinue:
|
||||
return "continue"
|
||||
case iterOutcomeBreak:
|
||||
return "break"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tool result image VLM description helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// TestTruncateForLangfuse verifies rune-aware truncation.
|
||||
// We specifically check CJK and ASCII inputs to ensure the "…" marker is
|
||||
// appended exactly once without splitting multi-byte characters.
|
||||
func TestTruncateForLangfuse(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in string
|
||||
n int
|
||||
want string
|
||||
}{
|
||||
{"empty", "", 10, ""},
|
||||
{"zero-budget", "abcdef", 0, "abcdef"},
|
||||
{"under-budget", "hello", 10, "hello"},
|
||||
{"at-budget", "hello", 5, "hello"},
|
||||
{"over-budget-ascii", "hello world", 5, "hello…"},
|
||||
{"over-budget-cjk", "你好世界测试", 3, "你好世…"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := truncateForLangfuse(tc.in, tc.n)
|
||||
if got != tc.want {
|
||||
t.Fatalf("truncateForLangfuse(%q, %d) = %q; want %q", tc.in, tc.n, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestArgKeysSorted ensures argKeys returns a deterministic (sorted) list so
|
||||
// Langfuse span inputs diff cleanly across runs.
|
||||
func TestArgKeysSorted(t *testing.T) {
|
||||
args := map[string]any{"zeta": 1, "alpha": 2, "mu": 3}
|
||||
got := argKeys(args)
|
||||
want := []string{"alpha", "mu", "zeta"}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("argKeys order = %v; want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDataKeysSorted mirrors TestArgKeysSorted for tool result Data maps.
|
||||
func TestDataKeysSorted(t *testing.T) {
|
||||
data := map[string]interface{}{"z": nil, "a": nil, "m": nil}
|
||||
got := dataKeys(data)
|
||||
want := []string{"a", "m", "z"}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("dataKeys order = %v; want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFinishToolSpanNilSafe is a regression check: callers always invoke
|
||||
// finishToolSpan unconditionally (Langfuse is may-be-disabled) so a nil span
|
||||
// must not panic.
|
||||
func TestFinishToolSpanNilSafe(t *testing.T) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
t.Fatalf("finishToolSpan panicked on nil span: %v", r)
|
||||
}
|
||||
}()
|
||||
finishToolSpan(nil, types.ToolCall{}, errors.New("boom"), 123)
|
||||
finishToolSpan(nil, types.ToolCall{Result: &types.ToolResult{Success: true}}, nil, 0)
|
||||
}
|
||||
|
||||
// TestIterOutcomeString locks in the string labels that surface in Langfuse
|
||||
// span outputs so dashboards built on those values don't silently break.
|
||||
func TestIterOutcomeString(t *testing.T) {
|
||||
cases := []struct {
|
||||
o iterOutcome
|
||||
want string
|
||||
}{
|
||||
{iterOutcomeNext, "next"},
|
||||
{iterOutcomeContinue, "continue"},
|
||||
{iterOutcomeBreak, "break"},
|
||||
{iterOutcome(42), "unknown"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if got := tc.o.String(); got != tc.want {
|
||||
t.Errorf("iterOutcome(%d).String() = %q; want %q", tc.o, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestTruncateRunesEngine verifies the engine-local rune truncator (separate
|
||||
// budget from act.go's helper) behaves identically for the common cases.
|
||||
func TestTruncateRunesEngine(t *testing.T) {
|
||||
if got := truncateRunes("abcde", 3); got != "abc…" {
|
||||
t.Fatalf("truncateRunes short = %q; want %q", got, "abc…")
|
||||
}
|
||||
if got := truncateRunes("abc", 10); got != "abc" {
|
||||
t.Fatalf("truncateRunes under-budget = %q; want abc", got)
|
||||
}
|
||||
if got := truncateRunes("", 10); got != "" {
|
||||
t.Fatalf("truncateRunes empty = %q; want empty", got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user