diff --git a/agent/assistant/llm.go b/agent/assistant/llm.go index de58aac0..b4e9d955 100644 --- a/agent/assistant/llm.go +++ b/agent/assistant/llm.go @@ -47,6 +47,8 @@ func (ast *Assistant) executeLLMStream( // Create LLM instance with connector and options llmInstance, err := llm.New(conn, completionOptions) if err != nil { + // Mark LLM Request as failed in trace + ast.traceLLMFail(ctx, err) return nil, err } @@ -70,6 +72,8 @@ func (ast *Assistant) executeLLMStream( log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err) if err != nil { log.Trace("[AGENT] Calling sendStreamEndOnError") + // Mark LLM Request as failed in trace + ast.traceLLMFail(ctx, err) return nil, err } @@ -108,6 +112,8 @@ func (ast *Assistant) executeLLMForToolRetry( // Create LLM instance with connector and options llmInstance, err := llm.New(conn, completionOptions) if err != nil { + // Mark LLM Retry Request as failed in trace + ast.traceLLMFail(ctx, err) return nil, err } @@ -116,6 +122,8 @@ func (ast *Assistant) executeLLMForToolRetry( completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler) log.Trace("[AGENT] LLM tool retry stream returned: assistant=%s, err=%v", ast.ID, err) if err != nil { + // Mark LLM Retry Request as failed in trace + ast.traceLLMFail(ctx, err) return nil, err } diff --git a/agent/assistant/trace.go b/agent/assistant/trace.go index da15e621..c8807f60 100644 --- a/agent/assistant/trace.go +++ b/agent/assistant/trace.go @@ -84,6 +84,16 @@ func (ast *Assistant) traceLLMComplete(ctx *context.Context, completionResponse trace.Complete(completionResponse) } +// traceLLMFail marks the LLM request as failed in the trace +func (ast *Assistant) traceLLMFail(ctx *context.Context, err error) { + trace, _ := ctx.Trace() + if trace == nil { + return + } + + trace.Fail(err) +} + // traceAgentCompletion creates a completion node to report the final output func (ast *Assistant) traceAgentCompletion(ctx *context.Context, createResponse *context.HookCreateResponse, nextResponse *context.NextHookResponse, completionResponse *context.CompletionResponse, finalResponse interface{}) { trace, _ := ctx.Trace()