diff --git a/internal/models/chat/langfuse_wrapper.go b/internal/models/chat/langfuse_wrapper.go index e628969f..bc75fdc8 100644 --- a/internal/models/chat/langfuse_wrapper.go +++ b/internal/models/chat/langfuse_wrapper.go @@ -43,11 +43,9 @@ func (l *langfuseChat) Chat(ctx context.Context, messages []Message, opts *ChatO var output interface{} if resp != nil { usage = convertUsage(&resp.Usage) - output = map[string]interface{}{ - "content": resp.Content, - "tool_calls": resp.ToolCalls, - "finish_reason": resp.FinishReason, - } + output = buildLangfuseGenerationOutput( + resp.Content, resp.ReasoningContent, resp.FinishReason, resp.ToolCalls, + ) } gen.Finish(output, usage, err) return resp, err @@ -85,12 +83,20 @@ func (l *langfuseChat) ChatStream(ctx context.Context, messages []Message, opts go func() { defer close(wrapped) var contentBuf []byte + var reasoningBuf []byte var usage *types.TokenUsage var toolCalls []types.LLMToolCall var finishReason string var firstToken bool for resp := range ch { + if resp.ResponseType == types.ResponseTypeThinking && resp.Content != "" { + if !firstToken { + gen.MarkCompletionStart(time.Now()) + firstToken = true + } + reasoningBuf = append(reasoningBuf, resp.Content...) + } if resp.ResponseType == types.ResponseTypeAnswer && resp.Content != "" { if !firstToken { gen.MarkCompletionStart(time.Now()) @@ -110,11 +116,9 @@ func (l *langfuseChat) ChatStream(ctx context.Context, messages []Message, opts wrapped <- resp } - output := map[string]interface{}{ - "content": string(contentBuf), - "tool_calls": toolCalls, - "finish_reason": finishReason, - } + output := buildLangfuseGenerationOutput( + string(contentBuf), string(reasoningBuf), finishReason, toolCalls, + ) gen.Finish(output, convertUsage(usage), nil) }() return wrapped, nil @@ -141,11 +145,29 @@ func buildLangfuseMessages(messages []Message) []map[string]interface{} { if len(m.ToolCalls) > 0 { entry["tool_calls"] = m.ToolCalls } + if m.ReasoningContent != "" { + entry["reasoning_content"] = m.ReasoningContent + } out = append(out, entry) } return out } +func buildLangfuseGenerationOutput( + content, reasoningContent, finishReason string, + toolCalls []types.LLMToolCall, +) map[string]interface{} { + output := map[string]interface{}{ + "content": content, + "tool_calls": toolCalls, + "finish_reason": finishReason, + } + if reasoningContent != "" { + output["reasoning_content"] = reasoningContent + } + return output +} + func buildLangfuseModelParams(opts *ChatOptions) map[string]interface{} { if opts == nil { return nil diff --git a/internal/models/chat/langfuse_wrapper_test.go b/internal/models/chat/langfuse_wrapper_test.go new file mode 100644 index 00000000..283d4d88 --- /dev/null +++ b/internal/models/chat/langfuse_wrapper_test.go @@ -0,0 +1,45 @@ +package chat + +import ( + "reflect" + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func TestBuildLangfuseGenerationOutput(t *testing.T) { + toolCalls := []types.LLMToolCall{{ID: "call_1", Type: "function"}} + + got := buildLangfuseGenerationOutput("", "", "tool_calls", toolCalls) + want := map[string]interface{}{ + "content": "", + "tool_calls": toolCalls, + "finish_reason": "tool_calls", + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("output without reasoning = %#v; want %#v", got, want) + } + + got = buildLangfuseGenerationOutput("answer", "thinking", "stop", nil) + want = map[string]interface{}{ + "content": "answer", + "tool_calls": []types.LLMToolCall(nil), + "finish_reason": "stop", + "reasoning_content": "thinking", + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("output with reasoning = %#v; want %#v", got, want) + } +} + +func TestBuildLangfuseMessagesReasoningContent(t *testing.T) { + msgs := buildLangfuseMessages([]Message{ + {Role: "assistant", ReasoningContent: "chain of thought", ToolCalls: []ToolCall{{ID: "tc1"}}}, + }) + if len(msgs) != 1 { + t.Fatalf("len(messages) = %d; want 1", len(msgs)) + } + if msgs[0]["reasoning_content"] != "chain of thought" { + t.Fatalf("reasoning_content = %v; want chain of thought", msgs[0]["reasoning_content"]) + } +}