init deepseek parser
This commit is contained in:
Grace Guo
parent
e3731fb160
commit
ec65cc3690
|
|
@ -0,0 +1,292 @@
|
|||
package parsers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"unicode"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type DeepSeekParserState int
|
||||
|
||||
const (
|
||||
DeepSeekCollectingThinking DeepSeekParserState = iota
|
||||
DeepSeekCollectingContent
|
||||
DeepSeekCollectingToolCalls
|
||||
DeepSeekCollectingToolOutput
|
||||
)
|
||||
|
||||
const (
|
||||
deepseekThinkingCloseTag = "</think>"
|
||||
deepseekToolCallsBeginTag = "<｜tool▁calls▁begin｜>"
|
||||
deepseekToolCallsEndTag = "<｜tool▁calls▁end｜>"
|
||||
deepseekToolCallBeginTag = "<｜tool▁call▁begin｜>"
|
||||
deepseekToolCallEndTag = "<｜tool▁call▁end｜>"
|
||||
deepseekToolSepTag = "<｜tool▁sep｜>"
|
||||
deepseekToolOutputBeginTag = "<｜tool▁output▁begin｜>"
|
||||
deepseekToolOutputEndTag = "<｜tool▁output▁end｜>"
|
||||
)
|
||||
|
||||
type DeepSeekParser struct {
|
||||
state DeepSeekParserState
|
||||
buffer strings.Builder
|
||||
hasThinkingSupport bool
|
||||
}
|
||||
|
||||
func (p *DeepSeekParser) HasToolSupport() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (p *DeepSeekParser) HasThinkingSupport() bool {
|
||||
return p.hasThinkingSupport
|
||||
}
|
||||
|
||||
func (p *DeepSeekParser) setInitialState(lastMessage *api.Message, tools []api.Tool, thinkValue *api.ThinkValue) {
|
||||
prefill := lastMessage != nil && lastMessage.Role == "assistant"
|
||||
|
||||
// Check both model capability AND request preference
|
||||
thinkingEnabled := p.HasThinkingSupport() && (thinkValue == nil || thinkValue.Bool())
|
||||
|
||||
if !thinkingEnabled {
|
||||
p.state = DeepSeekCollectingContent
|
||||
return
|
||||
}
|
||||
|
||||
if prefill && lastMessage.Content != "" {
|
||||
p.state = DeepSeekCollectingContent
|
||||
return
|
||||
}
|
||||
|
||||
p.state = DeepSeekCollectingThinking
|
||||
}
|
||||
|
||||
func (p *DeepSeekParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
|
||||
p.setInitialState(lastMessage, tools, thinkValue)
|
||||
return tools
|
||||
}
|
||||
|
||||
type deepseekEvent interface {
|
||||
isDeepSeekEvent()
|
||||
}
|
||||
|
||||
type deepseekEventThinkingContent struct {
|
||||
content string
|
||||
}
|
||||
|
||||
type deepseekEventContent struct {
|
||||
content string
|
||||
}
|
||||
|
||||
type deepseekEventToolCall struct {
|
||||
toolCall api.ToolCall
|
||||
}
|
||||
|
||||
func (deepseekEventThinkingContent) isDeepSeekEvent() {}
|
||||
func (deepseekEventContent) isDeepSeekEvent() {}
|
||||
func (deepseekEventToolCall) isDeepSeekEvent() {}
|
||||
|
||||
func (p *DeepSeekParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
|
||||
p.buffer.WriteString(s)
|
||||
events := p.parseEvents()
|
||||
|
||||
var toolCalls []api.ToolCall
|
||||
var contentSb strings.Builder
|
||||
var thinkingSb strings.Builder
|
||||
for _, event := range events {
|
||||
switch event := event.(type) {
|
||||
case deepseekEventToolCall:
|
||||
toolCalls = append(toolCalls, event.toolCall)
|
||||
case deepseekEventThinkingContent:
|
||||
thinkingSb.WriteString(event.content)
|
||||
case deepseekEventContent:
|
||||
contentSb.WriteString(event.content)
|
||||
}
|
||||
}
|
||||
|
||||
return contentSb.String(), thinkingSb.String(), toolCalls, nil
|
||||
}
|
||||
|
||||
func (p *DeepSeekParser) parseEvents() []deepseekEvent {
|
||||
var all []deepseekEvent
|
||||
|
||||
keepLooping := true
|
||||
for keepLooping {
|
||||
var events []deepseekEvent
|
||||
events, keepLooping = p.eat()
|
||||
if len(events) > 0 {
|
||||
all = append(all, events...)
|
||||
}
|
||||
}
|
||||
|
||||
return all
|
||||
}
|
||||
|
||||
func (p *DeepSeekParser) eat() ([]deepseekEvent, bool) {
|
||||
var events []deepseekEvent
|
||||
bufStr := p.buffer.String()
|
||||
if bufStr == "" {
|
||||
return events, false
|
||||
}
|
||||
|
||||
switch p.state {
|
||||
case DeepSeekCollectingThinking:
|
||||
if strings.Contains(bufStr, deepseekThinkingCloseTag) { // thinking[</think>] -> content
|
||||
split := strings.SplitN(bufStr, deepseekThinkingCloseTag, 2)
|
||||
thinking := split[0]
|
||||
thinking = strings.TrimRightFunc(thinking, unicode.IsSpace)
|
||||
|
||||
remaining := split[1]
|
||||
remaining = strings.TrimLeftFunc(remaining, unicode.IsSpace)
|
||||
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(remaining)
|
||||
p.state = DeepSeekCollectingContent
|
||||
|
||||
if len(thinking) > 0 {
|
||||
events = append(events, deepseekEventThinkingContent{content: thinking})
|
||||
}
|
||||
return events, true
|
||||
} else if overlapLen := overlap(bufStr, deepseekThinkingCloseTag); overlapLen > 0 { // partial </think>
|
||||
beforePartialTag := bufStr[:len(bufStr)-overlapLen]
|
||||
trailingLen := trailingWhitespaceLen(beforePartialTag)
|
||||
ambiguousStart := len(beforePartialTag) - trailingLen
|
||||
|
||||
unambiguous := bufStr[:ambiguousStart]
|
||||
ambiguous := bufStr[ambiguousStart:]
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(ambiguous)
|
||||
if len(unambiguous) > 0 {
|
||||
events = append(events, deepseekEventThinkingContent{content: unambiguous})
|
||||
}
|
||||
return events, false
|
||||
} else { // otherwise its thinking content
|
||||
whitespaceLen := trailingWhitespaceLen(bufStr)
|
||||
ambiguousStart := len(bufStr) - whitespaceLen
|
||||
|
||||
unambiguous := bufStr[:ambiguousStart]
|
||||
ambiguous := bufStr[ambiguousStart:]
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(ambiguous)
|
||||
if len(unambiguous) > 0 {
|
||||
events = append(events, deepseekEventThinkingContent{content: unambiguous})
|
||||
}
|
||||
return events, false
|
||||
}
|
||||
|
||||
case DeepSeekCollectingContent:
|
||||
switch {
|
||||
case strings.Contains(bufStr, deepseekToolCallsBeginTag): // content[<｜tool▁calls▁begin｜>] -> tool calls
|
||||
split := strings.SplitN(bufStr, deepseekToolCallsBeginTag, 2)
|
||||
contentBefore := strings.TrimRightFunc(split[0], unicode.IsSpace)
|
||||
remaining := split[1]
|
||||
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(remaining)
|
||||
p.state = DeepSeekCollectingToolCalls
|
||||
|
||||
if len(contentBefore) > 0 {
|
||||
events = append(events, deepseekEventContent{content: contentBefore})
|
||||
}
|
||||
return events, true
|
||||
case strings.Contains(bufStr, deepseekToolOutputBeginTag): // content[<｜tool▁output▁begin｜>] -> tool output
|
||||
split := strings.SplitN(bufStr, deepseekToolOutputBeginTag, 2)
|
||||
contentBefore := split[0] // Don't trim whitespace - preserve spaces
|
||||
remaining := split[1]
|
||||
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(remaining)
|
||||
p.state = DeepSeekCollectingToolOutput
|
||||
|
||||
if len(contentBefore) > 0 {
|
||||
events = append(events, deepseekEventContent{content: contentBefore})
|
||||
}
|
||||
return events, true
|
||||
default: // otherwise its content
|
||||
p.buffer.Reset()
|
||||
if len(bufStr) > 0 {
|
||||
events = append(events, deepseekEventContent{content: bufStr})
|
||||
}
|
||||
return events, false
|
||||
}
|
||||
|
||||
case DeepSeekCollectingToolCalls:
|
||||
if idx := strings.Index(bufStr, deepseekToolCallBeginTag); idx != -1 {
|
||||
startIdx := idx + len(deepseekToolCallBeginTag)
|
||||
if endIdx := strings.Index(bufStr[startIdx:], deepseekToolCallEndTag); endIdx != -1 {
|
||||
toolCallContent := bufStr[startIdx : startIdx+endIdx]
|
||||
|
||||
if toolCall, err := p.parseToolCallContent(toolCallContent); err == nil {
|
||||
remaining := bufStr[startIdx+endIdx+len(deepseekToolCallEndTag):]
|
||||
remaining = strings.TrimLeftFunc(remaining, unicode.IsSpace)
|
||||
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(remaining)
|
||||
|
||||
events = append(events, deepseekEventToolCall{toolCall: toolCall})
|
||||
return events, true
|
||||
} else {
|
||||
slog.Warn("deepseek tool call parsing failed", "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if idx := strings.Index(bufStr, deepseekToolCallsEndTag); idx != -1 {
|
||||
remaining := bufStr[idx+len(deepseekToolCallsEndTag):]
|
||||
remaining = strings.TrimLeftFunc(remaining, unicode.IsSpace)
|
||||
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(remaining)
|
||||
p.state = DeepSeekCollectingContent
|
||||
|
||||
return events, true
|
||||
}
|
||||
|
||||
return events, false
|
||||
|
||||
case DeepSeekCollectingToolOutput:
|
||||
if idx := strings.Index(bufStr, deepseekToolOutputEndTag); idx != -1 {
|
||||
toolOutputContent := bufStr[:idx]
|
||||
remaining := bufStr[idx+len(deepseekToolOutputEndTag):]
|
||||
remaining = strings.TrimLeftFunc(remaining, unicode.IsSpace)
|
||||
|
||||
p.buffer.Reset()
|
||||
p.buffer.WriteString(remaining)
|
||||
p.state = DeepSeekCollectingContent
|
||||
|
||||
if len(toolOutputContent) > 0 {
|
||||
events = append(events, deepseekEventContent{content: toolOutputContent})
|
||||
}
|
||||
return events, true
|
||||
}
|
||||
|
||||
return events, false
|
||||
}
|
||||
|
||||
return events, false
|
||||
}
|
||||
|
||||
func (p *DeepSeekParser) parseToolCallContent(content string) (api.ToolCall, error) {
|
||||
// Expected format: tool_name<｜tool▁sep｜>{args}
|
||||
parts := strings.SplitN(content, deepseekToolSepTag, 2)
|
||||
if len(parts) < 2 {
|
||||
return api.ToolCall{}, errors.New("invalid format")
|
||||
}
|
||||
|
||||
toolName := strings.TrimSpace(parts[0])
|
||||
argsJSON := strings.TrimSpace(parts[1])
|
||||
|
||||
var args api.ToolCallFunctionArguments
|
||||
if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
|
||||
return api.ToolCall{}, err
|
||||
}
|
||||
|
||||
return api.ToolCall{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: toolName,
|
||||
Arguments: args,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
|
@ -0,0 +1,435 @@
|
|||
package parsers
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
func TestDeepSeekParser(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expectedContent string
|
||||
expectedThinking string
|
||||
expectedCalls []api.ToolCall
|
||||
hasThinking bool
|
||||
}{
|
||||
{
|
||||
name: "simple_content",
|
||||
input: "Hello, how are you?",
|
||||
expectedContent: "Hello, how are you?",
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "thinking_content",
|
||||
input: "I need to think about this...</think>The answer is 42.",
|
||||
expectedThinking: "I need to think about this...",
|
||||
expectedContent: "The answer is 42.",
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "no_thinking_simple",
|
||||
input: "Just a regular response.",
|
||||
expectedContent: "Just a regular response.",
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "thinking_with_newlines",
|
||||
input: "Let me think:\n- Point 1\n- Point 2</think>\n\nHere's my answer.",
|
||||
expectedThinking: "Let me think:\n- Point 1\n- Point 2",
|
||||
expectedContent: "Here's my answer.",
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "tool_call_simple",
|
||||
input: "I'll check the weather.<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>{\"location\":\"Paris\"}<｜tool▁call▁end｜><｜tool▁calls▁end｜>",
|
||||
expectedContent: "I'll check the weather.",
|
||||
expectedCalls: []api.ToolCall{
|
||||
{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"location": "Paris",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "multiple_tool_calls",
|
||||
input: "Getting weather for both cities.<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>{\"location\":\"Paris\"}<｜tool▁call▁end｜><｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>{\"location\":\"London\"}<｜tool▁call▁end｜><｜tool▁calls▁end｜>",
|
||||
expectedContent: "Getting weather for both cities.",
|
||||
expectedCalls: []api.ToolCall{
|
||||
{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"location": "Paris",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"location": "London",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "tool_output",
|
||||
input: "Here's the weather: <｜tool▁output▁begin｜>Temperature: 22°C, Sunny<｜tool▁output▁end｜> Hope that helps!",
|
||||
expectedContent: "Here's the weather: Temperature: 22°C, SunnyHope that helps!",
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "complex_tool_arguments",
|
||||
input: "Processing data.<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>process_data<｜tool▁sep｜>{\"items\":[\"item1\",\"item2\"],\"config\":{\"enabled\":true,\"threshold\":0.95}}<｜tool▁call▁end｜><｜tool▁calls▁end｜>",
|
||||
expectedContent: "Processing data.",
|
||||
expectedCalls: []api.ToolCall{
|
||||
{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "process_data",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"items": []interface{}{"item1", "item2"},
|
||||
"config": map[string]interface{}{"enabled": true, "threshold": 0.95},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "thinking_with_tool_call", // technically this can't happen, but the parser can handle it
|
||||
input: "Let me check the weather...</think>I'll get that for you.<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>{\"location\":\"Paris\"}<｜tool▁call▁end｜><｜tool▁calls▁end｜>",
|
||||
expectedThinking: "Let me check the weather...",
|
||||
expectedContent: "I'll get that for you.",
|
||||
expectedCalls: []api.ToolCall{
|
||||
{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"location": "Paris",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "empty_content",
|
||||
input: "",
|
||||
expectedContent: "",
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "only_thinking",
|
||||
input: "Just thinking content</think>",
|
||||
expectedThinking: "Just thinking content",
|
||||
expectedContent: "",
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "multiple_tool_outputs",
|
||||
input: "Results: <｜tool▁output▁begin｜>Paris: 22°C<｜tool▁output▁end｜> and <｜tool▁output▁begin｜>London: 18°C<｜tool▁output▁end｜>",
|
||||
expectedContent: "Results: Paris: 22°Cand London: 18°C",
|
||||
hasThinking: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
parser := &DeepSeekParser{hasThinkingSupport: tt.hasThinking}
|
||||
parser.Init([]api.Tool{}, nil, &api.ThinkValue{Value: tt.hasThinking})
|
||||
|
||||
content, thinking, calls, err := parser.Add(tt.input, true)
|
||||
if err != nil {
|
||||
t.Fatalf("Add() error = %v", err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedContent, content); diff != "" {
|
||||
t.Errorf("Content mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedThinking, thinking); diff != "" {
|
||||
t.Errorf("Thinking mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedCalls, calls); diff != "" {
|
||||
t.Errorf("Tool calls mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepSeekParser_Streaming(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
chunks []string
|
||||
expectedContent string
|
||||
expectedThinking string
|
||||
expectedCalls []api.ToolCall
|
||||
hasThinking bool
|
||||
}{
|
||||
{
|
||||
name: "streaming_simple_content",
|
||||
chunks: []string{"Hello, ", "how are ", "you?"},
|
||||
expectedContent: "Hello, how are you?",
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "streaming_thinking",
|
||||
chunks: []string{"I need to ", "think about this", "...</think>", "The answer is 42."},
|
||||
expectedThinking: "I need to think about this...",
|
||||
expectedContent: "The answer is 42.",
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "streaming_tool_call",
|
||||
chunks: []string{"I'll check weather.", "<｜tool▁calls▁begin｜>", "<｜tool▁call▁begin｜>get_weather", "<｜tool▁sep｜>{\"location\":\"Paris\"}", "<｜tool▁call▁end｜><｜tool▁calls▁end｜>"},
|
||||
expectedContent: "I'll check weather.",
|
||||
expectedCalls: []api.ToolCall{
|
||||
{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"location": "Paris",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
hasThinking: false,
|
||||
},
|
||||
{
|
||||
name: "streaming_thinking_with_partial_tag",
|
||||
chunks: []string{"Thinking about this", "...</", "think>", "Done thinking."},
|
||||
expectedThinking: "Thinking about this...",
|
||||
expectedContent: "Done thinking.",
|
||||
hasThinking: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
parser := &DeepSeekParser{hasThinkingSupport: tt.hasThinking}
|
||||
parser.Init([]api.Tool{}, nil, &api.ThinkValue{Value: tt.hasThinking})
|
||||
|
||||
var allContent, allThinking string
|
||||
var allCalls []api.ToolCall
|
||||
|
||||
for i, chunk := range tt.chunks {
|
||||
done := i == len(tt.chunks)-1
|
||||
content, thinking, calls, err := parser.Add(chunk, done)
|
||||
if err != nil {
|
||||
t.Fatalf("Add() error = %v", err)
|
||||
}
|
||||
|
||||
allContent += content
|
||||
allThinking += thinking
|
||||
allCalls = append(allCalls, calls...)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedContent, allContent); diff != "" {
|
||||
t.Errorf("Content mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedThinking, allThinking); diff != "" {
|
||||
t.Errorf("Thinking mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedCalls, allCalls); diff != "" {
|
||||
t.Errorf("Tool calls mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepSeekParser_HasThinkingSupport(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
hasThinking bool
|
||||
expectedSupport bool
|
||||
}{
|
||||
{
|
||||
name: "thinking_enabled",
|
||||
hasThinking: true,
|
||||
expectedSupport: true,
|
||||
},
|
||||
{
|
||||
name: "thinking_disabled",
|
||||
hasThinking: false,
|
||||
expectedSupport: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
parser := &DeepSeekParser{hasThinkingSupport: tt.hasThinking}
|
||||
if got := parser.HasThinkingSupport(); got != tt.expectedSupport {
|
||||
t.Errorf("HasThinkingSupport() = %v, want %v", got, tt.expectedSupport)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepSeekParser_HasToolSupport(t *testing.T) {
|
||||
parser := &DeepSeekParser{}
|
||||
if !parser.HasToolSupport() {
|
||||
t.Error("HasToolSupport() should return true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepSeekParser_Init(t *testing.T) {
|
||||
parser := &DeepSeekParser{hasThinkingSupport: true}
|
||||
tools := []api.Tool{
|
||||
{
|
||||
Type: "function",
|
||||
Function: api.ToolFunction{
|
||||
Name: "test_tool",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
returnedTools := parser.Init(tools, nil, &api.ThinkValue{Value: true})
|
||||
|
||||
if diff := cmp.Diff(tools, returnedTools); diff != "" {
|
||||
t.Errorf("Init() returned tools mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
// Test initial state is set to thinking when enabled
|
||||
if parser.state != DeepSeekCollectingThinking {
|
||||
t.Errorf("Expected initial state to be DeepSeekCollectingThinking, got %v", parser.state)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepSeekParser_parseToolCallContent(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
expected api.ToolCall
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "valid_tool_call",
|
||||
content: "get_weather<｜tool▁sep｜>{\"location\":\"Paris\"}",
|
||||
expected: api.ToolCall{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"location": "Paris",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "complex_arguments",
|
||||
content: "process_data<｜tool▁sep｜>{\"items\":[\"a\",\"b\"],\"config\":{\"enabled\":true}}",
|
||||
expected: api.ToolCall{
|
||||
Function: api.ToolCallFunction{
|
||||
Name: "process_data",
|
||||
Arguments: api.ToolCallFunctionArguments{
|
||||
"items": []interface{}{"a", "b"},
|
||||
"config": map[string]interface{}{"enabled": true},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "invalid_format_no_separator",
|
||||
content: "get_weather{\"location\":\"Paris\"}",
|
||||
expectError: true,
|
||||
},
|
||||
{
|
||||
name: "invalid_json",
|
||||
content: "get_weather<｜tool▁sep｜>{invalid json}",
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
parser := &DeepSeekParser{}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result, err := parser.parseToolCallContent(tt.content)
|
||||
|
||||
if tt.expectError {
|
||||
if err == nil {
|
||||
t.Error("Expected error but got none")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("Unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expected, result); diff != "" {
|
||||
t.Errorf("parseToolCallContent() mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepSeekParser_EdgeCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expectedContent string
|
||||
expectedThinking string
|
||||
hasThinking bool
|
||||
}{
|
||||
{
|
||||
name: "nested_think_tags_in_thinking",
|
||||
input: "Outer thinking <think>inner</think> content</think>Final content",
|
||||
expectedThinking: "Outer thinking <think>inner",
|
||||
expectedContent: "content</think>Final content",
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "multiple_think_close_tags",
|
||||
input: "First thought</think>Second thought</think>Final content",
|
||||
expectedThinking: "First thought",
|
||||
expectedContent: "Second thought</think>Final content",
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "empty_thinking_content",
|
||||
input: "</think>Just content",
|
||||
expectedThinking: "",
|
||||
expectedContent: "Just content",
|
||||
hasThinking: true,
|
||||
},
|
||||
{
|
||||
name: "thinking_disabled_with_think_tags",
|
||||
input: "Some content</think>More content",
|
||||
expectedContent: "Some content</think>More content",
|
||||
hasThinking: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
parser := &DeepSeekParser{hasThinkingSupport: tt.hasThinking}
|
||||
parser.Init([]api.Tool{}, nil, &api.ThinkValue{Value: tt.hasThinking})
|
||||
|
||||
content, thinking, _, err := parser.Add(tt.input, true)
|
||||
if err != nil {
|
||||
t.Fatalf("Add() error = %v", err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedContent, content); diff != "" {
|
||||
t.Errorf("Content mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tt.expectedThinking, thinking); diff != "" {
|
||||
t.Errorf("Thinking mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -58,6 +58,8 @@ func ParserForName(name string) Parser {
|
|||
return harmony.NewHarmonyMessageHandler()
|
||||
case "cogito":
|
||||
return &CogitoParser{}
|
||||
case "deepseek":
|
||||
return &DeepSeekParser{hasThinkingSupport: true}
|
||||
case "olmo3":
|
||||
return &Olmo3Parser{}
|
||||
case "olmo3-think":
|
||||
|
|
|
|||
Loading…
Reference in New Issue